app.py

import os
import shutil
import random
import sys
import tempfile
from typing import Sequence, Mapping, Any, Union

import spaces
import torch
import gradio as gr
from PIL import Image
from huggingface_hub import hf_hub_download
from comfy import model_management

def hf_hub_download_local(repo_id, filename, local_dir, **kwargs):
    downloaded_path = hf_hub_download(repo_id=repo_id, filename=filename, **kwargs)
    os.makedirs(local_dir, exist_ok=True)
    base_filename = os.path.basename(filename)
    target_path = os.path.join(local_dir, base_filename)
    
    if os.path.exists(target_path) or os.path.islink(target_path):
        os.remove(target_path)
    
    os.symlink(downloaded_path, target_path)
    return target_path

# --- Model Downloads ---
print("Downloading models from Hugging Face Hub...")
hf_hub_download_local(repo_id="Comfy-Org/Wan_2.1_ComfyUI_repackaged", filename="split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors", local_dir="models/text_encoders")
hf_hub_download_local(repo_id="Comfy-Org/Wan_2.2_ComfyUI_Repackaged", filename="split_files/diffusion_models/wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors", local_dir="models/unet")
hf_hub_download_local(repo_id="Comfy-Org/Wan_2.2_ComfyUI_Repackaged", filename="split_files/diffusion_models/wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors", local_dir="models/unet")
hf_hub_download_local(repo_id="Comfy-Org/Wan_2.1_ComfyUI_repackaged", filename="split_files/vae/wan_2.1_vae.safetensors", local_dir="models/vae")
hf_hub_download_local(repo_id="Comfy-Org/Wan_2.1_ComfyUI_repackaged", filename="split_files/clip_vision/clip_vision_h.safetensors", local_dir="models/clip_vision")
hf_hub_download_local(repo_id="Kijai/WanVideo_comfy", filename="Wan22-Lightning/Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors", local_dir="models/loras")
hf_hub_download_local(repo_id="Kijai/WanVideo_comfy", filename="Wan22-Lightning/Wan2.2-Lightning_I2V-A14B-4steps-lora_LOW_fp16.safetensors", local_dir="models/loras")
print("Downloads complete.")

model_management.vram_state = model_management.VRAMState.HIGH_VRAM

# --- Image Processing Functions ---
def calculate_video_dimensions(width, height, max_size=832, min_size=480):
    """
    Calculate video dimensions based on input image size.
    Larger dimension becomes max_size, smaller becomes proportional.
    If square, use min_size x min_size.
    Results are rounded to nearest multiple of 16.
    """
    # Handle square images
    if width == height:
        video_width = min_size
        video_height = min_size
    else:
        # Calculate aspect ratio
        aspect_ratio = width / height
        
        if width > height:
            # Landscape orientation
            video_width = max_size
            video_height = int(max_size / aspect_ratio)
        else:
            # Portrait orientation
            video_height = max_size
            video_width = int(max_size * aspect_ratio)
    
    # Round to nearest multiple of 16
    video_width = round(video_width / 16) * 16
    video_height = round(video_height / 16) * 16
    
    # Ensure minimum size
    video_width = max(video_width, 16)
    video_height = max(video_height, 16)
    
    return video_width, video_height

def resize_and_crop_to_match(target_image, reference_image):
    """
    Resize and center crop target_image to match reference_image dimensions.
    """
    ref_width, ref_height = reference_image.size
    target_width, target_height = target_image.size
    
    # Calculate scaling factor to ensure target covers reference dimensions
    scale = max(ref_width / target_width, ref_height / target_height)
    
    # Resize target image
    new_width = int(target_width * scale)
    new_height = int(target_height * scale)
    resized = target_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    # Center crop to match reference dimensions
    left = (new_width - ref_width) // 2
    top = (new_height - ref_height) // 2
    right = left + ref_width
    bottom = top + ref_height
    
    cropped = resized.crop((left, top, right, bottom))
    return cropped

# --- Boilerplate code from the original script ---
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
    """Returns the value at the given index of a sequence or mapping.

    If the object is a sequence (like list or string), returns the value at the given index.
    If the object is a mapping (like a dictionary), returns the value at the index-th key.

    Some return a dictionary, in these cases, we look for the "results" key

    Args:
        obj (Union[Sequence, Mapping]): The object to retrieve the value from.
        index (int): The index of the value to retrieve.

    Returns:
        Any: The value at the given index.

    Raises:
        IndexError: If the index is out of bounds for the object and the object is not a mapping.
    """
    try:
        return obj[index]
    except KeyError:
        # This is a fallback for custom node outputs that might be dictionaries
        if isinstance(obj, Mapping) and "result" in obj:
            return obj["result"][index]
        raise

