Optimize video generation by adding torch.no_grad() context to reduce memory usage

app.py

@@ -198,23 +198,24 @@ def generate_video_from_text(
     def gradio_progress_callback(self, step, timestep, kwargs):
         progress((step + 1) / num_inference_steps)
 
-    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
+    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")
     print(images.shape)
@@ -268,23 +269,24 @@ def generate_video_from_image(
     def gradio_progress_callback(self, step, timestep, kwargs):
         progress((step + 1) / num_inference_steps)
 
-    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
+    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()