Update app.py

app.py

@@ -2,159 +2,47 @@ import spaces
 import torch
 import gc
 import gradio as gr
-from diffusers import AutoModel, WanPipeline
+from diffusers import WanPipeline
 from diffusers.utils import export_to_video
 import tempfile
 import time
 
-# ============================================================
-#  SHOTARCH VIDEO GEN — Wan2.1-1.3B (Fully Optimized)
-# ============================================================
-#  Optimizations applied:
-#   1. bfloat16 transformer (native on Ampere/Hopper GPUs)
-#   2. float32 VAE (required for sharp decode)
-#   3. VAE tiling (low peak VRAM during decode)
-#   4. torch.inference_mode (faster than no_grad)
-#   5. Pre-loaded on CPU → instant GPU transfer
-# ============================================================
-
-print("📦 Loading Wan2.1-1.3B on CPU (one-time)...")
-load_start = time.time()
-
-vae = AutoModel.from_pretrained(
-    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-    subfolder="vae",
-    torch_dtype=torch.float32,
-)
-
-pipe = WanPipeline.from_pretrained(
-    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-    vae=vae,
-    torch_dtype=torch.bfloat16,
-)
-
-print(f"✅ Model loaded in {time.time()-load_start:.0f}s")
-
+pipe = None
 
 @spaces.GPU(duration=240)
-def generate_video(
-    prompt,
-    negative_prompt,
-    num_frames,
-    height,
-    width,
-    num_inference_steps,
-    guidance_scale,
-):
-    """Generate video with Wan2.1 on ZeroGPU."""
-
-    # Move to GPU
-    pipe.to("cuda")
-    pipe.vae.enable_tiling()
-
-    print(f"🎥 Generating: {width}x{height}, {num_frames} frames, {num_inference_steps} steps")
-    start = time.time()
-
+def generate_video(prompt, negative_prompt, num_frames, height, width, num_inference_steps, guidance_scale):
+    global pipe
+    if pipe is None:
+        print("📦 Loading Wan2.1-1.3B to GPU...")
+        pipe = WanPipeline.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers", torch_dtype=torch.float16, low_cpu_mem_usage=True)
+        pipe.to("cuda")
+        pipe.vae.enable_tiling()
+        print("✅ Loaded!")
+    
     with torch.inference_mode():
-        result = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_frames=int(num_frames),
-            height=int(height),
-            width=int(width),
-            num_inference_steps=int(num_inference_steps),
-            guidance_scale=float(guidance_scale),
-        ).frames[0]
-
-    elapsed = time.time() - start
-    print(f"✅ Generated in {elapsed:.1f}s")
-
-    # Save video
+        result = pipe(prompt=prompt, negative_prompt=negative_prompt, num_frames=int(num_frames), height=int(height), width=int(width), num_inference_steps=int(num_inference_steps), guidance_scale=float(guidance_scale)).frames[0]
+    
     output_path = tempfile.mktemp(suffix=".mp4")
     export_to_video(result, output_path, fps=16)
-
-    # Cleanup GPU memory
-    gc.collect()
-    torch.cuda.empty_cache()
-
+    gc.collect(); torch.cuda.empty_cache()
     return output_path
 
-
-# ============================================================
-#  GRADIO UI
-# ============================================================
-css = """
-#main-container { max-width: 1200px; margin: auto; }
-.generate-btn { height: 50px !important; font-size: 18px !important; }
-"""
-
-with gr.Blocks(title="Shotarch Video Gen", css=css, theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # 🎬 Shotarch Video Generator
-        ### Powered by Wan2.1-1.3B — Ultra-Fast AI Video Generation
-        """
-    )
-
+with gr.Blocks(title="Shotarch Video Gen", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🎬 Shotarch Video Generator\n### Wan2.1-1.3B on ZeroGPU")
     with gr.Row():
-        with gr.Column(scale=1):
-            prompt = gr.Textbox(
-                label="✏️ Prompt",
-                lines=4,
-                placeholder="A cinematic slow-motion shot of a futuristic cyberpunk sports car drifting through a neon-lit rain-soaked city at night...",
-            )
-            negative_prompt = gr.Textbox(
-                label="🚫 Negative Prompt",
-                lines=2,
-                value="Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
-            )
-
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Describe your video...")
+            negative_prompt = gr.Textbox(label="Negative Prompt", lines=2, value="Bright tones, overexposed, static, blurred details, worst quality, low quality, ugly, deformed, still picture")
             with gr.Row():
-                width = gr.Slider(
-                    minimum=480, maximum=1280, value=1280, step=16, label="Width"
-                )
-                height = gr.Slider(
-                    minimum=320, maximum=720, value=720, step=16, label="Height"
-                )
-
+                width = gr.Slider(480, 1280, value=1280, step=16, label="Width")
+                height = gr.Slider(320, 720, value=720, step=16, label="Height")
             with gr.Row():
-                num_frames = gr.Slider(
-                    minimum=17,
-                    maximum=81,
-                    value=81,
-                    step=4,
-                    label="Frames (81 = 5 sec @ 16fps)",
-                )
-                steps = gr.Slider(
-                    minimum=10,
-                    maximum=50,
-                    value=25,
-                    step=1,
-                    label="Inference Steps",
-                )
-
-            guidance = gr.Slider(
-                minimum=1.0,
-                maximum=15.0,
-                value=5.0,
-                step=0.5,
-                label="Guidance Scale",
-            )
-
-            generate_btn = gr.Button(
-                "🎬 Generate Video",
-                variant="primary",
-                elem_classes="generate-btn",
-            )
-
-        with gr.Column(scale=1):
-            output_video = gr.Video(label="🎥 Generated Video")
-
-    # API-friendly: this function is also callable via /api/predict
-    generate_btn.click(
-        fn=generate_video,
-        inputs=[prompt, negative_prompt, num_frames, height, width, steps, guidance],
-        outputs=output_video,
-    )
+                num_frames = gr.Slider(17, 81, value=81, step=4, label="Frames (81=5sec)")
+                steps = gr.Slider(10, 50, value=25, step=1, label="Steps")
+            guidance = gr.Slider(1.0, 15.0, value=5.0, step=0.5, label="Guidance Scale")
+            btn = gr.Button("🎬 Generate Video", variant="primary")
+        with gr.Column():
+            output = gr.Video(label="Generated Video")
+    btn.click(fn=generate_video, inputs=[prompt, negative_prompt, num_frames, height, width, steps, guidance], outputs=output)
 
 demo.launch()