app.py

import torch
from diffusers import AutoencoderKLWan, WanVACEPipeline UniPCMultistepScheduler
from diffusers.utils import export_to_video
from transformers import CLIPVisionModel
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"

image_encoder = CLIPVisionModel.from_pretrained(MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float32)
vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
pipe = WanVACEPipeline.from_pretrained(
    MODEL_ID, vae=vae, image_encoder=image_encoder, torch_dtype=torch.bfloat16
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=2.0)
pipe.to("cuda")

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 = 24
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 81 

# Default prompts for different modes
MODE_PROMPTS = {
    "Ref2V": "",
    "FLF2V": "",
    "Random2V": ""
}

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]
        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 process_images_for_mode(images, mode):
    """Process images based on the selected mode"""
    if not images or len(images) == 0:
        return None
    
    if mode == "Ref2V":
        # Use the first image as reference
        return images[0]
    
    elif mode == "FLF2V":
        # First and Last Frame: blend or interpolate between first and last image
        if len(images) >= 2:
            return None
        else:
            return images[0]
    
    elif mode == "Random2V":
        # Randomly select one image from the gallery
        return images[0]
    
    return images[0]

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.")

    # Process images based on the selected mode
    input_image = process_images_for_mode(gallery_images, mode)
    
    if input_image is None:
        raise gr.Error("Failed to process images for the selected mode.")

    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)

    resized_image = input_image.resize((target_w, target_h))

    # Mode-specific processing can be added here if needed
    if mode == "FLF2V" and len(gallery_images) >= 2:
        # You can add special handling for FLF2V mode here
        # For example, use both first and last frames in some way
        pass

    with torch.inference_mode():
        output_frames_list = pipe(
            image=resized_image, 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("# Fast 4 steps Wan 2.1 I2V (14B) with CausVid LoRA - Multi-Image Gallery")
    gr.Markdown("[CausVid](https://github.com/tianweiy/CausVid) is a distilled version of Wan 2.1 to run faster in just 4-8 steps, [extracted as LoRA by Kijai](https://huggingface.co/Kijai/WanVideo_comfy/blob/main/Wan21_CausVid_14B_T2V_lora_rank32.safetensors) and is compatible with 🧨 diffusers")
    
    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 Image 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=True):
                gr.Markdown("""
                **Processing Modes:**
                - **Ref2V**: Uses the first image as reference for video generation
                - **FLF2V**: Blends first and last images for interpolation (requires at least 2 images)
                - **Random2V**: Randomly selects one image from the gallery for 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)