force clean

app.py

@@ -257,7 +257,7 @@ 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)
+    media_items = load_image_to_tensor_with_resize(image_path, height, width).to(device).detach()
 
     sample = {
         "prompt": prompt,
@@ -271,41 +271,42 @@ def generate_video_from_image(
 
     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
-
-    output_path = tempfile.mktemp(suffix=".mp4")
-    video_np = images.squeeze(0).permute(1, 2, 3, 0).cpu().float().numpy()
-    video_np = (video_np * 255).astype(np.uint8)
-    height, width = video_np.shape[1:3]
-    out = cv2.VideoWriter(
-        output_path, cv2.VideoWriter_fourcc(*"mp4v"), frame_rate, (width, height)
-    )
-    for frame in video_np[..., ::-1]:
-        out.write(frame)
-    out.release()
-
-    # Explicitly delete tensors and clear cache
-    del images
-    del video_np
-    torch.cuda.empty_cache()
+    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
+
+        output_path = tempfile.mktemp(suffix=".mp4")
+        video_np = images.squeeze(0).permute(1, 2, 3, 0).cpu().float().numpy()
+        video_np = (video_np * 255).astype(np.uint8)
+        height, width = video_np.shape[1:3]
+        out = cv2.VideoWriter(
+            output_path, cv2.VideoWriter_fourcc(*"mp4v"), frame_rate, (width, height)
+        )
+        for frame in video_np[..., ::-1]:
+            out.write(frame)
+        out.release()
+    finally:
+        del media_items
+        del images
+        del video_np
+        gc.collect()
+        torch.cuda.empty_cache()
     
     return output_path