return state for examples

app-huggingface.py

@@ -37,7 +37,8 @@ def add_text(gif_file: str, mp4_file: str, text_index: int):
     except EOFError:
         pass
 
-    text = ['Happy New Year!', 'Happy New Year!', 'Happy New Year!', 'Happy New Year!'][text_index]
+    text = ['Happy New Year!', 'Happy New Year!',
+            'Happy New Year!', 'Happy New Year!'][text_index]
     size = [36, 36, 36, 36][text_index]
 
     for i, frame in enumerate(frames):
@@ -55,6 +56,7 @@ def add_text(gif_file: str, mp4_file: str, text_index: int):
     mp4_file = convert_gif_to_mp4(gif_file, mp4_file)
     return gif_file, mp4_file
 
+
 css = """
 .toolbutton {
     margin-buttom: 0em 0em 0em 0em;
@@ -255,6 +257,83 @@ class AnimateController:
         return n_prompt, ip_adapter_scale
 
     def animate(
+        self,
+        init_img,
+        motion_scale,
+        prompt_textbox,
+        negative_prompt_textbox,
+        sample_step_slider,
+        cfg_scale_slider,
+        seed_textbox,
+        ip_adapter_scale,
+        style,
+        progress=gr.Progress(),
+    ):
+
+        global sample_idx
+
+        if init_img is None:
+            gr.Warning('Please upload image or use example images.')
+
+        if seed_textbox != -1 and seed_textbox != "":
+            torch.manual_seed(int(seed_textbox))
+            seed = int(seed_textbox)
+        else:
+            seed = torch.initial_seed()
+        generator = torch.Generator(device='cuda')
+        generator.manual_seed(seed)
+        seed_everything(seed)
+
+        print(f'Seed: {seed}')
+
+        pipeline = self.pipeline_dict[style]
+        init_img, h, w = preprocess_img(init_img)
+        print(f'img size: {h, w}')
+
+        sample = pipeline(
+            image=init_img,
+            prompt=prompt_textbox,
+            negative_prompt=negative_prompt_textbox,
+            generator=generator,
+            num_inference_steps=sample_step_slider,
+            guidance_scale=cfg_scale_slider,
+            width=w,
+            height=h,
+            video_length=16,
+            mask_sim_template_idx=motion_scale - 1,
+            ip_adapter_scale=ip_adapter_scale,
+            progress_fn=progress,
+        ).videos
+
+        save_sample_path_mp4 = os.path.join(
+            self.savedir_sample, f"{sample_idx}.mp4")
+        save_sample_path_gif = os.path.join(
+            self.savedir_sample, f"{sample_idx}.gif")
+
+        save_videos_grid(sample, save_sample_path_mp4)
+        save_videos_grid(sample, save_sample_path_gif)
+
+        sample_config = {
+            "prompt": prompt_textbox,
+            "n_prompt": negative_prompt_textbox,
+            "num_inference_steps": sample_step_slider,
+            "guidance_scale": cfg_scale_slider,
+            "width": w,
+            "height": h,
+            "seed": seed,
+            "motion": motion_scale,
+        }
+        print(sample_config)
+        json_str = json.dumps(sample_config, indent=4)
+        with open(os.path.join(self.savedir, "logs.json"), "a") as f:
+            f.write(json_str)
+            f.write("\n\n")
+
+        sample_idx += 1
+        return (save_sample_path_mp4,
+                [save_sample_path_mp4, save_sample_path_gif])
+
+    def animate_example(
         self,
         init_img,
         motion_scale,
@@ -269,17 +348,11 @@ class AnimateController:
         with_text=False,
         text_idx=0,
         progress=gr.Progress(),
-    ):
 
+    ):
         global sample_idx
-        print(example_img)
-
         if init_img is None:
-            if example_img is None:
-                gr.Warning('Please upload image or use example images.')
-            else:
-                gr.Info('Use exmaple image for fast try!')
-                init_img = np.array(Image.open(example_img))
+            gr.Warning('Please upload image or use example images.')
 
         if seed_textbox != -1 and seed_textbox != "":
             torch.manual_seed(int(seed_textbox))
@@ -339,11 +412,13 @@ class AnimateController:
             add_text(save_sample_path_gif, save_sample_path_mp4, text_idx)
 
         sample_idx += 1
-        # return (save_sample_path_mp4,
-        #         gr.update(link=f'file/{save_sample_path_mp4}'),
-        #         gr.update(link=f'file/{save_sample_path_gif}'))
         return (save_sample_path_mp4,
-                [save_sample_path_mp4, save_sample_path_gif])
+                [save_sample_path_mp4, save_sample_path_gif],
+                example_img,
+                seed,
+                motion_scale,
+                cfg_scale_slider,
+                )
 
 
 controller = AnimateController()
@@ -472,7 +547,8 @@ def ui():
                 result_video = gr.Video(
                     label="Generated Animation", interactive=False)
                 with gr.Row():
-                    download = gr.Files(file_types=['gif', 'mp4'], label='Donwload Output')
+                    download = gr.Files(
+                        file_types=['gif', 'mp4'], label='Donwload Output')
 
         generate_button.click(
             fn=controller.animate,
@@ -486,12 +562,11 @@ def ui():
                 seed_textbox,
                 ip_adapter_scale,
                 style_dropdown,
-                # None,
             ],
             outputs=[result_video, download])
 
         default_1.click(
-            fn=controller.animate,
+            fn=controller.animate_example,
             inputs=[
                 init_img,
                 default_motion,
@@ -506,9 +581,16 @@ def ui():
                 with_wishes,
                 default_1_idx,
             ],
-            outputs=[result_video, download])
+            outputs=[
+                result_video,
+                download,
+                example_img,
+                default_seed,
+                default_motion,
+                default_cfg,
+            ])
         default_2.click(
-            fn=controller.animate,
+            fn=controller.animate_example,
             inputs=[
                 init_img,
                 default_motion,
@@ -523,9 +605,16 @@ def ui():
                 with_wishes,
                 default_2_idx,
             ],
-            outputs=[result_video, download])
+            outputs=[
+                result_video,
+                download,
+                example_img,
+                default_seed,
+                default_motion,
+                default_cfg,
+            ])
         default_3.click(
-            fn=controller.animate,
+            fn=controller.animate_example,
             inputs=[
                 init_img,
                 default_motion,
@@ -540,9 +629,16 @@ def ui():
                 with_wishes,
                 default_3_idx,
             ],
-            outputs=[result_video, download])
+            outputs=[
+                result_video,
+                download,
+                example_img,
+                default_seed,
+                default_motion,
+                default_cfg,
+            ])
         default_4.click(
-            fn=controller.animate,
+            fn=controller.animate_example,
             inputs=[
                 init_img,
                 default_motion,
@@ -557,13 +653,14 @@ def ui():
                 with_wishes,
                 default_4_idx,
             ],
-            outputs=[result_video, download])
-
-        # save_as_gif.click(
-        #     fn=None,
-        #     inputs=[save_as_gif],
-        #     js=_js_open_if_not_logged_in,
-        # )
+            outputs=[
+                result_video,
+                download,
+                example_img,
+                default_seed,
+                default_motion,
+                default_cfg,
+            ])
 
         def create_example(input_list):
             return gr.Examples(