app.py

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

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

# --- 1. Model Download and Setup ---

def hf_hub_download_local(repo_id, filename, local_dir, **kwargs):
    """Downloads a file from Hugging Face Hub and symlinks it to a local directory."""
    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)
    
    # Remove existing symlink or file to avoid errors
    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

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


# --- 2. ComfyUI Backend Initialization ---

def find_path(name: str, path: str = None) -> str:
    """Recursively finds a directory with a given name."""
    if path is None: path = os.getcwd()
    if name in os.listdir(path): return os.path.join(path, name)
    parent_directory = os.path.dirname(path)
    return find_path(name, parent_directory) if parent_directory != path else None

def add_comfyui_directory_to_sys_path() -> None:
    """Adds the ComfyUI directory to sys.path for imports."""
    comfyui_path = find_path("ComfyUI")
    if comfyui_path and os.path.isdir(comfyui_path):
        sys.path.append(comfyui_path)
        print(f"'{comfyui_path}' added to sys.path")

def add_extra_model_paths() -> None:
    """Initializes ComfyUI's folder_paths with custom paths."""
    from main import apply_custom_paths
    apply_custom_paths()

def import_custom_nodes() -> None:
    """Initializes all ComfyUI custom nodes."""
    import nodes
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(nodes.init_extra_nodes(init_custom_nodes=True))

print("Setting up ComfyUI paths and nodes...")
add_comfyui_directory_to_sys_path()
add_extra_model_paths()
import_custom_nodes()
print("ComfyUI setup complete.")


# --- 3. Global Model & Node Loading and Patching ---

from nodes import NODE_CLASS_MAPPINGS
import folder_paths
from comfy import model_management

# Set VRAM mode to HIGH to prevent models from being offloaded from GPU after use.
# model_management.vram_state = model_management.VRAMState.HIGH_VRAM

MODELS_AND_NODES = {}

def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
    """Helper to safely access outputs from ComfyUI nodes, which are often tuples."""
    try:
        return obj[index]
    except (KeyError, TypeError):
        # Fallback for custom nodes that might return a dictionary with a 'result' key
        if isinstance(obj, Mapping) and "result" in obj:
            return obj["result"][index]
        raise

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

# Instantiate Node Classes that will be used for loading and patching
cliploader = NODE_CLASS_MAPPINGS["CLIPLoader"]()
unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
loraloadermodelonly = NODE_CLASS_MAPPINGS["LoraLoaderModelOnly"]()
modelsamplingsd3 = NODE_CLASS_MAPPINGS["ModelSamplingSD3"]()
pathchsageattentionkj = NODE_CLASS_MAPPINGS["PathchSageAttentionKJ"]()

# Load base models into CPU RAM initially
MODELS_AND_NODES["clip"] = cliploader.load_clip(clip_name="umt5_xxl_fp8_e4m3fn_scaled.safetensors", type="wan")
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")
MODELS_AND_NODES["vae"] = vaeloader.load_vae(vae_name="wan_2.1_vae.safetensors")
MODELS_AND_NODES["clip_vision"] = clipvisionloader.load_clip(clip_name="clip_vision_h.safetensors")

# Chain all patching operations together for the final models
print("Applying all patches to models...")

# --- Low Noise Model Chain ---
model_low_with_lora = 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))
model_low_patched = modelsamplingsd3.patch(shift=8, model=get_value_at_index(model_low_with_lora, 0))
MODELS_AND_NODES["model_low_noise"] = pathchsageattentionkj.patch(sage_attention="auto", model=get_value_at_index(model_low_patched, 0))

# --- High Noise Model Chain ---
model_high_with_lora = 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))
model_high_patched = modelsamplingsd3.patch(shift=8, model=get_value_at_index(model_high_with_lora, 0))
MODELS_AND_NODES["model_high_noise"] = pathchsageattentionkj.patch(sage_attention="auto", model=get_value_at_index(model_high_patched, 0))

# Instantiate all other node classes ONCE and store them
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["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"]()

# Move all final, fully-patched models to the GPU
print("Moving final models to GPU...")
model_loaders_final = [
    MODELS_AND_NODES["clip"],
#    MODELS_AND_NODES["vae"],
    MODELS_AND_NODES["model_low_noise"],
    MODELS_AND_NODES["model_high_noise"],
    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_final
], force_patch_weights=True) # force_patch_weights permanently merges the LoRA

print("All models loaded, patched, and on GPU. Gradio app is ready.")


# --- 4. Application Logic and Gradio Interface ---

def calculate_video_dimensions(width, height, max_size=832, min_size=480):
    """Calculates video dimensions, ensuring they are multiples of 16."""
    if width == height:
        return min_size, min_size
    aspect_ratio = width / height
    if width > height:
        video_width = max_size
        video_height = int(max_size / aspect_ratio)
    else:
        video_height = max_size
        video_width = int(max_size * aspect_ratio)
    video_width = max(16, round(video_width / 16) * 16)
    video_height = max(16, round(video_height / 16) * 16)
    return video_width, video_height

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="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走,过曝，",
    duration=33,
    progress=gr.Progress(track_tqdm=True)
):
    """
    Generates a video by interpolating between a start and end image, guided by a text prompt.
    This function relies on globally pre-loaded models and pre-instantiated ComfyUI nodes.
    """
    FPS = 16

    # --- 1. Retrieve Pre-loaded and Pre-patched Models & Node Instances ---
    # These are not re-instantiated; we are just getting references to the global objects.
    clip = MODELS_AND_NODES["clip"]
    vae = MODELS_AND_NODES["vae"]
    model_low_final = MODELS_AND_NODES["model_low_noise"]
    model_high_final = 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"]
    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"]

    # --- 2. Image Preprocessing for the Current Run ---
    print("Preprocessing images with Pillow...")
    processed_start_image = start_image_pil.copy()
    processed_end_image = resize_and_crop_to_match(end_image_pil, start_image_pil)
    video_width, video_height = calculate_video_dimensions(processed_start_image.width, processed_start_image.height)

    # Save processed images to temporary files for the LoadImage node
    temp_dir = "input" # ComfyUI's default input directory
    os.makedirs(temp_dir, exist_ok=True)
    
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False, dir=temp_dir) as start_file, \
         tempfile.NamedTemporaryFile(suffix=".png", delete=False, dir=temp_dir) as end_file:
        processed_start_image.save(start_file.name)
        processed_end_image.save(end_file.name)
        start_image_path = os.path.basename(start_file.name)
        end_image_path = os.path.basename(end_file.name)
    print(f"Images resized to {video_width}x{video_height} and saved temporarily.")
    
    # --- 3. Execute the ComfyUI Workflow in Inference Mode ---
    with torch.inference_mode():
        progress(0.1, desc="Encoding text and images...")
        
        # Encode prompts and vision models
        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),
        )

        ksampler_positive = get_value_at_index(initial_latents, 0)
        ksampler_negative = get_value_at_index(initial_latents, 1)
        ksampler_latent = get_value_at_index(initial_latents, 2)
        
        progress(0.5, desc="Denoising (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=ksampler_positive,
            negative=ksampler_negative,
            latent_image=ksampler_latent,
        )
        
        progress(0.7, desc="Denoising (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=ksampler_positive,
            negative=ksampler_negative,
            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 default 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!")

        # --- 4. Cleanup and Return ---
        try:
            os.remove(start_file.name)
            os.remove(end_file.name)
        except Exception as e:
            print(f"Error cleaning up temporary files: {e}")

        # Gradio video component expects a filepath relative to the root of the app
        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=False):
                    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)