Switch to 1.3B model for faster generation

app.py

@@ -9,11 +9,12 @@ from diffusers import AutoencoderKLWan, WanPipeline
 from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
 
 # ── Load model once at startup ──────────────────────────────────────────
-MODEL_ID = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
+# 1.3B model — fast and fits easily in A10G (24GB)
+MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 
 vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
 pipe = WanPipeline.from_pretrained(MODEL_ID, vae=vae, torch_dtype=torch.bfloat16)
-pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=5.0)
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=3.0)
 pipe.to("cuda")
 
 # ── Resolution presets ──────────────────────────────────────────────────
@@ -21,12 +22,10 @@ RESOLUTIONS = {
     "9:16 (480p)": (480, 832),
     "16:9 (480p)": (832, 480),
     "1:1 (480p)": (624, 624),
-    "9:16 (720p)": (720, 1280),
-    "16:9 (720p)": (1280, 720),
 }
 
 # ── Generation function ─────────────────────────────────────────────────
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=180)
 def generate(prompt, negative_prompt, resolution, steps, guidance, seed, progress=gr.Progress(track_tqdm=True)):
     width, height = RESOLUTIONS.get(resolution, (480, 832))
 
@@ -35,8 +34,8 @@ def generate(prompt, negative_prompt, resolution, steps, guidance, seed, progres
     output = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt or None,
-        num_inference_steps=steps,
-        guidance_scale=guidance,
+        num_inference_steps=int(steps),
+        guidance_scale=float(guidance),
         height=height,
         width=width,
         num_frames=81,
@@ -52,7 +51,7 @@ def generate(prompt, negative_prompt, resolution, steps, guidance, seed, progres
 
 # ── Gradio UI ────────────────────────────────────────────────────────────
 with gr.Blocks(title="Wan 2.1 Video Generator") as demo:
-    gr.Markdown("# Wan 2.1 T2V (14B) — ZeroGPU")
+    gr.Markdown("# Wan 2.1 T2V (1.3B) — ZeroGPU")
 
     with gr.Row():
         with gr.Column(scale=3):