app1.py

import os
# Set environment variables before any imports to suppress inductor warnings
os.environ["TORCHINDUCTOR_CUDA_GRAPHS"] = "0"
os.environ["TORCHINDUCTOR_MAX_AUTOTUNE_GEMM"] = "0"

# Install dependencies as specified
os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')

import spaces
import torch
from diffusers import WanPipeline, AutoencoderKLWan
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
from diffusers.utils.export_utils import export_to_video
import gradio as gr
import tempfile
import numpy as np
from PIL import Image
import random
import gc

# Debug: Log PyTorch version and environment variables
print(f"PyTorch version: {torch.__version__}")
print(f"TORCHINDUCTOR_MAX_AUTOTUNE_GEMM: {os.environ.get('TORCHINDUCTOR_MAX_AUTOTUNE_GEMM')}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

# Assuming optimize_pipeline_ is a custom function; if not available, define a no-op
try:
    from optimization import optimize_pipeline_
except ImportError:
    def optimize_pipeline_(pipe, **kwargs):
        pass  # No-op if optimization is not available

# Model configurations
T2V_MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
I2V_MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
LANDSCAPE_WIDTH = 832
LANDSCAPE_HEIGHT = 480
MAX_SEED = np.iinfo(np.int32).max
FIXED_FPS = 16
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 81
MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS, 1)
MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS, 1)

# Cache for pipelines
t2v_pipe_cache = [None]
i2v_pipe_cache = [None]

def clear_memory():
    """Aggressively clear memory and CUDA cache."""
    for _ in range(3):
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()

def load_t2v_pipeline():
    """Load and optimize the T2V pipeline."""
    if t2v_pipe_cache[0] is None:
        try:
            print("Loading T2V pipeline...")
            vae = AutoencoderKLWan.from_pretrained(T2V_MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
            t2v_pipe_cache[0] = WanPipeline.from_pretrained(T2V_MODEL_ID,
                transformer=WanTransformer3DModel.from_pretrained('linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
                    subfolder='transformer',
                    torch_dtype=torch.bfloat16,
                    device_map='cuda',
                ),
                transformer_2=WanTransformer3DModel.from_pretrained('linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
                    subfolder='transformer_2',
                    torch_dtype=torch.bfloat16,
                    device_map='cuda',
                ),
                vae=vae,
                torch_dtype=torch.bfloat16,
            ).to('cuda')
            optimize_pipeline_(t2v_pipe_cache[0],
                prompt='prompt',
                height=LANDSCAPE_HEIGHT,
                width=LANDSCAPE_WIDTH,
                num_frames=MAX_FRAMES_MODEL,
            )
            t2v_pipe_cache[0].enable_model_cpu_offload()  # Enable CPU offload for memory optimization
            print("T2V pipeline loaded successfully")
        except Exception as e:
            print(f"Error loading T2V pipeline: {e}")
            t2v_pipe_cache[0] = None
            raise
        clear_memory()
    return t2v_pipe_cache[0]

def load_i2v_pipeline():
    """Load and optimize the I2V pipeline."""
    if i2v_pipe_cache[0] is None:
        try:
            print("Loading I2V pipeline...")
            i2v_pipe_cache[0] = WanImageToVideoPipeline.from_pretrained(I2V_MODEL_ID,
                transformer=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
                    subfolder='transformer',
                    torch_dtype=torch.bfloat16,
                    device_map='cuda',
                ),
                transformer_2=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
                    subfolder='transformer_2',
                    torch_dtype=torch.bfloat16,
                    device_map='cuda',
                ),
                torch_dtype=torch.bfloat16,
            ).to('cuda')
            optimize_pipeline_(i2v_pipe_cache[0],
                image=Image.new('RGB', (LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT)),
                prompt='prompt',
                height=LANDSCAPE_HEIGHT,
                width=LANDSCAPE_WIDTH,
                num_frames=MAX_FRAMES_MODEL,
            )
            i2v_pipe_cache[0].enable_model_cpu_offload()  # Enable CPU offload for memory optimization
            print("I2V pipeline loaded successfully")
        except Exception as e:
            print(f"Error loading I2V pipeline: {e}")
            i2v_pipe_cache[0] = None
            raise
        clear_memory()
    return i2v_pipe_cache[0]

def unload_t2v_pipeline():
    """Unload the T2V pipeline to free memory."""
    if t2v_pipe_cache[0] is not None:
        t2v_pipe_cache[0].to("cpu")
        del t2v_pipe_cache[0]
        t2v_pipe_cache[0] = None
        clear_memory()

