Update app_gradio.py

app_gradio.py

@@ -117,47 +117,67 @@ pipe = TextToVideoSDPipelineModded.from_pretrained(
 @torch.no_grad()
 def process_video(num_frames, num_seeds, generator, exp_dir, load_name, caption, lambda_):
     pipe_inversion.to(device)
-    id_latents = invert(pipe_inversion, inv, load_name).to(device, dtype=dtype)
+    try:
+        id_latents = invert(pipe_inversion, inv, load_name).to(device, dtype=dtype)
+    except Exception as e:
+        torch.cuda.empty_cache()  # Clear CUDA cache in case of failure
+        gc.collect()
+        raise gr.Error(f"Invert latents failed: {str(e)}") from e
+    
     latents = id_latents.repeat(num_seeds, 1, 1, 1, 1)
     generator = [torch.Generator(device="cuda").manual_seed(i) for i in range(num_seeds)]
-    video_frames = pipe(
-        prompt=caption,
-        negative_prompt="",
-        num_frames=num_frames,
-        num_inference_steps=25,
-        inv_latents=latents,
-        guidance_scale=9,
-        generator=generator,
-        lambda_=lambda_,
-    ).frames
-
+    
+    try:
+        video_frames = pipe(
+            prompt=caption,
+            negative_prompt="",
+            num_frames=num_frames,
+            num_inference_steps=25,
+            inv_latents=latents,
+            guidance_scale=9,
+            generator=generator,
+            lambda_=lambda_,
+        ).frames
+    except Exception as e:
+        torch.cuda.empty_cache()
+        gc.collect()
+        raise RuntimeError(f"Failed to process video: {e}") from e
+    
     gifs = []
-    for seed in range(num_seeds):
-        vid_name = f"{exp_dir}/mp4_logs/vid_{os.path.basename(load_name)[:-4]}-rand{seed}.mp4"
-        gif_name = f"{exp_dir}/gif_logs/vid_{os.path.basename(load_name)[:-4]}-rand{seed}.gif"
-        
-        os.makedirs(os.path.dirname(vid_name), exist_ok=True)
-        os.makedirs(os.path.dirname(gif_name), exist_ok=True)
-        
-        video_path = export_to_video(video_frames[seed], output_video_path=vid_name)
-        VideoFileClip(vid_name).write_gif(gif_name)
-        
-        with Image.open(gif_name) as im:
-            frames = load_frames(im)
+    try:
+        for seed in range(num_seeds):
+            vid_name = f"{exp_dir}/mp4_logs/vid_{os.path.basename(load_name)[:-4]}-rand{seed}.mp4"
+            gif_name = f"{exp_dir}/gif_logs/vid_{os.path.basename(load_name)[:-4]}-rand{seed}.gif"
+            
+            os.makedirs(os.path.dirname(vid_name), exist_ok=True)
+            os.makedirs(os.path.dirname(gif_name), exist_ok=True)
+            
+            video_path = export_to_video(video_frames[seed], output_video_path=vid_name)
+            VideoFileClip(vid_name).write_gif(gif_name)
+            
+            with Image.open(gif_name) as im:
+                frames = load_frames(im)
 
-        frames_collect = np.empty((0, 1024, 1024), int)
-        for frame in frames:
-            frame = cv2.resize(frame, (1024, 1024))[:, :, :3]
-            frame = cv2.cvtColor(255 - frame, cv2.COLOR_RGB2GRAY)
-            _, frame = cv2.threshold(255 - frame, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-            frames_collect = np.append(frames_collect, [frame], axis=0)
+            frames_collect = np.empty((0, 1024, 1024), int)
+            for frame in frames:
+                frame = cv2.resize(frame, (1024, 1024))[:, :, :3]
+                frame = cv2.cvtColor(255 - frame, cv2.COLOR_RGB2GRAY)
+                _, frame = cv2.threshold(255 - frame, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+                frames_collect = np.append(frames_collect, [frame], axis=0)
 
-        save_gif(frames_collect, gif_name)
-        gifs.append(gif_name)
+            save_gif(frames_collect, gif_name)
+            gifs.append(gif_name)
+    except Exception as e:
+        torch.cuda.empty_cache()
+        raise RuntimeError(f"Failed during GIF generation: {e}") from e
 
     return gifs
 
 def generate_output(image, apply_filter, prompt: str, num_seeds: int = 3, lambda_value: float = 0.5, progress=gr.Progress(track_tqdm=True)) -> List[str]:
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+    
     if prompt is None:
         raise gr.Error("You forgot to describe the motion !")
     """Main function to generate output GIFs"""
@@ -175,21 +195,17 @@ def generate_output(image, apply_filter, prompt: str, num_seeds: int = 3, lambda
 
     image.save(temp_image_path)
     
-    try:
-        # Attempt to process video
-        generated_gifs = process_video(
-            num_frames=10,
-            num_seeds=num_seeds,
-            generator=None,
-            exp_dir=exp_dir,
-            load_name=temp_image_path,
-            caption=prompt,
-            lambda_=1 - lambda_value
-        )
-    except Exception as e:
-        torch.cuda.empty_cache()  # Clear CUDA cache in case of failure
-        gc.collect()
-        raise gr.Error(f"Video processing failed: {str(e)}") from e
+    
+    # Attempt to process video
+    generated_gifs = process_video(
+        num_frames=10,
+        num_seeds=num_seeds,
+        generator=None,
+        exp_dir=exp_dir,
+        load_name=temp_image_path,
+        caption=prompt,
+        lambda_=1 - lambda_value
+    )
     
     if apply_filter:
         try: