Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 """
 Copyright NewGenAI
 Code can't be included in commercial app used for monetary gain. No derivative code allowed.
@@ -10,6 +11,7 @@ import time
 from datetime import datetime
 import os
 
+import torch
 from diffusers.utils import export_to_video
 from diffusers import LTXImageToVideoPipeline
 from transformers import T5EncoderModel, T5Tokenizer
@@ -144,20 +146,24 @@ pipe = LTXImageToVideoPipeline.from_single_file(
 )
 pipe.enable_model_cpu_offload()
 
+@spaces.GPU()
 def generate_video(image, prompt, negative_prompt, height, width, num_frames, num_inference_steps, fps, seed):
     if seed == 0:
         seed = random.randint(0, 999999)
-    
-    video = pipe(
-        image=image,
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        width=width,
-        height=height,
-        num_frames=num_frames,
-        num_inference_steps=num_inference_steps,
-        generator=torch.Generator(device='cuda').manual_seed(seed),
-    ).frames[0]
+
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+    with torch.inference_mode():
+        video = pipe(
+            image=image,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            width=width,
+            height=height,
+            num_frames=num_frames,
+            num_inference_steps=num_inference_steps,
+            generator=torch.Generator(device='cuda').manual_seed(seed),
+        ).frames[0]
 
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     filename = f"{prompt[:10]}_{timestamp}.mp4"