Enhance video generation error handling and memory management

app.py

@@ -17,6 +17,7 @@ import cv2
 from PIL import Image
 import tempfile
 import os
+import gc
 
 # Load Hugging Face token if needed
 hf_token = os.getenv("HF_TOKEN")
@@ -198,24 +199,34 @@ def generate_video_from_text(
     def gradio_progress_callback(self, step, timestep, kwargs):
         progress((step + 1) / num_inference_steps)
 
-    with torch.no_grad():
-        images = pipeline(
-            num_inference_steps=num_inference_steps,
-            num_images_per_prompt=1,
-            guidance_scale=guidance_scale,
-            generator=generator,
-            output_type="pt",
-            height=height,
-            width=width,
-            num_frames=num_frames,
-            frame_rate=frame_rate,
-            **sample,
-            is_video=True,
-            vae_per_channel_normalize=True,
-            conditioning_method=ConditioningMethod.FIRST_FRAME,
-            mixed_precision=True,
-            callback_on_step_end=gradio_progress_callback,
-        ).images
+    try:
+        with torch.no_grad():
+            images = pipeline(
+                num_inference_steps=num_inference_steps,
+                num_images_per_prompt=1,
+                guidance_scale=guidance_scale,
+                generator=generator,
+                output_type="pt",
+                height=height,
+                width=width,
+                num_frames=num_frames,
+                frame_rate=frame_rate,
+                **sample,
+                is_video=True,
+                vae_per_channel_normalize=True,
+                conditioning_method=ConditioningMethod.FIRST_FRAME,
+                mixed_precision=True,
+                callback_on_step_end=gradio_progress_callback,
+            ).images
+    except Exception as e:
+        raise gr.Error(
+            f"An error occurred while generating the video. Please try again. Error: {e}",
+            duration=5,
+        )
+    finally:
+        pipeline.to("cpu")
+        torch.cuda.empty_cache()
+        gc.collect()
 
     output_path = tempfile.mktemp(suffix=".mp4")
     print(images.shape)
@@ -257,7 +268,9 @@ def generate_video_from_image(
     if not image_path:
         raise gr.Error("Please provide an input image.", duration=5)
 
-    media_items = load_image_to_tensor_with_resize(image_path, height, width).to(device).detach()
+    media_items = (
+        load_image_to_tensor_with_resize(image_path, height, width).to(device).detach()
+    )
 
     sample = {
         "prompt": prompt,
@@ -271,6 +284,7 @@ def generate_video_from_image(
 
     def gradio_progress_callback(self, step, timestep, kwargs):
         progress((step + 1) / num_inference_steps)
+
     try:
         with torch.no_grad():
             images = pipeline(
@@ -301,13 +315,17 @@ def generate_video_from_image(
         for frame in video_np[..., ::-1]:
             out.write(frame)
         out.release()
+    except Exception as e:
+        raise gr.Error(
+            f"An error occurred while generating the video. Please try again. Error: {e}",
+            duration=5,
+        )
+
     finally:
-        del media_items
-        del images
-        del video_np
-        gc.collect()
+        pipeline.to("cpu")
         torch.cuda.empty_cache()
-    
+        gc.collect()
+
     return output_path