envs

app.py

@@ -20,13 +20,13 @@ import spaces
 from PIL import Image
 from omegaconf import OmegaConf
 from einops import rearrange, repeat
-from torchvision import transforms
+from torchvision import transforms,utils
 from transformers import CLIPTextModel, CLIPTokenizer
 from diffusers import AutoencoderKL, DDIMScheduler
 
 from pipelines.pipeline_imagecoductor import ImageConductorPipeline
 from modules.unet import UNet3DConditionFlowModel
-from utils.gradio_utils import ensure_dirname, split_filename, visualize_drag, image2pil, image2arr
+from utils.gradio_utils import ensure_dirname, split_filename, visualize_drag, image2pil, save_videos_grid
 from utils.utils import create_image_controlnet, create_flow_controlnet, interpolate_trajectory, load_weights, load_model, bivariate_Gaussian
 from utils.lora_utils import add_LoRA_to_controlnet
 from utils.visualizer import Visualizer, vis_flow_to_video
@@ -382,9 +382,12 @@ class ImageConductor:
                     eval_mode = True,
                 ).videos
 
-        outputs_path = os.path.join(output_dir, f'output_{i}_{id}.mp4')
-        vis_video = (rearrange(sample[0], 'c t h w -> t h w c') * 255.).clip(0, 255)
-        torchvision.io.write_video(outputs_path, vis_video, fps=8, video_codec='h264', options={'crf': '10'})
+        # outputs_path = os.path.join(output_dir, f'output_{i}_{id}.mp4')
+        # vis_video = (rearrange(sample[0], 'c t h w -> t h w c') * 255.).clip(0, 255)
+        # torchvision.io.write_video(outputs_path, vis_video, fps=8, video_codec='h264', options={'crf': '10'})
+        
+        outputs_path = os.path.join(output_dir, f'output_{i}_{id}.gif')
+        save_videos_grid(sample[0][None], outputs_path)
 
         return {output_image: visualized_drag, output_video: outputs_path}
 
@@ -619,10 +622,13 @@ with block:
                     examples_type = gr.Textbox(label="Examples Type (Ignore) ",  value="", visible=False)
 
         with gr.Column(scale=7):
-            output_video = gr.Video(
-                                    label="Output Video", 
-                                    width=384, 
-                                    height=256)
+            # output_video = gr.Video(
+            #                         label="Output Video", 
+            #                         width=384, 
+            #                         height=256)
+            output_video = gr.Image(label="Output Video",
+                                            height=256,
+                                            width=384,)
             
             
     with gr.Row():