app.py

import torch
from diffusers import AutoencoderKLWan, WanPipeline, UniPCMultistepScheduler
from diffusers.utils import export_to_video
import gradio as gr
import tempfile
import spaces
from huggingface_hub import hf_hub_download
import numpy as np
import random
import os

MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
LORA_REPO_ID = "Kijai/WanVideo_comfy"
LORA_FILENAME = "Wan21_CausVid_bidirect2_T2V_1_3B_lora_rank32.safetensors"

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=8.0)
pipe.to("cuda")

causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)
pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
pipe.set_adapters(["causvid_lora"], adapter_weights=[0.95])
pipe.fuse_lora()

MOD_VALUE = 32
DEFAULT_H_SLIDER_VALUE = 384 # 512
DEFAULT_W_SLIDER_VALUE = 640 # 896

# Environment variable check
IS_ORIGINAL_SPACE = os.environ.get("IS_ORIGINAL_SPACE", "False") == "True"

# Original limits
ORIGINAL_SLIDER_MIN_H, ORIGINAL_SLIDER_MAX_H = 128, 1280
ORIGINAL_SLIDER_MIN_W, ORIGINAL_SLIDER_MAX_W = 128, 1280
ORIGINAL_MAX_DURATION = round(81/24, 1)  # MAX_FRAMES_MODEL/FIXED_FPS

# Limited space constants
LIMITED_MAX_RESOLUTION = 640
LIMITED_MAX_DURATION = 2.0
LIMITED_MAX_STEPS = 4

# Set limits based on environment variable
if IS_ORIGINAL_SPACE:
    SLIDER_MIN_H, SLIDER_MAX_H = 128, LIMITED_MAX_RESOLUTION
    SLIDER_MIN_W, SLIDER_MAX_W = 128, LIMITED_MAX_RESOLUTION
    MAX_DURATION = LIMITED_MAX_DURATION
    MAX_STEPS = LIMITED_MAX_STEPS
else:
    SLIDER_MIN_H, SLIDER_MAX_H = ORIGINAL_SLIDER_MIN_H, ORIGINAL_SLIDER_MAX_H
    SLIDER_MIN_W, SLIDER_MAX_W = ORIGINAL_SLIDER_MIN_W, ORIGINAL_SLIDER_MAX_W
    MAX_DURATION = ORIGINAL_MAX_DURATION
    MAX_STEPS = 8

MAX_SEED = np.iinfo(np.int32).max

FIXED_FPS = 24
FIXED_OUTPUT_FPS = 18 # we downspeed the output video as a temporary "trick"
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 81 

default_prompt_t2v = "A wide aerial view of a bustling open-air street market in a densely built urban neighborhood, narrow lanes packed with colorful umbrellas and stalls, crowds of people walking, talking, and shopping, vibrant fabrics and goods displayed on tables, vendors gesturing to customers, children weaving through the crowd, occasional bicycles passing, the scene alive with constant movement and chatter, distant city skyline visible under a hazy sky, smooth cinematic camera pan across the lively street"
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, watermark, text, signature"

def get_duration(prompt, height, width, 
                   negative_prompt, duration_seconds,
                   guidance_scale, steps,
                   seed, randomize_seed, 
                   progress):
    if steps > 4 and duration_seconds > 2:
        return 90
    elif steps > 4 or duration_seconds > 2:
        return 75
    else:
        return 60

@spaces.GPU(duration=get_duration)
def generate_video(prompt, height, width, 
                   negative_prompt=default_negative_prompt, duration_seconds = 2,
                   guidance_scale = 1, steps = 4,
                   seed = 42, randomize_seed = False, 
                   progress=gr.Progress(track_tqdm=True)):
    """
    Generate a video from a text prompt using the Wan 2.1 T2V model with CausVid LoRA.
    
    This function takes a text prompt and generates a video based on the provided
    prompt and parameters. It uses the Wan 2.1 1.3B Text-to-Video model with CausVid LoRA
    for fast generation in 3-8 steps.
    
    Args:
        prompt (str): Text prompt describing the desired video content.
        height (int): Target height for the output video. Will be adjusted to multiple of MOD_VALUE (32).
        width (int): Target width for the output video. Will be adjusted to multiple of MOD_VALUE (32).
        negative_prompt (str, optional): Negative prompt to avoid unwanted elements. 
            Defaults to default_negative_prompt (contains unwanted visual artifacts).
        duration_seconds (float, optional): Duration of the generated video in seconds.
            Defaults to 2. Clamped between MIN_FRAMES_MODEL/FIXED_FPS and MAX_FRAMES_MODEL/FIXED_FPS.
        guidance_scale (float, optional): Controls adherence to the prompt. Higher values = more adherence.
            Defaults to 1.0. Range: 0.0-20.0.
        steps (int, optional): Number of inference steps. More steps = higher quality but slower.
            Defaults to 4. Range: 1-30.
        seed (int, optional): Random seed for reproducible results. Defaults to 42.
            Range: 0 to MAX_SEED (2147483647).
        randomize_seed (bool, optional): Whether to use a random seed instead of the provided seed.
            Defaults to False.
        progress (gr.Progress, optional): Gradio progress tracker. Defaults to gr.Progress(track_tqdm=True).
    
    Returns:
        tuple: A tuple containing:
            - video_path (str): Path to the generated video file (.mp4)
            - current_seed (int): The seed used for generation (useful when randomize_seed=True)
    
    Raises:
        gr.Error: If prompt is empty or None.
    
    Note:
        - Frame count is calculated as duration_seconds * FIXED_FPS (24)
        - Output dimensions are adjusted to be multiples of MOD_VALUE (32)
        - The function uses GPU acceleration via the @spaces.GPU decorator
        - Generation time varies based on steps and duration (see get_duration function)
    """
    if not prompt or prompt.strip() == "":
        raise gr.Error("Please enter a text prompt. Try to use long and precise descriptions.")

    # Apply limits based on environment variable
    if IS_ORIGINAL_SPACE:
        height = min(height, LIMITED_MAX_RESOLUTION)
        width = min(width, LIMITED_MAX_RESOLUTION)
        duration_seconds = min(duration_seconds, LIMITED_MAX_DURATION)
        steps = min(steps, LIMITED_MAX_STEPS)

    target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
    target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
    
    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
    
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    with torch.inference_mode():
        output_frames_list = pipe(
            prompt=prompt, negative_prompt=negative_prompt,
            height=target_h, width=target_w, num_frames=num_frames,
            guidance_scale=float(guidance_scale), num_inference_steps=int(steps),
            generator=torch.Generator(device="cuda").manual_seed(current_seed)
        ).frames[0]

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name
    export_to_video(output_frames_list, video_path, fps=FIXED_OUTPUT_FPS)
    return video_path, current_seed