def find_path(name: str, path: str = None) -> str:
    """
    Recursively looks at parent folders starting from the given path until it finds the given name.
    Returns the path as a Path object if found, or None otherwise.
    """
    if path is None:
        path = os.getcwd()

    if name in os.listdir(path):
        path_name = os.path.join(path, name)
        print(f"'{name}' found: {path_name}")
        return path_name

    parent_directory = os.path.dirname(path)
    if parent_directory == path:
        return None

    return find_path(name, parent_directory)


def add_comfyui_directory_to_sys_path() -> None:
    """
    Add 'ComfyUI' to the sys.path
    """
    comfyui_path = find_path("ComfyUI")
    if comfyui_path is not None and os.path.isdir(comfyui_path):
        sys.path.append(comfyui_path)
        print(f"'{comfyui_path}' added to sys.path")
    else:
        print("Could not find ComfyUI directory. Please run from a parent folder of ComfyUI.")

def add_extra_model_paths() -> None:
    """
    Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path.
    """
    try:
        from main import load_extra_path_config
    except ImportError:
        print(
            "Could not import load_extra_path_config from main.py. This might be okay if you don't use it."
        )
        return

    extra_model_paths = find_path("extra_model_paths.yaml")
    if extra_model_paths is not None:
        load_extra_path_config(extra_model_paths)
    else:
        print("Could not find an optional 'extra_model_paths.yaml' config file.")

def import_custom_nodes() -> None:
    """Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS
    This function sets up a new asyncio event loop, initializes the PromptServer,
    creates a PromptQueue, and initializes the custom nodes.
    """
    import asyncio
    import execution
    from nodes import init_extra_nodes
    import server

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    server_instance = server.PromptServer(loop)
    execution.PromptQueue(server_instance)
    loop.run_until_complete(init_extra_nodes(init_custom_nodes=True))


# --- Model Loading and Caching ---
MODELS_AND_NODES = {}

print("Setting up ComfyUI paths...")
add_comfyui_directory_to_sys_path()
add_extra_model_paths()

print("Importing custom nodes...")
import_custom_nodes()

# Now that paths are set up, we can import from nodes
from nodes import NODE_CLASS_MAPPINGS
global folder_paths # Make folder_paths globally accessible
import folder_paths

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

# Load Text-to-Image models (CLIP, UNETs, VAE)
cliploader = NODE_CLASS_MAPPINGS["CLIPLoader"]()
MODELS_AND_NODES["clip"] = cliploader.load_clip(
    clip_name="umt5_xxl_fp8_e4m3fn_scaled.safetensors", type="wan", device="cpu"
)

unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
unet_low_noise = unetloader.load_unet(
    unet_name="wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors",
    weight_dtype="default",
)
unet_high_noise = unetloader.load_unet(
    unet_name="wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors",
    weight_dtype="default",
)

vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
MODELS_AND_NODES["vae"] = vaeloader.load_vae(vae_name="wan_2.1_vae.safetensors")

# Load LoRAs
loraloadermodelonly = NODE_CLASS_MAPPINGS["LoraLoaderModelOnly"]()
MODELS_AND_NODES["model_low_noise"] = loraloadermodelonly.load_lora_model_only(
    lora_name="Wan2.2-Lightning_I2V-A14B-4steps-lora_LOW_fp16.safetensors",
    strength_model=0.8,
    model=get_value_at_index(unet_low_noise, 0),
)
MODELS_AND_NODES["model_high_noise"] = loraloadermodelonly.load_lora_model_only(
    lora_name="Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors",
    strength_model=0.8,
    model=get_value_at_index(unet_high_noise, 0),
)

# Load Vision model
clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
MODELS_AND_NODES["clip_vision"] = clipvisionloader.load_clip(
    clip_name="clip_vision_h.safetensors"
)

