app.py

import os
# PyTorch 2.8 (temporary hack)
os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')

# --- 1. Model Download and Setup (Diffusers Backend) ---
import spaces
import torch
from diffusers import FlowMatchEulerDiscreteScheduler
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

# Import the optimization function from the separate file
from optimization import optimize_pipeline_

# --- Constants and Model Loading ---
MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"

# --- NEW: Flexible Dimension Constants ---
MAX_DIMENSION = 832
MIN_DIMENSION = 480
DIMENSION_MULTIPLE = 16
SQUARE_SIZE = 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)

default_negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走,过曝，"

print("Loading models into memory. This may take a few minutes...")

pipe = WanImageToVideoPipeline.from_pretrained(
    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,
)
pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config, shift=8.0)
pipe.to('cuda')



print("Optimizing pipeline...")
for i in range(3):
    gc.collect()
    torch.cuda.synchronize()
    torch.cuda.empty_cache()

# Calling the imported optimization function with a placeholder image for compilation tracing
optimize_pipeline_(pipe,
    image=Image.new('RGB', (MAX_DIMENSION, MIN_DIMENSION)), # Use representative dims
    prompt='prompt',
    height=MIN_DIMENSION,
    width=MAX_DIMENSION,
    num_frames=MAX_FRAMES_MODEL,
)
print("All models loaded and optimized. Gradio app is ready.")


# --- 2. Image Processing and Application Logic ---