def unload_i2v_pipeline():
    """Unload the I2V pipeline to free memory."""
    if i2v_pipe_cache[0] is not None:
        i2v_pipe_cache[0].to("cpu")
        del i2v_pipe_cache[0]
        i2v_pipe_cache[0] = None
        clear_memory()

# Load T2V pipeline at startup
try:
    load_t2v_pipeline()
except Exception as e:
    print(f"Failed to load T2V pipeline at startup: {e}")

# Default prompts
default_prompt_t2v = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"

def resize_image(image: Image.Image) -> Image.Image:
    if image.height > image.width:
        transposed = image.transpose(Image.Transpose.ROTATE_90)
        resized = resize_image_landscape(transposed)
        return resized.transpose(Image.Transpose.ROTATE_270)
    return resize_image_landscape(image)

def resize_image_landscape(image: Image.Image) -> Image.Image:
    target_aspect = LANDSCAPE_WIDTH / LANDSCAPE_HEIGHT
    width, height = image.size
    in_aspect = width / height
    if in_aspect > target_aspect:
        new_width = round(height * target_aspect)
        left = (width - new_width) // 2
        image = image.crop((left, 0, left + new_width, height))
    else:
        new_height = round(width / target_aspect)
        top = (height - new_height) // 2
        image = image.crop((0, top, width, top + new_height))
    return image.resize((LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT), Image.LANCZOS)

def get_duration(
    mode,
    input_image,
    prompt,
    negative_prompt,
    duration_seconds,
    guidance_scale,
    guidance_scale_2,
    steps,
    seed,
    randomize_seed,
    progress,
):
    return int(steps) * 15

@spaces.GPU(duration=get_duration)
@torch.no_grad()
def generate_video(
    mode,
    input_image,
    prompt,
    negative_prompt=default_negative_prompt,
    duration_seconds=MAX_DURATION,
    guidance_scale=1,
    guidance_scale_2=1,
    steps=4,
    seed=42,
    randomize_seed=False,
    progress=gr.Progress(track_tqdm=True),
):
    # Debug: Log inputs to verify steps
    print(f"Generating video with mode={mode}, steps={steps}, num_frames={int(round(duration_seconds * FIXED_FPS))}")

    if mode == "Image-to-Video" and input_image is None:
        raise gr.Error("Please upload an input image for Image-to-Video mode.")

    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)

    if mode == "Text-to-Video":
        unload_i2v_pipeline()  # Unload I2V to free memory
        pipe = load_t2v_pipeline()
        if pipe is None:
            raise gr.Error("T2V pipeline failed to load. Please check logs.")
        output_frames_list = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=LANDSCAPE_HEIGHT,
            width=LANDSCAPE_WIDTH,
            num_frames=num_frames,
            guidance_scale=float(guidance_scale),
            guidance_scale_2=float(guidance_scale_2),
            num_inference_steps=int(steps),
            generator=torch.Generator(device="cuda").manual_seed(current_seed),
        ).frames[0]
    else:  # Image-to-Video
        unload_t2v_pipeline()  # Unload T2V to free memory
        pipe = load_i2v_pipeline()
        if pipe is None:
            raise gr.Error("I2V pipeline failed to load. Please check logs.")
        resized_image = resize_image(input_image)
        output_frames_list = pipe(
            image=resized_image,
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=resized_image.height,
            width=resized_image.width,
            num_frames=num_frames,
            guidance_scale=float(guidance_scale),
            guidance_scale_2=float(guidance_scale_2),
            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_FPS)
    clear_memory()  # Clean up after generation
    return video_path, current_seed