# Instantiate all required node classes
MODELS_AND_NODES["CLIPTextEncode"] = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
MODELS_AND_NODES["LoadImage"] = NODE_CLASS_MAPPINGS["LoadImage"]()
MODELS_AND_NODES["CLIPVisionEncode"] = NODE_CLASS_MAPPINGS["CLIPVisionEncode"]()
MODELS_AND_NODES["ModelSamplingSD3"] = NODE_CLASS_MAPPINGS["ModelSamplingSD3"]()
MODELS_AND_NODES["PathchSageAttentionKJ"] = NODE_CLASS_MAPPINGS["PathchSageAttentionKJ"]()
MODELS_AND_NODES["WanFirstLastFrameToVideo"] = NODE_CLASS_MAPPINGS["WanFirstLastFrameToVideo"]()
MODELS_AND_NODES["KSamplerAdvanced"] = NODE_CLASS_MAPPINGS["KSamplerAdvanced"]()
MODELS_AND_NODES["VAEDecode"] = NODE_CLASS_MAPPINGS["VAEDecode"]()
MODELS_AND_NODES["CreateVideo"] = NODE_CLASS_MAPPINGS["CreateVideo"]()
MODELS_AND_NODES["SaveVideo"] = NODE_CLASS_MAPPINGS["SaveVideo"]()

print("Pre-loading main models onto GPU...")
model_loaders = [
    MODELS_AND_NODES["clip"],
    MODELS_AND_NODES["vae"],
    MODELS_AND_NODES["model_low_noise"],  # This is the UNET + LoRA
    MODELS_AND_NODES["model_high_noise"], # This is the other UNET + LoRA
    MODELS_AND_NODES["clip_vision"],
]
model_management.load_models_gpu([
    loader[0].patcher if hasattr(loader[0], 'patcher') else loader[0] for loader in model_loaders
])
print("All models loaded successfully!")

