Update app.py

app.py

@@ -1,30 +1,35 @@
 import os
 import shutil
-import random
 import sys
+import subprocess
+import asyncio
+import uuid
+import random
 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
+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
 
-# --- 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")
@@ -35,241 +40,151 @@ hf_hub_download_local(repo_id="Kijai/WanVideo_comfy", filename="Wan22-Lightning/
 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
+# --- 2. ComfyUI Backend Initialization ---
 
 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
-
+    """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)
-    if parent_directory == path:
-        return None
-
-    return find_path(name, parent_directory)
-
+    return find_path(name, parent_directory) if parent_directory != path else None
 
 def add_comfyui_directory_to_sys_path() -> None:
-    """
-    Add 'ComfyUI' to the sys.path
-    """
+    """Adds the ComfyUI directory to sys.path for imports."""
     comfyui_path = find_path("ComfyUI")
-    if comfyui_path is not None and os.path.isdir(comfyui_path):
+    if comfyui_path 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.")
+    """Initializes ComfyUI's folder_paths with custom paths."""
+    from main import apply_custom_paths
+    apply_custom_paths()
 
 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
-
+    """Initializes all ComfyUI custom nodes."""
+    import nodes
     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))
-
+    loop.run_until_complete(nodes.init_extra_nodes(init_custom_nodes=True))
 
-# --- Model Loading and Caching ---
-MODELS_AND_NODES = {}
-
-print("Setting up ComfyUI paths...")
+print("Setting up ComfyUI paths and nodes...")
 add_comfyui_directory_to_sys_path()
 add_extra_model_paths()
-
-print("Importing custom nodes...")
 import_custom_nodes()
+print("ComfyUI setup complete.")
+
+
+# --- 3. Global Model & Node Loading and Patching ---
 
-# 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
+from comfy import model_management
 
-print("Loading models into memory. This may take a few minutes...")
+# Set VRAM mode to HIGH to prevent models from being offloaded from GPU after use.
+model_management.vram_state = model_management.VRAMState.HIGH_VRAM
 
-# 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"
-)
+MODELS_AND_NODES = {}
 
-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",
-)
+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")
 
-# Load LoRAs
-loraloadermodelonly = NODE_CLASS_MAPPINGS["LoraLoaderModelOnly"]()
-MODELS_AND_NODES["model_low_noise"] = loraloadermodelonly.load_lora_model_only(
+# 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),
-)
-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),
-)
+    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))
 
-# Load Vision model
-clipvisionloader = NODE_CLASS_MAPPINGS["CLIPVisionLoader"]()
-MODELS_AND_NODES["clip_vision"] = clipvisionloader.load_clip(
-    clip_name="clip_vision_h.safetensors"
-)
+# --- 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 required node classes
+# 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["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 = [
+# 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"],  # This is the UNET + LoRA
-    MODELS_AND_NODES["model_high_noise"], # This is the other UNET + LoRA
+    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
-])
-print("All models loaded successfully!")
+    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))
 
-# --- Main Video Generation Logic ---
 @spaces.GPU(duration=120)
 def generate_video(
     start_image_pil,
@@ -280,65 +195,59 @@ def generate_video(
     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.
+    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
-    
-    # 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}")
-    
+
+    # --- 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_noise = MODELS_AND_NODES["model_low_noise"]
-    model_high_noise = MODELS_AND_NODES["model_high_noise"]
+    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"]
-    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:
+    # --- 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 = start_file.name
-        end_image_path = 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...")
-        # --- Workflow execution ---
+        
+        # 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)
-        )
+        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(
@@ -352,30 +261,27 @@ def generate_video(
             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)...")
+        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=get_value_at_index(initial_latents, 0),
-            negative=get_value_at_index(initial_latents, 1),
-            latent_image=get_value_at_index(initial_latents, 2),
+            positive=ksampler_positive,
+            negative=ksampler_negative,
+            latent_image=ksampler_latent,
         )
         
-        progress(0.7, desc="Running KSampler (Step 2/2)...")
+        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=get_value_at_index(initial_latents, 0),
-            negative=get_value_at_index(initial_latents, 1),
+            positive=ksampler_positive,
+            negative=ksampler_negative,
             latent_image=get_value_at_index(latent_step1, 0),
         )
 
@@ -385,15 +291,23 @@ def generate_video(
         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 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!")
-        return f"output/{save_result['ui']['images'][0]['filename']}"
 
+        # --- 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 = '''