app.py

import os
import shutil
import sys
import subprocess
import asyncio
import uuid
from typing import Sequence, Mapping, Any, Union
from huggingface_hub import hf_hub_download
import spaces

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...")
text_encoder_repo = 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")
print(text_encoder_repo)
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. Let ComfyUI's main.py handle all initial setup ---
print("Importing ComfyUI's main.py for setup...")
import main
print("ComfyUI main imported.")

# --- 3. Now we can import the rest of the necessary modules ---
import torch
import gradio as gr
from comfy import model_management
from PIL import Image
import random
import nodes

# --- 4. Manually trigger the node initialization ---
print("Initializing ComfyUI nodes...")
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
loop.run_until_complete(nodes.init_extra_nodes())
print("Nodes initialized.")

# --- Helper function ---
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
    try:
        return obj[index]
    except KeyError:
        return obj["result"][index]


# --- ZeroGPU: Pre-load models and instantiate nodes globally ---
cliploader = nodes.NODE_CLASS_MAPPINGS["CLIPLoader"]()
cliptextencode = nodes.NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
unetloader = nodes.NODE_CLASS_MAPPINGS["UNETLoader"]()
vaeloader = nodes.NODE_CLASS_MAPPINGS["VAELoader"]()
clipvisionloader = nodes.NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
loadimage = nodes.NODE_CLASS_MAPPINGS["LoadImage"]()
clipvisionencode = nodes.NODE_CLASS_MAPPINGS["CLIPVisionEncode"]()
loraloadermodelonly = nodes.NODE_CLASS_MAPPINGS["LoraLoaderModelOnly"]()
modelsamplingsd3 = nodes.NODE_CLASS_MAPPINGS["ModelSamplingSD3"]()
pathchsageattentionkj = nodes.NODE_CLASS_MAPPINGS["PathchSageAttentionKJ"]()
wanfirstlastframetovideo = nodes.NODE_CLASS_MAPPINGS["WanFirstLastFrameToVideo"]()
ksampleradvanced = nodes.NODE_CLASS_MAPPINGS["KSamplerAdvanced"]()
vaedecode = nodes.NODE_CLASS_MAPPINGS["VAEDecode"]()
createvideo = nodes.NODE_CLASS_MAPPINGS["CreateVideo"]()
savevideo = nodes.NODE_CLASS_MAPPINGS["SaveVideo"]()

cliploader_38 = cliploader.load_clip(clip_name="umt5_xxl_fp8_e4m3fn_scaled.safetensors", type="wan", device="cpu")
unetloader_37_low_noise = unetloader.load_unet(unet_name="wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors", weight_dtype="default")
unetloader_91_high_noise = unetloader.load_unet(unet_name="wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors", weight_dtype="default")
vaeloader_39 = vaeloader.load_vae(vae_name="wan_2.1_vae.safetensors")
clipvisionloader_49 = clipvisionloader.load_clip(clip_name="clip_vision_h.safetensors")

loraloadermodelonly_94_high = 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(unetloader_91_high_noise, 0))
loraloadermodelonly_95_low = 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(unetloader_37_low_noise, 0))
modelsamplingsd3_93_low = modelsamplingsd3.patch(shift=8, model=get_value_at_index(loraloadermodelonly_95_low, 0))
pathchsageattentionkj_98_low = pathchsageattentionkj.patch(sage_attention="auto", model=get_value_at_index(modelsamplingsd3_93_low, 0))
modelsamplingsd3_79_high = modelsamplingsd3.patch(shift=8, model=get_value_at_index(loraloadermodelonly_94_high, 0))
pathchsageattentionkj_96_high = pathchsageattentionkj.patch(sage_attention="auto", model=get_value_at_index(modelsamplingsd3_79_high, 0))

model_loaders = [cliploader_38, unetloader_37_low_noise, unetloader_91_high_noise, vaeloader_39, clipvisionloader_49, loraloadermodelonly_94_high, loraloadermodelonly_95_low]
valid_models = [getattr(loader[0], 'patcher', loader[0]) for loader in model_loaders if not isinstance(loader[0], dict) and not isinstance(getattr(loader[0], 'patcher', None), dict)]
model_management.load_models_gpu(valid_models)

