Create app.py

app.py

@@ -0,0 +1,137 @@
+import torch
+from diffusers import AutoencoderKLWan, WanPipeline
+from diffusers.utils import export_to_video
+import gradio as gr
+import tempfile
+import os
+import spaces
+from huggingface_hub import hf_hub_download
+
+# --- Global Model Loading (runs once when the script starts) ---
+MODEL_ID = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
+
+
+print(f"Loading VAE for {MODEL_ID}...")
+# Use float32 for VAE for stability, as recommended in some cases
+vae = AutoencoderKLWan.from_pretrained(
+    MODEL_ID,
+    subfolder="vae",
+    torch_dtype=torch.float32
+)
+print(f"Loading Pipeline {MODEL_ID}...")
+# Use bfloat16 for the main pipeline for memory efficiency and speed
+pipe = WanPipeline.from_pretrained(
+    MODEL_ID,
+    vae=vae,
+    torch_dtype=torch.bfloat16
+)
+print("Moving pipeline to CUDA...")
+pipe.to("cuda")
+causvid_path = hf_hub_download(repo_id="Kijai/WanVideo_comfy", filename="Wan21_CausVid_14B_T2V_lora_rank32.safetensors")
+pipe.load_lora_weights(causvid_path)
+
+# --- Gradio Interface Function ---
+@spaces.GPU
+def generate_video(prompt, negative_prompt, height, width, num_frames, guidance_scale, fps):
+
+    print("Starting video generation...")
+    print(f"  Prompt: {prompt}")
+    print(f"  Negative Prompt: {negative_prompt if negative_prompt else 'None'}")
+    print(f"  Height: {height}, Width: {width}")
+    print(f"  Num Frames: {num_frames}, FPS: {fps}")
+    print(f"  Guidance Scale: {guidance_scale}")
+
+    # Ensure height and width are multiples of 8 (common requirement for VAEs)
+    height = (int(height) // 8) * 8
+    width = (int(width) // 8) * 8
+    num_frames = int(num_frames)
+    fps = int(fps)
+
+    with torch.inference_mode(): # Conserve memory
+        output_frames_list = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            guidance_scale=float(guidance_scale),
+            # num_inference_steps=25 # Default is 25, can be exposed if needed
+        ).frames
+
+    if not output_frames_list or not output_frames_list[0]:
+        raise gr.Error("Model returned empty frames. Check parameters or try a different prompt.")
+
+    output_frames = output_frames_list[0] # The actual list of PIL Image frames
+
+    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
+        video_path = tmpfile.name
+
+    export_to_video(output_frames, video_path, fps=fps)
+    print(f"Video successfully generated and saved to {video_path}")
+
+    return video_path
+
+
+# --- Gradio UI Definition ---
+default_prompt = "A cat walks on the grass, realistic"
+default_negative_prompt = "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.Blocks() as demo:
+    gr.Markdown(f"""
+    # Text-to-Video with Wan 2.1 (14B)
+    Powered by `diffusers` and `Wan-AI/{MODEL_ID}`.
+    Model is loaded into memory when the app starts. This might take a few minutes.
+    Ensure you have a GPU with sufficient VRAM (e.g., ~24GB+ for these default settings).
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            prompt_input = gr.Textbox(label="Prompt", value=default_prompt, lines=3)
+            negative_prompt_input = gr.Textbox(
+                label="Negative Prompt (Optional)",
+                value=default_negative_prompt,
+                lines=3
+            )
+            with gr.Row():
+                height_input = gr.Slider(minimum=256, maximum=768, step=64, value=480, label="Height (multiple of 8)")
+                width_input = gr.Slider(minimum=256, maximum=1024, step=64, value=832, label="Width (multiple of 8)")
+            with gr.Row():
+                num_frames_input = gr.Slider(minimum=16, maximum=100, step=1, value=25, label="Number of Frames") # Reduced default for faster demo
+                fps_input = gr.Slider(minimum=5, maximum=30, step=1, value=15, label="Output FPS")
+            guidance_scale_input = gr.Slider(minimum=1.0, maximum=20.0, step=0.5, value=5.0, label="Guidance Scale")
+
+            generate_button = gr.Button("Generate Video", variant="primary")
+
+        with gr.Column(scale=3):
+            video_output = gr.Video(label="Generated Video")
+
+    generate_button.click(
+        fn=generate_video,
+        inputs=[
+            prompt_input,
+            negative_prompt_input,
+            height_input,
+            width_input,
+            num_frames_input,
+            guidance_scale_input,
+            fps_input
+        ],
+        outputs=video_output
+    )
+
+    gr.Examples(
+        examples=[
+            ["A panda eating bamboo in a lush forest, cinematic lighting", default_negative_prompt, 480, 832, 25, 5.0, 15],
+            ["A majestic eagle soaring over snowy mountains", default_negative_prompt, 512, 768, 30, 7.0, 12],
+            ["Timelapse of a flower blooming, vibrant colors", "static, ugly", 384, 640, 40, 6.0, 20],
+            ["Astronaut walking on the moon, Earth in the background, highly detailed", default_negative_prompt, 480, 832, 20, 5.5, 10],
+        ],
+        inputs=[prompt_input, negative_prompt_input, height_input, width_input, num_frames_input, guidance_scale_input, fps_input],
+        outputs=video_output,
+        fn=generate_video,
+        cache_examples=False # Caching videos can take a lot of space quickly
+    )
+
+if __name__ == "__main__":
+    # The share=True option will create a public temporary link if you run this on Colab or similar
+    demo.queue().launch(share=True, debug=True)