# --- Main Video Generation Logic ---
@spaces.GPU(duration=120)
def generate_video(
    start_image_pil,
    end_image_pil,
    prompt,
    negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走,过曝，",
    duration=33,
    progress=gr.Progress(track_tqdm=True)
):
    """
    The main function to generate a video based on user inputs.
    This function is called every time the user clicks the 'Generate' button.
    """
    FPS = 16
    
    # Process images: resize and crop second image to match first
    # The first image determines the dimensions
    processed_start_image = start_image_pil.copy()
    processed_end_image = resize_and_crop_to_match(end_image_pil, start_image_pil)
    
    # Calculate video dimensions based on the first image
    video_width, video_height = calculate_video_dimensions(
        processed_start_image.width, 
        processed_start_image.height
    )
    
    print(f"Input image size: {processed_start_image.width}x{processed_start_image.height}")
    print(f"Video dimensions: {video_width}x{video_height}")
    
    clip = MODELS_AND_NODES["clip"]
    vae = MODELS_AND_NODES["vae"]
    model_low_noise = MODELS_AND_NODES["model_low_noise"]
    model_high_noise = MODELS_AND_NODES["model_high_noise"]
    clip_vision = MODELS_AND_NODES["clip_vision"]
    
    cliptextencode = MODELS_AND_NODES["CLIPTextEncode"]
    loadimage = MODELS_AND_NODES["LoadImage"]
    clipvisionencode = MODELS_AND_NODES["CLIPVisionEncode"]
    modelsamplingsd3 = MODELS_AND_NODES["ModelSamplingSD3"]
    pathchsageattentionkj = MODELS_AND_NODES["PathchSageAttentionKJ"]
    wanfirstlastframetovideo = MODELS_AND_NODES["WanFirstLastFrameToVideo"]
    ksampleradvanced = MODELS_AND_NODES["KSamplerAdvanced"]
    vaedecode = MODELS_AND_NODES["VAEDecode"]
    createvideo = MODELS_AND_NODES["CreateVideo"]
    savevideo = MODELS_AND_NODES["SaveVideo"]

    # Save processed images to temporary files
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as start_file, \
         tempfile.NamedTemporaryFile(suffix=".png", delete=False) as end_file:
        processed_start_image.save(start_file.name)
        processed_end_image.save(end_file.name)
        start_image_path = start_file.name
        end_image_path = end_file.name

    with torch.inference_mode():
        progress(0.1, desc="Encoding text and images...")
        # --- Workflow execution ---
        positive_conditioning = cliptextencode.encode(text=prompt, clip=get_value_at_index(clip, 0))
        negative_conditioning = cliptextencode.encode(text=negative_prompt, clip=get_value_at_index(clip, 0))

        start_image_loaded = loadimage.load_image(image=start_image_path)
        end_image_loaded = loadimage.load_image(image=end_image_path)
        
        clip_vision_encoded_start = clipvisionencode.encode(
            crop="none", clip_vision=get_value_at_index(clip_vision, 0), image=get_value_at_index(start_image_loaded, 0)
        )
        clip_vision_encoded_end = clipvisionencode.encode(
            crop="none", clip_vision=get_value_at_index(clip_vision, 0), image=get_value_at_index(end_image_loaded, 0)
        )

        progress(0.2, desc="Preparing initial latents...")
        initial_latents = wanfirstlastframetovideo.EXECUTE_NORMALIZED(
            width=video_width, height=video_height, length=duration, batch_size=1,
            positive=get_value_at_index(positive_conditioning, 0),
            negative=get_value_at_index(negative_conditioning, 0),
            vae=get_value_at_index(vae, 0),
            clip_vision_start_image=get_value_at_index(clip_vision_encoded_start, 0),
            clip_vision_end_image=get_value_at_index(clip_vision_encoded_end, 0),
            start_image=get_value_at_index(start_image_loaded, 0),
            end_image=get_value_at_index(end_image_loaded, 0),
        )

        progress(0.3, desc="Patching models...")
        model_low_patched = modelsamplingsd3.patch(shift=8, model=get_value_at_index(model_low_noise, 0))
        model_low_final = pathchsageattentionkj.patch(sage_attention="auto", model=get_value_at_index(model_low_patched, 0))

        model_high_patched = modelsamplingsd3.patch(shift=8, model=get_value_at_index(model_high_noise, 0))
        model_high_final = pathchsageattentionkj.patch(sage_attention="auto", model=get_value_at_index(model_high_patched, 0))

        progress(0.5, desc="Running KSampler (Step 1/2)...")
        latent_step1 = ksampleradvanced.sample(
            add_noise="enable", noise_seed=random.randint(1, 2**64), steps=8, cfg=1,
            sampler_name="euler", scheduler="simple", start_at_step=0, end_at_step=4,
            return_with_leftover_noise="enable", model=get_value_at_index(model_high_final, 0),
            positive=get_value_at_index(initial_latents, 0),
            negative=get_value_at_index(initial_latents, 1),
            latent_image=get_value_at_index(initial_latents, 2),
        )
        
        progress(0.7, desc="Running KSampler (Step 2/2)...")
        latent_step2 = ksampleradvanced.sample(
            add_noise="disable", noise_seed=random.randint(1, 2**64), steps=8, cfg=1,
            sampler_name="euler", scheduler="simple", start_at_step=4, end_at_step=10000,
            return_with_leftover_noise="disable", model=get_value_at_index(model_low_final, 0),
            positive=get_value_at_index(initial_latents, 0),
            negative=get_value_at_index(initial_latents, 1),
            latent_image=get_value_at_index(latent_step1, 0),
        )

        progress(0.8, desc="Decoding VAE...")
        decoded_images = vaedecode.decode(samples=get_value_at_index(latent_step2, 0), vae=get_value_at_index(vae, 0))

        progress(0.9, desc="Creating and saving video...")
        video_data = createvideo.create_video(fps=FPS, images=get_value_at_index(decoded_images, 0))
        
        # Save the video to ComfyUI's output directory
        save_result = savevideo.save_video(
            filename_prefix="GradioVideo", format="mp4", codec="h264",
            video=get_value_at_index(video_data, 0),
        )

        progress(1.0, desc="Done!")
        return f"output/{save_result['ui']['images'][0]['filename']}"



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("Running the [Wan 2.2 First/Last Frame ComfyUI workflow](https://www.reddit.com/r/StableDiffusion/comments/1me4306/psa_wan_22_does_first_frame_last_frame_out_of_the/) and the [lightx2v/Wan2.2-Lightning](https://huggingface.co/lightx2v/Wan2.2-Lightning) 8-step LoRA on ZeroGPU")
    
    with gr.Row():
        with gr.Column():
            with gr.Group():
                with gr.Row():
                    start_image = gr.Image(type="pil", label="Start Frame")
                    end_image = gr.Image(type="pil", label="End Frame")
                
                prompt = gr.Textbox(label="Prompt", info="Describe the transition between the two images")
    
                with gr.Accordion("Advanced Settings", open=False, visible=True):
                    duration = gr.Radio(
                        [("Short (2s)", 33), ("Mid (4s)", 66)],
                        value=33, 
                        label="Video Duration",
                        visible=False
                    )
                    negative_prompt = gr.Textbox(
                        label="Negative Prompt", 
                        value="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走,过曝，",
                        visible=False
                    )
                
                generate_button = gr.Button("Generate Video", variant="primary")
        
        with gr.Column():
            output_video = gr.Video(label="Generated Video", autoplay=True)

    generate_button.click(
        fn=generate_video,
        inputs=[start_image, end_image, prompt, negative_prompt, duration],
        outputs=output_video
    )

    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", "capybara.webp", "a dramatic dolly zoom"],
        ],
        inputs=[start_image, end_image, prompt],
        outputs=output_video,
        fn=generate_video,
        cache_examples="lazy",
    )

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