with gr.Blocks() as demo:
    gr.Markdown("# ⚡ DEval VideoGen")
    gr.Markdown("This Gradio space is a fork of [wan2-1-fast from jbilcke-hf](https://huggingface.co/spaces/jbilcke-hf/InstaVideo), and is powered by the Wan CausVid LoRA")

    # Add notice for limited spaces
    if IS_ORIGINAL_SPACE:
        gr.Markdown("⚠️ **Limit the resolution to 640px, duration to 2s, and inference steps to 4. For full capabilities please duplicate this space.**")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_t2v, placeholder="Describe the video you want to generate...")
            
            with gr.Accordion("Advanced Settings", open=False):
                negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
                seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
                randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
                with gr.Row():
                    height_input = gr.Slider(
                        minimum=SLIDER_MIN_H, 
                        maximum=SLIDER_MAX_H, 
                        step=MOD_VALUE, 
                        value=min(DEFAULT_H_SLIDER_VALUE, SLIDER_MAX_H), 
                        label=f"Output Height (multiple of {MOD_VALUE})"
                    )
                    width_input = gr.Slider(
                        minimum=SLIDER_MIN_W, 
                        maximum=SLIDER_MAX_W, 
                        step=MOD_VALUE, 
                        value=min(DEFAULT_W_SLIDER_VALUE, SLIDER_MAX_W), 
                        label=f"Output Width (multiple of {MOD_VALUE})"
                    )
                duration_seconds_input = gr.Slider(
                    minimum=round(MIN_FRAMES_MODEL/FIXED_FPS,1), 
                    maximum=MAX_DURATION, 
                    step=0.1, 
                    value=2, 
                    label="Duration (seconds)", 
                    info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps."
                )
                steps_slider = gr.Slider(minimum=1, maximum=MAX_STEPS, step=1, value=4, label="Inference Steps") 
                guidance_scale_input = gr.Slider(minimum=0.0, maximum=20.0, step=0.5, value=1.0, label="Guidance Scale", visible=False)

            generate_button = gr.Button("Generate Video", variant="primary")
        with gr.Column():
            video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
    
    ui_inputs = [
        prompt_input, height_input, width_input,
        negative_prompt_input, duration_seconds_input,
        guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox
    ]
    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])

    # Adjust examples based on space limits
    example_configs = [
        ["a majestic condor soaring through mountain peaks, cinematic aerial view", 896, 512],
        ["A wide aerial view of a bustling open-air street market in a densely built urban neighborhood, narrow lanes packed with colorful umbrellas and stalls, crowds of people walking, talking, and shopping, vibrant fabrics and goods displayed on tables, vendors gesturing to customers, children weaving through the crowd, occasional bicycles passing, the scene alive with constant movement and chatter, distant city skyline visible under a hazy sky, smooth cinematic camera pan across the lively street", 448, 832],
    ]
    
    if IS_ORIGINAL_SPACE:
        # Limit example resolutions for limited spaces
        example_configs = [
            [example[0], min(example[1], LIMITED_MAX_RESOLUTION), min(example[2], LIMITED_MAX_RESOLUTION)]
            for example in example_configs
        ]
    
    gr.Examples(
        examples=example_configs,
        inputs=[prompt_input, height_input, width_input], 
        outputs=[video_output, seed_input], 
        fn=generate_video, 
        cache_examples="lazy"
    )

if __name__ == "__main__":
    auth_user = os.getenv("USER")
    auth_pass = os.getenv("PASS")

    demo.queue().launch(
        auth=(auth_user, auth_pass), share=True,  ssr_mode=False
    )