def process_image_for_video(image: Image.Image) -> Image.Image:
    """
    Resizes an image based on the following rules for video generation:
    1. The longest side will be scaled down to MAX_DIMENSION if it's larger.
    2. The shortest side will be scaled up to MIN_DIMENSION if it's smaller.
    3. The final dimensions will be rounded to the nearest multiple of DIMENSION_MULTIPLE.
    4. Square images are resized to a fixed SQUARE_SIZE.
    The aspect ratio is preserved as closely as possible.
    """
    width, height = image.size

    # Rule 4: Handle square images
    if width == height:
        return image.resize((SQUARE_SIZE, SQUARE_SIZE), Image.Resampling.LANCZOS)

    # Determine target dimensions while preserving aspect ratio
    aspect_ratio = width / height
    new_width, new_height = width, height

    # Rule 1: Scale down if too large
    if new_width > MAX_DIMENSION or new_height > MAX_DIMENSION:
        if aspect_ratio > 1:  # Landscape
            scale = MAX_DIMENSION / new_width
        else:  # Portrait
            scale = MAX_DIMENSION / new_height
        new_width *= scale
        new_height *= scale

    # Rule 2: Scale up if too small
    if new_width < MIN_DIMENSION or new_height < MIN_DIMENSION:
        if aspect_ratio > 1:  # Landscape
            scale = MIN_DIMENSION / new_height
        else:  # Portrait
            scale = MIN_DIMENSION / new_width
        new_width *= scale
        new_height *= scale

    # Rule 3: Round to the nearest multiple of DIMENSION_MULTIPLE
    final_width = int(round(new_width / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
    final_height = int(round(new_height / DIMENSION_MULTIPLE) * DIMENSION_MULTIPLE)
    
    # Ensure final dimensions are at least the minimum
    final_width = max(final_width, MIN_DIMENSION if aspect_ratio < 1 else SQUARE_SIZE)
    final_height = max(final_height, MIN_DIMENSION if aspect_ratio > 1 else SQUARE_SIZE)


    return image.resize((final_width, final_height), Image.Resampling.LANCZOS)

def resize_and_crop_to_match(target_image, reference_image):
    """Resizes and center-crops the target image to match the reference image's dimensions."""
    ref_width, ref_height = reference_image.size
    target_width, target_height = target_image.size
    scale = max(ref_width / target_width, ref_height / target_height)
    new_width, new_height = int(target_width * scale), int(target_height * scale)
    resized = target_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    left, top = (new_width - ref_width) // 2, (new_height - ref_height) // 2
    return resized.crop((left, top, left + ref_width, top + ref_height))

@spaces.GPU(duration=120)
def generate_video(
    start_image_pil,
    end_image_pil,
    prompt,
    negative_prompt=default_negative_prompt,
    duration_seconds=2.1,
    steps=8,
    guidance_scale=1,
    guidance_scale_2=1,
    seed=42,
    randomize_seed=False,
    progress=gr.Progress(track_tqdm=True)
):
    """
    Generates a video by interpolating between a start and end image, guided by a text prompt,
    using the diffusers Wan2.2 pipeline.
    """
    if start_image_pil is None or end_image_pil is None:
        raise gr.Error("Please upload both a start and an end image.")

    progress(0.1, desc="Preprocessing images...")

    # Step 1: Process the start image to get our target dimensions based on the new rules.
    processed_start_image = process_image_for_video(start_image_pil)
    
    # Step 2: Make the end image match the *exact* dimensions of the processed start image.
    processed_end_image = resize_and_crop_to_match(end_image_pil, processed_start_image)
    
    target_height, target_width = processed_start_image.height, processed_start_image.width

    # Handle seed and frame count
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)

    progress(0.2, desc=f"Generating {num_frames} frames at {target_width}x{target_height} (seed: {current_seed})...")

    output_frames_list = pipe(
        image=processed_start_image,
        last_image=processed_end_image,
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=target_height,
        width=target_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]

    progress(0.9, desc="Encoding and saving video...")

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name

    export_to_video(output_frames_list, video_path, fps=FIXED_FPS)

    progress(1.0, desc="Done!")
    return video_path, current_seed


# --- 3. Gradio User Interface --- (No changes needed here)

css = '''
.fillable{max-width: 1100px !important}
.dark .progress-text {color: white}
'''
with gr.Blocks(theme=gr.themes.Citrus(), css=css) as app:
    gr.Markdown("# Wan 2.2 First/Last Frame Video Fast")
    gr.Markdown("Based on the [Wan 2.2 First/Last Frame workflow](https://www.reddit.com/r/StableDiffusion/comments/1me4306/psa_wan_22_does_first_frame_last_frame_out_of_the/), applied to 🧨 Diffusers + [lightx2v/Wan2.2-Lightning](https://huggingface.co/lightx2v/Wan2.2-Lightning) 8-step LoRA")

    with gr.Row():
        with gr.Column():
            with gr.Group():
                with gr.Row():
                    start_image = gr.Image(type="pil", label="Start Frame", sources=["upload", "clipboard"])
                    end_image = gr.Image(type="pil", label="End Frame", sources=["upload", "clipboard"])

                prompt = gr.Textbox(label="Prompt", info="Describe the transition between the two images")

                with gr.Accordion("Advanced Settings", open=False):
                    duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=2.1, label="Video Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
                    negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
                    steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=8, label="Inference Steps")
                    guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale - high noise")
                    guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Scale - low noise")
                    with gr.Row():
                        seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
                        randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True)

                generate_button = gr.Button("Generate Video", variant="primary")

        with gr.Column():
            output_video = gr.Video(label="Generated Video", autoplay=True)

    # Define the inputs list for the click event
    ui_inputs = [
        start_image,
        end_image,
        prompt,
        negative_prompt_input,
        duration_seconds_input,
        steps_slider,
        guidance_scale_input,
        guidance_scale_2_input,
        seed_input,
        randomize_seed_checkbox
    ]
    # The seed_input is both an input and an output to reflect the randomly generated seed
    ui_outputs = [output_video, seed_input]

    generate_button.click(
        fn=generate_video,
        inputs=ui_inputs,
        outputs=ui_outputs
    )

    gr.Examples(
        examples=[
            ["poli_tower.png", "tower_takes_off.png", "the man turns around"],
            ["ugly_sonic.jpeg", "squatting_sonic.png", "the character dodges the missiles"],
            ["capyabara_zoomed.png", "capyabara.webp", "a dramatic dolly zoom"],
        ],
        inputs=[start_image, end_image, prompt],
        outputs=ui_outputs,
        fn=generate_video,
        cache_examples="lazy",
    )

if __name__ == "__main__":
    app.launch(share=True)