Update app.py

app.py

@@ -1,59 +1,79 @@
 import torch
-import gradio as gr
 from diffusers import CogVideoXPipeline
 from diffusers.utils import export_to_video
+import gradio as gr
 
-# 1. Initialize the CogVideoX1.5‑5B‑SAT pipeline on ZeroGPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# —————————————————————————————————————————————————————————————————————————
+# 1️⃣ Load & optimize the CogVideoX pipeline
+# —————————————————————————————————————————————————————————————————————————
+device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe = CogVideoXPipeline.from_pretrained(
-    "Wan-AI/Wan2.1-T2V-14B",      # High‑quality 10s model :contentReference[oaicite:5]{index=5}
-    torch_dtype=torch.float16
-).to(device)
-pipe.enable_model_cpu_offload()        # Memory optimization :contentReference[oaicite:6]{index=6}
-pipe.enable_vae_slicing()              # Further memory savings :contentReference[oaicite:7]{index=7}
-
-# 2. Define the video generation function
-def generate_video(prompt: str):
+    "THUDM/CogVideoX1.5-5B",
+    torch_dtype=torch.bfloat16
+).to(device)                                                       # :contentReference[oaicite:5]{index=5}
+
+# 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
+# —————————————————————————————————————————————————————————————————————————
+def generate_video(prompt: str) -> str:
     """
-    Generates a 10‑second video (160 frames at 16 fps) for the given prompt.
-    Returns the path to the saved MP4 file.
+    Generates a ~10s MP4 video (161 frames @16 FPS) for the given prompt.
+    Returns the filepath to the saved video.
     """
-    # 160 frames @16 fps → 10-second duration :contentReference[oaicite:8]{index=8}
-    result = pipe(
-        prompt,
-        num_frames=160,
-        fps=16,
+    # Inference settings
+    guidance_scale   = 6.0
+    num_inference_steps = 50
+    num_frames       = 161    # yields ~10s @16 FPS        :contentReference[oaicite:8]{index=8}
+    fps              = 16
+
+    # Run the pipeline
+    output = pipe(
+        prompt=prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        num_frames=num_frames,
+        fps=fps,
     )
-    frames = result.frames[0]
-    
-    # Export frames to MP4; Gradio's Video component accepts file paths
-    video_path = export_to_video(frames, "generated.mp4", fps=16)
-    return video_path
-
-# 3. Build the Gradio interface
-with gr.Blocks(title="CogVideoX1.5‑5B‑SAT Text‑to‑Video") as demo:
+    frames = output.frames[0]  # batch_size=1
+
+    # 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
+
+# —————————————————————————————————————————————————————————————————————————
+# 3️⃣ Gradio UI
+# —————————————————————————————————————————————————————————————————————————
+with gr.Blocks(title="CogVideoX Text‑to‑Video Demo") as demo:
     gr.Markdown(
         """
-        # 🎥 Text‑to‑Video with CogVideoX1.5‑5B‑SAT
-        Enter a prompt and generate up to **10 s** of high‑quality video.
-        Powered by Wan-AI/Wan2.1-T2V-14B on Hugging Face ZeroGPU.
+        # 🎞️ CogVideoX1.5‑5B Text→Video  
+        Enter a descriptive prompt and generate **up to 10 s** of video.
         """
     )
     with gr.Row():
         prompt_input = gr.Textbox(
-            label="Video Prompt",
-            placeholder="e.g., A majestic eagle soaring over mountain peaks",
+            label="Prompt",
+            placeholder="e.g., A vibrant hot air balloon festival at sunrise",
             lines=2,
         )
-        generate_btn = gr.Button("Generate Video")
-    video_output = gr.Video(label="Your Generated Video")
-    
-    generate_btn.click(
+        gen_button = gr.Button("Generate")
+    video_output = gr.Video(
+        label="Generated Video",
+        format="mp4"                  # ensure browser‑playable :contentReference[oaicite:10]{index=10}
+    )
+
+    gen_button.click(
         fn=generate_video,
         inputs=prompt_input,
         outputs=video_output,
     )
 
-# 4. Launch (Spaces will handle server settings automatically)
+# —————————————————————————————————————————————————————————————————————————
+# 4️⃣ Launch (Spaces auto‑handles server config & auth)
+# —————————————————————————————————————————————————————————————————————————
 if __name__ == "__main__":
     demo.launch()