update

app.py

@@ -164,6 +164,8 @@ def generate(
     randomize_cfg: bool,
     text_cfg_scale: float,
     image_cfg_scale: float,
+    weather_hide_video: bool,
+    decode_image_batch: int
 ):
     seed = random.randint(0, 100000) if randomize_seed else seed
     text_cfg_scale = round(random.uniform(6.0, 9.0), ndigits=2) if randomize_cfg else text_cfg_scale
@@ -222,35 +224,30 @@ def generate(
         edited_mask = Image.fromarray(x_1.type(torch.uint8).cpu().numpy())
 
 
-        image_video = []
-
-        batch_size = 10
-        for i in range(0, len(image_list), batch_size):
-            if i + batch_size < len(image_list):
-                tmp_image_list = image_list[i:i+batch_size]
-            else:
-                tmp_image_list = image_list[i:]
-            tmp_image_list = model.decode_first_stage(tmp_image_list)
-            tmp_image_list = torch.clamp((tmp_image_list + 1.0) / 2.0, min=0.0, max=1.0)
-            tmp_image_list = 255.0 * rearrange(tmp_image_list, "b c h w -> b h w c")
-            tmp_image_list = tmp_image_list.type(torch.uint8).cpu().numpy()
-            # image list to image
-            for image in tmp_image_list:
-                image_video.append(image)
-
-        
-        # for i,image in enumerate(mask_list):
-        #     Image.fromarray(image).save(f"test/mask_{i}.png")
-
-        image_video_path = "image.mp4"
-        fps = 30
-        with imageio.get_writer(image_video_path, fps=fps) as video:
-            for image in image_video:
-                video.append_data(image)
-
+        image_video_path = None
+        if not weather_hide_video:
+            image_video = []
+            
+            for i in range(0, len(image_list), decode_image_batch):
+                if i + decode_image_batch < len(image_list):
+                    tmp_image_list = image_list[i:i+decode_image_batch]
+                else:
+                    tmp_image_list = image_list[i:]
+                tmp_image_list = model.decode_first_stage(tmp_image_list)
+                tmp_image_list = torch.clamp((tmp_image_list + 1.0) / 2.0, min=0.0, max=1.0)
+                tmp_image_list = 255.0 * rearrange(tmp_image_list, "b c h w -> b h w c")
+                tmp_image_list = tmp_image_list.type(torch.uint8).cpu().numpy()
+                # image list to image
+                for image in tmp_image_list:
+                    image_video.append(image)
+
+            image_video_path = "image.mp4"
+            fps = 30
+            with imageio.get_writer(image_video_path, fps=fps) as video:
+                for image in image_video:
+                    video.append_data(image)
 
         # 对edited_mask做膨胀
-
         edited_mask_copy = edited_mask.copy()
         kernel = np.ones((3, 3), np.uint8)
         edited_mask = cv2.dilate(np.array(edited_mask), kernel, iterations=3)
@@ -286,7 +283,7 @@ def generate(
         return [int(seed), text_cfg_scale, image_cfg_scale, edited_image, mix_image, edited_mask_copy, mask_video_path, image_video_path, input_image_copy, mix_result_with_red_mask]
 
 def reset():
-    return [100, "Randomize Seed", 1372, "Fix CFG", 7.5, 1.5, None, None, None, None, None, None, None]
+    return [100, "Randomize Seed", 1372, "Fix CFG", 7.5, 1.5, None, None, None, None, None, None, None, "Hide Image Video", 10]
 
 def get_example():
     return [
@@ -336,8 +333,8 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
                 text_cfg_scale = gr.Number(value=7.5, label=f"Text CFG", interactive=True)
                 image_cfg_scale = gr.Number(value=1.5, label=f"Image CFG", interactive=True)
             with gr.Row():
-                generate_button = gr.Button("Generate")
                 reset_button = gr.Button("Reset")
+                generate_button = gr.Button("Generate")
         with gr.Column(scale=1, min_width=100):
             with gr.Column():
                 mix_image = gr.Image(label=f"Mix Image", type="pil", interactive=False)
@@ -346,6 +343,18 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
     
     
     with gr.Accordion('More outputs', open=False):
+        with gr.Row():
+            # 单选框 选择是否显示视频
+            weather_hide_video = gr.Radio(
+                ["Show Image Video", "Hide Image Video"],
+                value="Hide Image Video",
+                type="index",
+                label="Dynamic Process",
+                show_label=False,
+                interactive=True,
+            )
+            # decode的batch size，最大为steps的值
+            decode_image_batch = gr.Number(value=10, precision=0, label="Decode Image Batch (<steps)", interactive=True)
         with gr.Row():
             image_video = gr.Video(label="Real-time Image Output")
             mask_video = gr.Video(label="Real-time Mask Output")
@@ -353,6 +362,7 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
             original_image = gr.Image(label=f"Original Image", type="pil", interactive=False)
             edited_image = gr.Image(label=f"Output Image", type="pil", interactive=False)
             mix_result_with_red_mask = gr.Image(label=f"Mix Image With Red Mask", type="pil", interactive=False)
+            
     
     with gr.Row():
         gr.Examples(
@@ -374,13 +384,15 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
             randomize_cfg,
             text_cfg_scale,
             image_cfg_scale,
+            weather_hide_video,
+            decode_image_batch
         ],
         outputs=[seed, text_cfg_scale, image_cfg_scale, edited_image, mix_image, edited_mask, mask_video, image_video, original_image, mix_result_with_red_mask],
     )
     reset_button.click(
         fn=reset,
         inputs=[],
-        outputs=[steps, randomize_seed, seed, randomize_cfg, text_cfg_scale, image_cfg_scale, edited_image, mix_image, edited_mask, mask_video, image_video, original_image, mix_result_with_red_mask],
+        outputs=[steps, randomize_seed, seed, randomize_cfg, text_cfg_scale, image_cfg_scale, edited_image, mix_image, edited_mask, mask_video, image_video, original_image, mix_result_with_red_mask, weather_hide_video, decode_image_batch],
     )