# --- App Logic ---
def calculate_dimensions(width, height):
    if width == height: return 480, 480
    if width > height: new_width, new_height = 832, int(height * (832 / width))
    else: new_height, new_width = 832, int(width * (832 / height))
    return (new_width // 16) * 16, (new_height // 16) * 16

@spaces.GPU(duration=120)
def generate_video(prompt, first_image_path, last_image_path, duration_seconds, progress=gr.Progress(track_tqdm=True)):
    # Create a temporary directory for resized images
    temp_dir = "input"
    os.makedirs(temp_dir, exist_ok=True)
    
    with torch.inference_mode():
        
        # --- Python Image Preprocessing using Pillow ---
        print("Preprocessing images with Pillow...")
        with Image.open(first_image_path) as img:
            orig_width, orig_height = img.size
        
        target_width, target_height = calculate_dimensions(orig_width, orig_height)
        
        # Resize first image
        with Image.open(first_image_path) as img:
            img_resized = img.resize((target_width, target_height), Image.Resampling.LANCZOS)
            resized_first_path = os.path.join(temp_dir, f"first_frame_resized_{uuid.uuid4().hex}.png")
            print(resized_first_path)
            img_resized.save(resized_first_path)

        # Resize second image to match the target dimensions
        with Image.open(last_image_path) as img:
            img_resized = img.resize((target_width, target_height), Image.Resampling.LANCZOS)
            resized_last_path = os.path.join(temp_dir, f"last_frame_resized_{uuid.uuid4().hex}.png")
            print(resized_last_path)
            img_resized.save(resized_last_path)
        print(f"Images resized to {target_width}x{target_height} and saved temporarily.")
        # --- End Preprocessing ---

        FPS, MAX_FRAMES = 16, 81
        length_in_frames = max(1, min(int(duration_seconds * FPS), MAX_FRAMES))
        print(f"Requested duration: {duration_seconds}s. Calculated frames: {length_in_frames}")
        
        # Load the pre-processed images into ComfyUI
        loaded_first_image = loadimage.load_image(image=os.path.basename(resized_first_path))
        loaded_last_image = loadimage.load_image(image=os.path.basename(resized_last_path))
        
        cliptextencode_6 = cliptextencode.encode(text=prompt, clip=get_value_at_index(cliploader_38, 0))
        cliptextencode_7_negative = cliptextencode.encode(text="low quality, worst quality, jpeg artifacts, ugly, deformed, blurry", clip=get_value_at_index(cliploader_38, 0))
        clipvisionencode_51 = clipvisionencode.encode(crop="none", clip_vision=get_value_at_index(clipvisionloader_49, 0), image=get_value_at_index(loaded_first_image, 0))
        clipvisionencode_87 = clipvisionencode.encode(crop="none", clip_vision=get_value_at_index(clipvisionloader_49, 0), image=get_value_at_index(loaded_last_image, 0))

        wanfirstlastframetovideo_83 = wanfirstlastframetovideo.EXECUTE_NORMALIZED(width=target_width, height=target_height, length=length_in_frames, batch_size=1, positive=get_value_at_index(cliptextencode_6, 0), negative=get_value_at_index(cliptextencode_7_negative, 0), vae=get_value_at_index(vaeloader_39, 0), clip_vision_start_image=get_value_at_index(clipvisionencode_51, 0), clip_vision_end_image=get_value_at_index(clipvisionencode_87, 0), start_image=get_value_at_index(loaded_first_image, 0), end_image=get_value_at_index(loaded_last_image, 0))

        ksampler_positive = get_value_at_index(wanfirstlastframetovideo_83, 0)
        ksampler_negative = get_value_at_index(wanfirstlastframetovideo_83, 1)
        ksampler_latent = get_value_at_index(wanfirstlastframetovideo_83, 2)
        
        ksampleradvanced_101 = 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(pathchsageattentionkj_96_high, 0), positive=ksampler_positive, negative=ksampler_negative, latent_image=ksampler_latent)
        ksampleradvanced_102 = 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(pathchsageattentionkj_98_low, 0), positive=ksampler_positive, negative=ksampler_negative, latent_image=get_value_at_index(ksampleradvanced_101, 0))

        vaedecode_8 = vaedecode.decode(samples=get_value_at_index(ksampleradvanced_102, 0), vae=get_value_at_index(vaeloader_39, 0))
        createvideo_104 = createvideo.create_video(fps=16, images=get_value_at_index(vaedecode_8, 0))
        savevideo_103 = savevideo.save_video(filename_prefix="ComfyUI_Video", format="mp4", codec="h264", video=get_value_at_index(createvideo_104, 0))
        print("** DEBUG ** ", savevideo_103)
        return f"output/{savevideo_103['ui']['images'][0]['filename']}"

# --- Gradio Interface ---
with gr.Blocks() 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/) on ZeroGPU")
    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(label="Prompt", value="a man dancing in the street, cinematic")
            duration_slider = gr.Slider(minimum=1.0, maximum=5.0, value=2.0, step=0.1, label="Video Duration (seconds)")
            with gr.Row():
                first_image = gr.Image(label="First Frame", type="filepath")
                last_image = gr.Image(label="Last Frame", type="filepath")
            generate_btn = gr.Button("Generate Video")
        with gr.Column(scale=2):
            output_video = gr.Video(label="Generated Video")
    generate_btn.click(fn=generate_video, inputs=[prompt_input, first_image, last_image, duration_slider], outputs=[output_video])
    gr.Examples(examples=[["a beautiful woman, cinematic", "examples/start.png", "examples/end.png", 2.5]], inputs=[prompt_input, first_image, last_image, duration_slider])

if __name__ == "__main__":
    if not os.path.exists("examples"): os.makedirs("examples")
    if not os.path.exists("examples/start.png"): Image.new('RGB', (512, 512), color='red').save('examples/start.png')
    if not os.path.exists("examples/end.png"): Image.new('RGB', (512, 512), color='blue').save('examples/end.png')
    # Set the input directory for LoadImage to find the temp files
    import folder_paths
    folder_paths.add_model_folder_path("input", "temp_resized")
    app.launch()