Update app.py

app.py

@@ -1,86 +1,89 @@
 import torch
+import spaces
 from diffusers import CogVideoXPipeline
 from diffusers.utils import export_to_video
 import gradio as gr
 
-# —————————————————————————————————————————————————————————————————————————
-# 1️⃣ Load & optimize the CogVideoX pipeline
-# —————————————————————————————————————————————————————————————————————————
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# ------------------------------------------------------------------------------
+# 1. Load and optimize the CogVideoX pipeline on CPU by default
+# ------------------------------------------------------------------------------
+device = torch.device("cpu")
 pipe = CogVideoXPipeline.from_pretrained(
     "THUDM/CogVideoX1.5-5B",
-    torch_dtype=torch.bfloat16
-).to(device)                                                       # :contentReference[oaicite:5]{index=5}
+).to(device)
+# Memory optimizations for ZeroGPU or limited GPU environments
+pipe.enable_model_cpu_offload()
+pipe.vae.enable_slicing()
 
-# Memory optimizations for A100 (ZeroGPU) or any limited‑GPU environment
-pipe.enable_model_cpu_offload()                                    # ⬇️ from ~33 GB → ~19 GB :contentReference[oaicite:6]{index=6}
-pipe.vae.enable_slicing()                                          # ⬇️ further to ~11 GB    :contentReference[oaicite:7]{index=7}
-
-# —————————————————————————————————————————————————————————————————————————
-# 2️⃣ Video generation function
-# —————————————————————————————————————————————————————————————————————————
+# ------------------------------------------------------------------------------
+# 2. Decorated GPU function for ZeroGPU
+# ------------------------------------------------------------------------------
+@spaces.GPU(duration=180)  # request up to 180s of GPU time
 def generate_video(prompt: str) -> str:
     """
-    Generates a ~10s MP4 video (161 frames @16 FPS) for the given prompt.
+    Generates a ~10s MP4 video (161 frames @16 FPS) from the given prompt.
     Returns the filepath to the saved video.
     """
-    # Inference settings
-    guidance_scale   = 6.0
+    # Move pipeline to GPU for inference
+    pipe.to("cuda")
+    # Inference parameters
+    guidance_scale = 6.0
     num_inference_steps = 50
-    num_frames       = 161    # yields ~10s @16 FPS        :contentReference[oaicite:8]{index=8}
-    fps              = 16
+    num_frames = 161  # ~10s at 16 FPS
+    fps = 16
 
-    # Run the pipeline
-    output = pipe(
+    # Run the diffusion pipeline
+    result = pipe(
         prompt=prompt,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         num_frames=num_frames,
         fps=fps,
     )
-    frames = output.frames[0]  # batch_size=1
+    frames = result.frames[0]
+
+    # Move pipeline back to CPU to free GPU memory
+    pipe.to("cpu")
 
-    # Convert frames → MP4 (H.264) for browser playability :contentReference[oaicite:9]{index=9}
-    file_path = export_to_video(frames, "generated.mp4", fps=fps)
-    return file_path
+    # Export frames to MP4 for Gradio
+    return export_to_video(frames, "generated.mp4", fps=fps)
 
-# —————————————————————————————————————————————————————————————————————————
-# 3️⃣ Gradio UI
-# —————————————————————————————————————————————————————————————————————————
-with gr.Blocks(title="CogVideoX Text‑to‑Video Demo") as demo:
+# ------------------------------------------------------------------------------
+# 3. Define the Gradio interface
+# ------------------------------------------------------------------------------
+with gr.Blocks(title="CogVideoX Text-to-Video with ZeroGPU") as demo:
     gr.Markdown(
         """
-        # 🎞️ CogVideoX1.5‑5B Text→Video  
-        Enter a descriptive prompt and generate **up to 10 s** of video.
+        ## 🎞️ Text-to-Video Generator
+        Enter a descriptive prompt and generate up to 10s of video
+        powered by CogVideoX1.5-5B on Hugging Face ZeroGPU.
         """
     )
     with gr.Row():
         prompt_input = gr.Textbox(
             label="Prompt",
-            placeholder="e.g., A vibrant hot air balloon festival at sunrise",
+            placeholder="e.g., A majestic eagle soaring over mountain peaks",
             lines=2,
         )
-        gen_button = gr.Button("Generate")
+        generate_btn = gr.Button("Generate Video")
     video_output = gr.Video(
         label="Generated Video",
-        format="mp4"                  # ensure browser‑playable :contentReference[oaicite:10]{index=10}
+        format="mp4"
     )
 
-    gen_button.click(
+    generate_btn.click(
         fn=generate_video,
         inputs=prompt_input,
         outputs=video_output,
     )
 
-# —————————————————————————————————————————————————————————————————————————
-# 4️⃣ Launch (Spaces auto‑handles server config & auth)
-# —————————————————————————————————————————————————————————————————————————
+# ------------------------------------------------------------------------------
+# 4. Launch the Gradio app
+# ------------------------------------------------------------------------------
 if __name__ == "__main__":
     demo.launch(
-        server_name="0.0.0.0",   # listen externally
-        server_port=7860,        # default HF Spaces port
-        ssr_mode=False,          # disable SSR to block the main thread
-        share=True,           # optional for local tunnels
-        debug=True               # optional debug output
+        server_name="0.0.0.0",
+        server_port=7860,
+        ssr_mode=False,
+        debug=True,
     )
-