with gr.Blocks() as demo:
    gr.Markdown("# Fast 4 steps Wan 2.2 T2V/I2V (14B) with Lightning LoRA")
    gr.Markdown("Run Wan 2.2 in just 4-8 steps, with [Lightning LoRA](https://huggingface.co/Kijai/WanVideo_comfy/tree/main/Wan22-Lightning), fp8 quantization & AoT compilation - compatible with 🧨 diffusers and ZeroGPU⚡️")
    
    with gr.Tabs() as tabs:
        with gr.TabItem("Text-to-Video"):
            with gr.Row():
                with gr.Column():
                    t2v_prompt_input = gr.Textbox(label="Prompt", value=default_prompt_t2v)
                    t2v_duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=MAX_DURATION, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
                    with gr.Accordion("Advanced Settings", open=False):
                        t2v_negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
                        t2v_seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
                        t2v_randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
                        t2v_steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=4, label="Inference Steps")
                        t2v_guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale - high noise stage")
                        t2v_guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=3, label="Guidance Scale 2 - low noise stage")
                    t2v_generate_button = gr.Button("Generate Video", variant="primary")
                with gr.Column():
                    t2v_video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
            
            t2v_inputs = [
                gr.State(value="Text-to-Video"),
                gr.State(value=None),
                t2v_prompt_input,
                t2v_negative_prompt_input,
                t2v_duration_seconds_input,
                t2v_guidance_scale_input,
                t2v_guidance_scale_2_input,
                t2v_steps_slider,
                t2v_seed_input,
                t2v_randomize_seed_checkbox
            ]
            t2v_generate_button.click(fn=generate_video, inputs=t2v_inputs, outputs=[t2v_video_output, t2v_seed_input])

            gr.Examples(
                examples=[
                    ["POV selfie video, white cat with sunglasses standing on surfboard, relaxed smile, tropical beach behind (clear water, green hills, blue sky with clouds). Surfboard tips, cat falls into ocean, camera plunges underwater with bubbles and sunlight beams. Brief underwater view of cat’s face, then cat resurfaces, still filming selfie, playful summer vacation mood."],
                    ["Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."],
                    ["A cinematic shot of a boat sailing on a calm sea at sunset."],
                    ["Drone footage flying over a futuristic city with flying cars."],
                ],
                inputs=[t2v_prompt_input],
                outputs=[t2v_video_output, t2v_seed_input],
                fn=generate_video,
                cache_examples="lazy"
            )

        with gr.TabItem("Image-to-Video"):
            with gr.Row():
                with gr.Column():
                    i2v_input_image_component = gr.Image(type="pil", label="Input Image (auto-resized to target H/W)")
                    i2v_prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v)
                    i2v_duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=MAX_DURATION, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
                    with gr.Accordion("Advanced Settings", open=False):
                        i2v_negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
                        i2v_seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
                        i2v_randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
                        i2v_steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=6, label="Inference Steps")
                        i2v_guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale - high noise stage")
                        i2v_guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale 2 - low noise stage")
                    i2v_generate_button = gr.Button("Generate Video", variant="primary")
                with gr.Column():
                    i2v_video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
            
            i2v_inputs = [
                gr.State(value="Image-to-Video"),
                i2v_input_image_component,
                i2v_prompt_input,
                i2v_negative_prompt_input,
                i2v_duration_seconds_input,
                i2v_guidance_scale_input,
                i2v_guidance_scale_2_input,
                i2v_steps_slider,
                i2v_seed_input,
                i2v_randomize_seed_checkbox
            ]
            i2v_generate_button.click(fn=generate_video, inputs=i2v_inputs, outputs=[i2v_video_output, i2v_seed_input])

            gr.Examples(
                examples=[
                    ["wan_i2v_input.JPG", "POV selfie video, white cat with sunglasses standing on surfboard, relaxed smile, tropical beach behind (clear water, green hills, blue sky with clouds). Surfboard tips, cat falls into ocean, camera plunges underwater with bubbles and sunlight beams. Brief underwater view of cat’s face, then cat resurfaces, still filming selfie, playful summer vacation mood.", 4],
                    ["wan22_input_2.jpg", "A sleek lunar vehicle glides into view from left to right, kicking up moon dust as astronauts in white spacesuits hop aboard with characteristic lunar bouncing movements. In the distant background, a VTOL craft descends straight down and lands silently on the surface. Throughout the entire scene, ethereal aurora borealis ribbons dance across the star-filled sky, casting shimmering curtains of green, blue, and purple light that bathe the lunar landscape in an otherworldly, magical glow.", 4],
                    ["kill_bill.jpeg", "Uma Thurman's character, Beatrix Kiddo, holds her razor-sharp katana blade steady in the cinematic lighting. Suddenly, the polished steel begins to soften and distort, like heated metal starting to lose its structural integrity. The blade's perfect edge slowly warps and droops, molten steel beginning to flow downward in silvery rivulets while maintaining its metallic sheen. The transformation starts subtly at first - a slight bend in the blade - then accelerates as the metal becomes increasingly fluid. The camera holds steady on her face as her piercing eyes gradually narrow, not with lethal focus, but with confusion and growing alarm as she watches her weapon dissolve before her eyes. Her breathing quickens slightly as she witnesses this impossible transformation. The melting intensifies, the katana's perfect form becoming increasingly abstract, dripping like liquid mercury from her grip. Molten droplets fall to the ground with soft metallic impacts. Her expression shifts from calm readiness to bewilderment and concern as her legendary instrument of vengeance literally liquefies in her hands, leaving her defenseless and disoriented.", 6],
                ],
                inputs=[i2v_input_image_component, i2v_prompt_input, i2v_steps_slider],
                outputs=[i2v_video_output, i2v_seed_input],
                fn=generate_video,
                cache_examples="lazy"
            )

if __name__ == "__main__":
    demo.queue().launch(mcp_server=True)