Fix GPU context for Zero GPU

README.md

@@ -7,7 +7,7 @@ sdk: gradio
 sdk_version: 5.42.0
 app_file: app.py
 pinned: false
-suggested_hardware: a10g-small
+suggested_hardware: zero-a10g
 ---
 
 # LTX Video Studio

app.py

@@ -152,21 +152,6 @@ input[type="range"] {
 }
 """
 
-# ============================================
-# MODEL LOADING
-# ============================================
-pipe = None
-
-def load_model():
-    global pipe
-    if pipe is None:
-        pipe = LTXPipeline.from_pretrained(
-            "Lightricks/LTX-Video",
-            torch_dtype=torch.bfloat16
-        )
-        pipe.to("cuda")
-    return pipe
-
 # ============================================
 # VIDEO GENERATION FUNCTION
 # ============================================
@@ -188,32 +173,39 @@ def generate_video(
         return None
 
     progress(0, desc="Loading model...")
-    pipeline = load_model()
+
+    # Load model inside GPU context
+    pipe = LTXPipeline.from_pretrained(
+        "Lightricks/LTX-Video",
+        torch_dtype=torch.bfloat16
+    )
+    pipe.to("cuda")
+    pipe.enable_model_cpu_offload()
 
     # Handle seed
     if seed == -1:
         seed = random.randint(0, 2**32 - 1)
 
-    generator = torch.Generator(device="cuda").manual_seed(seed)
+    generator = torch.Generator(device="cuda").manual_seed(int(seed))
 
     progress(0.1, desc="Generating video frames...")
 
     # Generate video
-    output = pipeline(
+    output = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt if negative_prompt else None,
-        num_frames=num_frames,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
+        num_frames=int(num_frames),
+        guidance_scale=float(guidance_scale),
+        num_inference_steps=int(num_inference_steps),
         generator=generator,
-        height=height,
-        width=width,
+        height=int(height),
+        width=int(width),
     )
 
     progress(0.9, desc="Exporting video...")
 
     # Export to video file
-    video_path = export_to_video(output.frames[0], fps=fps)
+    video_path = export_to_video(output.frames[0], fps=int(fps))
 
     progress(1.0, desc="Complete!")