app.py

import torch
from typing import List
from diffusers import AutoencoderKLWan, WanVACEPipeline, 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
from PIL import Image
import random

model_id = "Wan-AI/Wan2.1-VACE-14B-diffusers"
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
pipe = WanVACEPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16).to("cuda")

pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=2.0)

pipe.load_lora_weights(
   "vrgamedevgirl84/Wan14BT2VFusioniX", 
   weight_name="FusionX_LoRa/Phantom_Wan_14B_FusionX_LoRA.safetensors", 
    adapter_name="phantom"
)
pipe.load_lora_weights(
   "vrgamedevgirl84/Wan14BT2VFusioniX", 
   weight_name="OtherLoRa's/DetailEnhancerV1.safetensors", adapter_name="detailer"
)
pipe.set_adapters(["phantom","detailer"], adapter_weights=[1, .9])
pipe.fuse_lora()

MOD_VALUE = 32
DEFAULT_H_SLIDER_VALUE = 512
DEFAULT_W_SLIDER_VALUE = 896
NEW_FORMULA_MAX_AREA = 480.0 * 832.0 

SLIDER_MIN_H, SLIDER_MAX_H = 128, 896
SLIDER_MIN_W, SLIDER_MAX_W = 128, 896
MAX_SEED = np.iinfo(np.int32).max

FIXED_FPS = 16
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 81 

# Default prompts for different modes
MODE_PROMPTS = {
    "Ref2V": "the playful penguin picks up the green cat eye sunglasses and puts them on",
    "FLF2V": "CG animation style, a small blue bird takes off from the ground, flapping its wings. The bird's feathers are delicate, with a unique pattern on its chest. The background shows a blue sky with white clouds under bright sunshine. The camera follows the bird upward, capturing its flight and the vastness of the sky from a close-up, low-angle perspective.",
    "Random2V": "Various different characters appear and disappear in a fast transition video showcasting their unique features and personalities. The video is about showcasing different dance styles, with each character performing a distinct dance move. The background is a vibrant, colorful stage with dynamic lighting that changes with each dance style. The camera captures close-ups of the dancers' expressions and movements. Highly dynamic, fast-paced music video, with quick cuts and transitions between characters, cinematic, vibrant colors"
}

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 _calculate_new_dimensions_wan(pil_image, mod_val, calculation_max_area,
                                 min_slider_h, max_slider_h,
                                 min_slider_w, max_slider_w,
                                 default_h, default_w):
    orig_w, orig_h = pil_image.size
    if orig_w <= 0 or orig_h <= 0:
        return default_h, default_w

    aspect_ratio = orig_h / orig_w
    
    calc_h = round(np.sqrt(calculation_max_area * aspect_ratio))
    calc_w = round(np.sqrt(calculation_max_area / aspect_ratio))

    calc_h = max(mod_val, (calc_h // mod_val) * mod_val)
    calc_w = max(mod_val, (calc_w // mod_val) * mod_val)
    
    new_h = int(np.clip(calc_h, min_slider_h, (max_slider_h // mod_val) * mod_val))
    new_w = int(np.clip(calc_w, min_slider_w, (max_slider_w // mod_val) * mod_val))
    
    return new_h, new_w

def handle_gallery_upload_for_dims_wan(gallery_images, current_h_val, current_w_val):
    if gallery_images is None or len(gallery_images) == 0:
        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
    try:
        # Use the first image to calculate dimensions
        first_image = gallery_images[0][0]
        new_h, new_w = _calculate_new_dimensions_wan(
            first_image, MOD_VALUE, NEW_FORMULA_MAX_AREA,
            SLIDER_MIN_H, SLIDER_MAX_H, SLIDER_MIN_W, SLIDER_MAX_W,
            DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE
        )
        return gr.update(value=new_h), gr.update(value=new_w)
    except Exception as e:
        gr.Warning("Error attempting to calculate new dimensions")
        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)

def update_prompt_from_mode(mode):
    """Update the prompt based on the selected mode"""
    return MODE_PROMPTS.get(mode, "")


def prepare_video_and_mask_Ref2V(height: int, width: int, num_frames: int):
    frames = []
    # Ideally, this should be 127.5 to match original code, but they perform computation on numpy arrays
    # whereas we are passing PIL images. If you choose to pass numpy arrays, you can set it to 127.5 to
    # match the original code.
    frames.extend([Image.new("RGB", (width, height), (128, 128, 128))] * (num_frames))
    mask_white = Image.new("L", (width, height), 255)
    mask = [mask_white] * (num_frames)
    return frames, mask

def prepare_video_and_mask_FLF2V(first_img: Image.Image, last_img: Image.Image, height: int, width: int, num_frames: int):
    first_img = first_img.resize((width, height))
    last_img = last_img.resize((width, height))
    frames = []
    frames.append(first_img)
    # Ideally, this should be 127.5 to match original code, but they perform computation on numpy arrays
    # whereas we are passing PIL images. If you choose to pass numpy arrays, you can set it to 127.5 to
    # match the original code.
    frames.extend([Image.new("RGB", (width, height), (128, 128, 128))] * (num_frames - 2))
    frames.append(last_img)
    mask_black = Image.new("L", (width, height), 0)
    mask_white = Image.new("L", (width, height), 255)
    mask = [mask_black, *[mask_white] * (num_frames - 2), mask_black]
    return frames, mask

def prepare_video_and_mask_Random2V(images: List[Image.Image], frame_indices: List[int], height: int, width: int, num_frames: int):
    images = [img.resize((width, height)) for img in images]
    # Ideally, this should be 127.5 to match original code, but they perform computation on numpy arrays
    # whereas we are passing PIL images. If you choose to pass numpy arrays, you can set it to 127.5 to
    # match the original code.
    frames = [Image.new("RGB", (width, height), (128, 128, 128))] * num_frames
    
    mask_black = Image.new("L", (width, height), 0)
    mask_white = Image.new("L", (width, height), 255)
    mask = [mask_white] * num_frames
    
    for img, idx in zip(images, frame_indices):
        assert idx < num_frames, f"Frame index {idx} exceeds num_frames {num_frames}"
        frames[idx] = img
        mask[idx] = mask_black
    
    return frames, mask

def get_duration(gallery_images, mode, 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(gallery_images, mode, 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 gallery images using the selected mode.
    
    Args:
        gallery_images (list): List of PIL images from the gallery
        mode (str): Processing mode - "Ref2V", "FLF2V", or "Random2V"
        prompt (str): Text prompt describing the desired animation
        height (int): Target height for the output video
        width (int): Target width for the output video
        negative_prompt (str): Negative prompt to avoid unwanted elements
        duration_seconds (float): Duration of the generated video in seconds
        guidance_scale (float): Controls adherence to the prompt
        steps (int): Number of inference steps
        seed (int): Random seed for reproducible results
        randomize_seed (bool): Whether to use a random seed
        progress (gr.Progress): Gradio progress tracker
    
    Returns:
        tuple: (video_path, current_seed)
    """
    if gallery_images is None or len(gallery_images) == 0:
        raise gr.Error("Please upload at least one image to the gallery.")
    else:
        gallery_images = [img[0] for img in gallery_images]

    if mode == "FLF2V" and len(gallery_images) >= 2:
        gallery_images = gallery_images[:2]
    elif mode == "FLF2V" and len(gallery_images) < 2:
        raise gr.Error("FLF2V mode requires at least 2 images, but only {} were supplied.".format(len(gallery_images)))

    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)

    # Process images based on the selected mode
    if mode == "FLF2V": 
        frames, mask = prepare_video_and_mask_FLF2V(
            first_img=gallery_images[0], 
            last_img=gallery_images[1], 
            height=target_h, 
            width=target_w, 
            num_frames=num_frames
        )
        reference_images = None
    elif mode == "Ref2V": 
        frames, mask = prepare_video_and_mask_Ref2V(height=target_h, width=target_w, num_frames=num_frames)
        reference_images = gallery_images
    else:  # mode == "Random2V"
        
        frames, mask = prepare_video_and_mask_Random2V(
            images=gallery_images, 
            frame_indices=[0,20,40], # todo - generalize
            height=target_h, 
            width=target_w, 
            num_frames=num_frames
        )
        reference_images = None

    with torch.inference_mode():
        output_frames_list = pipe(
            video=frames,
            mask=mask,
            reference_images=reference_images,
            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_FPS)
    return video_path, current_seed

with gr.Blocks() as demo:
    gr.Markdown("# Wan 2.1 VACE (14B) with Phantom & Detail Enhancer LoRAs - Multi-Image Gallery")
    gr.Markdown("Using [Wan2.1-VACE-14B](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) with Phantom FusionX and Detail Enhancer LoRAs for advanced video generation with multiple conditioning modes.")
    
    with gr.Row():
        with gr.Column():
            # Gallery component for multiple image upload
            gallery_component = gr.Gallery(
                label="Upload Images",
                show_label=True,
                elem_id="gallery",
                columns=3,
                rows=2,
                object_fit="contain",
                height="auto",
                type="pil",
                allow_preview=True
            )
            
            # Radio button for mode selection
            mode_radio = gr.Radio(
                choices=["Ref2V", "FLF2V", "Random2V"],
                value="Ref2V",
                label="Processing Mode",
                info="Ref2V: Reference to Video | FLF2V: First-Last Frame to Video | Random2V: Random frames to Video"
            )
            
            prompt_input = gr.Textbox(label="Prompt", value=MODE_PROMPTS["Ref2V"])
            duration_seconds_input = gr.Slider(
                minimum=round(MIN_FRAMES_MODEL/FIXED_FPS,1), 
                maximum=round(MAX_FRAMES_MODEL/FIXED_FPS,1), 
                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."
            )
            
            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=DEFAULT_H_SLIDER_VALUE, label=f"Output Height (multiple of {MOD_VALUE})")
                    width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label=f"Output Width (multiple of {MOD_VALUE})")
                steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=6, label="Inference Steps") 
                guidance_scale_input = gr.Slider(minimum=0.0, maximum=5.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)
            with gr.Accordion("Mode Information", open=False):
                gr.Markdown("""
                **Processing Modes:**
                - **Ref2V**: Uses uploaded images as style references for video generation. All frames are generated based on the reference images.
                - **FLF2V**: First-Last Frame mode - uses first and last images as keyframes and generates the frames in between (requires exactly 2 images)
                - **Random2V**: Places uploaded images at specific frames in the video and generates the rest. Images are distributed evenly across the video duration.
                
                **Note**: VACE pipeline supports advanced conditioning with masks and reference images for more control over generation.
                """)

    # Update prompt when mode changes
    mode_radio.change(
        fn=update_prompt_from_mode,
        inputs=[mode_radio],
        outputs=[prompt_input]
    )
    
    # Update dimensions when gallery changes
    gallery_component.change(
        fn=handle_gallery_upload_for_dims_wan,
        inputs=[gallery_component, height_input, width_input],
        outputs=[height_input, width_input]
    )
    
    ui_inputs = [
        gallery_component, mode_radio, 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])

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