Update app.py

app.py

@@ -1,382 +1,302 @@
 import gradio as gr
-import spaces
 import torch
 import numpy as np
+import random
 import os
 import yaml
-import random
-from PIL import Image
-import imageio # For export_to_video and reading video frames
 from pathlib import Path
+import imageio
+import tempfile
+from PIL import Image
 from huggingface_hub import hf_hub_download
+import shutil
 
-# --- LTX-Video Imports (from your provided codebase) ---
-from ltx_video.pipelines.pipeline_ltx_video import (
-    ConditioningItem,
-    LTXVideoPipeline,
-    LTXMultiScalePipeline,
-)
-from ltx_video.models.autoencoders.vae_encode import vae_decode, vae_encode, un_normalize_latents, normalize_latents
+# --- Import necessary classes from the provided files ---
 from inference import (
     create_ltx_video_pipeline,
     create_latent_upsampler,
-    load_image_to_tensor_with_resize_and_crop, # Re-using for image conditioning
-    load_media_file, # Re-using for video conditioning
-    get_device,
+    load_image_to_tensor_with_resize_and_crop,
     seed_everething,
+    get_device,
     calculate_padding,
+    load_media_file
 )
+from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem, LTXMultiScalePipeline, LTXVideoPipeline
 from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
-from ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler
-# --- End LTX-Video Imports ---
-
-# --- Diffusers/Original utils (keeping export_to_video for convenience if it works) ---
-from diffusers.utils import export_to_video # Keep if it works with PIL list
-# ---
-
-# --- Global Configuration & Model Loading ---
-DEVICE = get_device()
-MODEL_DIR = "downloaded_models" # Directory to store downloaded models
-Path(MODEL_DIR).mkdir(parents=True, exist_ok=True)
 
-# Load YAML configuration
-YAML_CONFIG_PATH = "configs/ltxv-13b-0.9.7-distilled.yaml" # Place this file in the same directory
-with open(YAML_CONFIG_PATH, "r") as f:
-    PIPELINE_CONFIG_YAML = yaml.safe_load(f)
+# --- Global constants from user's request and YAML ---
+YAML_CONFIG_STRING = """
+pipeline_type: multi-scale
+checkpoint_path: "ltxv-13b-0.9.7-distilled.safetensors" # This will be replaced by the rc3 version
+downscale_factor: 0.6666666
+spatial_upscaler_model_path: "ltxv-spatial-upscaler-0.9.7.safetensors"
+stg_mode: "attention_values" # options: "attention_values", "attention_skip", "residual", "transformer_block"
+decode_timestep: 0.05
+decode_noise_scale: 0.025
+text_encoder_model_name_or_path: "PixArt-alpha/PixArt-XL-2-1024-MS"
+precision: "bfloat16"
+sampler: "from_checkpoint" # options: "uniform", "linear-quadratic", "from_checkpoint"
+prompt_enhancement_words_threshold: 120
+prompt_enhancer_image_caption_model_name_or_path: "MiaoshouAI/Florence-2-large-PromptGen-v2.0"
+prompt_enhancer_llm_model_name_or_path: "unsloth/Llama-3.2-3B-Instruct"
+stochastic_sampling: false
+
+first_pass:
+  timesteps: [1.0000, 0.9937, 0.9875, 0.9812, 0.9750, 0.9094, 0.7250]
+  guidance_scale: 1
+  stg_scale: 0
+  rescaling_scale: 1
+  skip_block_list: [42]
+
+second_pass:
+  timesteps: [0.9094, 0.7250, 0.4219]
+  guidance_scale: 1
+  stg_scale: 0
+  rescaling_scale: 1
+  skip_block_list: [42]
+"""
+PIPELINE_CONFIG_YAML = yaml.safe_load(YAML_CONFIG_STRING)
+
+# Model specific paths (to be downloaded)
+DISTILLED_MODEL_REPO = "LTX-Colab/LTX-Video-Preview"
+DISTILLED_MODEL_FILENAME = "ltxv-13b-0.9.7-distilled-rc3.safetensors"
+
+UPSCALER_REPO = "Lightricks/LTX-Video"
+# SPATIAL_UPSCALER_FILENAME will be taken from PIPELINE_CONFIG_YAML after it's loaded
+
+MAX_IMAGE_SIZE = PIPELINE_CONFIG_YAML.get("max_resolution", 1280) # Max width/height from UI
+MAX_NUM_FRAMES = 257 # From inference.py
+
+# --- Global variables for loaded models ---
+pipeline_instance = None
+latent_upsampler_instance = None
+current_device = get_device()
+models_dir = "downloaded_models_gradio" # Use a distinct name
+Path(models_dir).mkdir(parents=True, exist_ok=True)
+
+# Download models and update config paths
+print(f"Using device: {current_device}")
+print("Downloading models...")
+
+distilled_model_actual_path = hf_hub_download(
+    repo_id=DISTILLED_MODEL_REPO,
+    filename=DISTILLED_MODEL_FILENAME,
+    local_dir=models_dir,
+    local_dir_use_symlinks=False
+)
+PIPELINE_CONFIG_YAML["checkpoint_path"] = distilled_model_actual_path
+print(f"Distilled model downloaded to: {distilled_model_actual_path}")
 
-# Download and prepare model paths from YAML
-LTXV_MODEL_FILENAME = PIPELINE_CONFIG_YAML["checkpoint_path"]
 SPATIAL_UPSCALER_FILENAME = PIPELINE_CONFIG_YAML["spatial_upscaler_model_path"]
-TEXT_ENCODER_PATH = PIPELINE_CONFIG_YAML["text_encoder_model_name_or_path"] # This is usually a repo name
-
-try:
-    # Main LTX-Video model
-    if not os.path.isfile(os.path.join(MODEL_DIR, LTXV_MODEL_FILENAME)):
-        print(f"Downloading {LTXV_MODEL_FILENAME}...")
-        ltxv_checkpoint_path = hf_hub_download(
-            repo_id="LTX-Colab/LTX-Video-Preview", # Assuming the distilled model is also here or adjust repo_id
-            filename=LTXV_MODEL_FILENAME,
-            local_dir=MODEL_DIR,
-            repo_type="model",
-        )
-    else:
-        ltxv_checkpoint_path = os.path.join(MODEL_DIR, LTXV_MODEL_FILENAME)
-
-    # Spatial Upsampler model
-    if not os.path.isfile(os.path.join(MODEL_DIR, SPATIAL_UPSCALER_FILENAME)):
-        print(f"Downloading {SPATIAL_UPSCALER_FILENAME}...")
-        spatial_upsampler_path = hf_hub_download(
-            repo_id="Lightricks/LTX-Video",
-            filename=SPATIAL_UPSCALER_FILENAME,
-            local_dir=MODEL_DIR,
-            repo_type="model",
-        )
-    else:
-        spatial_upsampler_path = os.path.join(MODEL_DIR, SPATIAL_UPSCALER_FILENAME)
-except Exception as e:
-    print(f"Error downloading models: {e}")
-    print("Please ensure model files are correctly specified and accessible.")
-    # Depending on severity, you might want to exit or disable GPU features
-    # For now, we'll let it proceed and potentially fail later if paths are invalid.
-    ltxv_checkpoint_path = LTXV_MODEL_FILENAME # Fallback to filename if download fails
-    spatial_upsampler_path = SPATIAL_UPSCALER_FILENAME
-
-
-print(f"Using LTX-Video checkpoint: {ltxv_checkpoint_path}")
-print(f"Using Spatial Upsampler: {spatial_upsampler_path}")
-print(f"Using Text Encoder: {TEXT_ENCODER_PATH}")
-
-# Create LTX-Video pipeline
-pipe = create_ltx_video_pipeline(
-    ckpt_path=ltxv_checkpoint_path,
-    precision=PIPELINE_CONFIG_YAML["precision"],
-    text_encoder_model_name_or_path=TEXT_ENCODER_PATH,
-    sampler=PIPELINE_CONFIG_YAML["sampler"], # "from_checkpoint" or specific sampler
-    device=DEVICE,
-    enhance_prompt=False, # Assuming Gradio controls this, or set based on YAML later
-)#.to(torch.bfloat16)
-
-# Create Latent Upsampler
-latent_upsampler = create_latent_upsampler(
-    latent_upsampler_model_path=spatial_upsampler_path,
-    device=DEVICE
+spatial_upscaler_actual_path = hf_hub_download(
+    repo_id=UPSCALER_REPO,
+    filename=SPATIAL_UPSCALER_FILENAME,
+    local_dir=models_dir,
+    local_dir_use_symlinks=False
 )
-#latent_upsampler = latent_upsampler.to(torch.bfloat16)
+PIPELINE_CONFIG_YAML["spatial_upscaler_model_path"] = spatial_upscaler_actual_path
+print(f"Spatial upscaler model downloaded to: {spatial_upscaler_actual_path}")
 
-
-# Multi-scale pipeline (wrapper)
-multi_scale_pipe = LTXMultiScalePipeline(
-    video_pipeline=pipe,
-    latent_upsampler=latent_upsampler
+# Load pipelines
+print("Creating LTX Video pipeline...")
+pipeline_instance = create_ltx_video_pipeline(
+    ckpt_path=PIPELINE_CONFIG_YAML["checkpoint_path"],
+    precision=PIPELINE_CONFIG_YAML["precision"],
+    text_encoder_model_name_or_path=PIPELINE_CONFIG_YAML["text_encoder_model_name_or_path"],
+    sampler=PIPELINE_CONFIG_YAML["sampler"],
+    device=current_device,
+    enhance_prompt=False, # Prompt enhancement handled by UI choice / Gradio logic if desired
+    prompt_enhancer_image_caption_model_name_or_path=PIPELINE_CONFIG_YAML["prompt_enhancer_image_caption_model_name_or_path"],
+    prompt_enhancer_llm_model_name_or_path=PIPELINE_CONFIG_YAML["prompt_enhancer_llm_model_name_or_path"],
 )
-# --- End Global Configuration & Model Loading ---
-
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 2048 # Not strictly used here, but good to keep in mind
-
-
-def round_to_nearest_resolution_acceptable_by_vae(height, width, vae_scale_factor):
-    # print("before rounding",height, width)
-    height = height - (height % vae_scale_factor)
-    width = width - (width % vae_scale_factor)
-    # print("after rounding",height, width)
-    return height, width
-
-@spaces.GPU
-def generate(prompt,
-             negative_prompt,
-             image_path, # Gradio gives filepath for Image component
-             video_path, # Gradio gives filepath for Video component
-             height,
-             width,
-             mode,
-             steps,      # This will map to num_inference_steps for the first pass
-             num_frames,
-             frames_to_use,
-             seed,
-             randomize_seed,
-             guidance_scale,
-             improve_texture=False, progress=gr.Progress(track_tqdm=True)):
+print("LTX Video pipeline created.")
 
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    seed_everething(seed)
-    
-    generator = torch.Generator(device=DEVICE).manual_seed(seed)
+if PIPELINE_CONFIG_YAML.get("spatial_upscaler_model_path"):
+    print("Creating latent upsampler...")
+    latent_upsampler_instance = create_latent_upsampler(
+        PIPELINE_CONFIG_YAML["spatial_upscaler_model_path"],
+        device=current_device
+    )
+    print("Latent upsampler created.")
 
-    # --- Prepare conditioning items ---
-    conditioning_items_list = []
-    input_media_for_vid2vid = None # For the specific vid2vid mode in LTX pipeline
 
-    # Pad target dimensions
-    # VAE scale factor is typically 8 for spatial, but LTX might have its own specific factor.
-    # CausalVideoAutoencoder has spatial_downscale_factor and temporal_downscale_factor
-    vae_spatial_scale_factor = pipe.vae.spatial_downscale_factor
-    vae_temporal_scale_factor = pipe.vae.temporal_downscale_factor
+def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath,
+             height_ui, width_ui, mode,
+             ui_steps, num_frames_ui,
+             ui_frames_to_use,
+             seed_ui, randomize_seed, ui_guidance_scale, improve_texture_flag,
+             progress=gr.Progress(track_ τότε=True)):
 
-    # Ensure target height/width are multiples of VAE spatial scale factor
-    height_padded_target = ((height - 1) // vae_spatial_scale_factor + 1) * vae_spatial_scale_factor
-    width_padded_target = ((width - 1) // vae_spatial_scale_factor + 1) * vae_spatial_scale_factor
-    
-    # Ensure num_frames is multiple of VAE temporal scale factor + 1 (for causal VAE)
-    # (num_frames - 1) should be multiple of temporal_scale_factor for non-causal parts
-    # For CausalVAE, it's often (N * temporal_factor) + 1 frames.
-    # The inference script uses: num_frames_padded = ((num_frames - 2) // 8 + 1) * 8 + 1
-    # Assuming 8 is the temporal scale factor here for simplicity, adjust if different
-    num_frames_padded_target = ((num_frames - 2) // vae_temporal_scale_factor + 1) * vae_temporal_scale_factor + 1
-
-
-    padding_target = calculate_padding(height, width, height_padded_target, width_padded_target)
-
-
-    if mode == "video-to-video" and video_path:
-        # LTX pipeline's vid2vid uses `media_items` argument for the full video to transform
-        # and `conditioning_items` for specific keyframes if needed.
-        # Here, the Gradio's "video-to-video" seems to imply transforming the input video.
-        input_media_for_vid2vid = load_media_file(
-            media_path=video_path,
-            height=height, # Original height before padding for loading
-            width=width,   # Original width
-            max_frames=min(num_frames_padded_target, frames_to_use if frames_to_use > 0 else num_frames_padded_target),
-            padding=padding_target, # Padding to make it compatible with VAE of target size
-        )
-        # If we also want to strongly condition on the first frame(s) of this video:
-        conditioning_media = load_media_file(
-            media_path=video_path,
-            height=height, width=width,
-            max_frames=min(frames_to_use if frames_to_use > 0 else 1, num_frames_padded_target), # Use specified frames or just the first
-            padding=padding_target,
-            just_crop=True # Crop to aspect ratio, then resize
-        )
-        conditioning_items_list.append(ConditioningItem(media_item=conditioning_media, media_frame_number=0, conditioning_strength=1.0))
-
-    elif mode == "image-to-video" and image_path:
-        conditioning_media = load_image_to_tensor_with_resize_and_crop(
-            image_input=image_path,
-            target_height=height, # Original height
-            target_width=width    # Original width
-        )
-        # Apply padding to the loaded tensor
-        conditioning_media = torch.nn.functional.pad(conditioning_media, padding_target)
-        conditioning_items_list.append(ConditioningItem(media_item=conditioning_media, media_frame_number=0, conditioning_strength=1.0))
+    if randomize_seed:
+        seed_ui = random.randint(0, 2**32 - 1)
+    seed_everething(int(seed_ui))
     
-    # else mode is "text-to-video", no explicit conditioning items unless defined elsewhere
-
-    # --- Get pipeline parameters from YAML ---
-    first_pass_config = PIPELINE_CONFIG_YAML.get("first_pass", {})
-    second_pass_config = PIPELINE_CONFIG_YAML.get("second_pass", {})
-    downscale_factor = PIPELINE_CONFIG_YAML.get("downscale_factor", 2/3)
-
-    # Override steps from Gradio if provided, for the first pass
-    if steps:
-        # The YAML timesteps are specific, so overriding num_inference_steps might not be what we want
-        # If YAML has `timesteps`, `num_inference_steps` is ignored by LTXVideoPipeline.
-        # If YAML does not have `timesteps`, then `num_inference_steps` from Gradio will be used for the first pass.
-        first_pass_config["num_inference_steps"] = steps
-        # For distilled model, the second pass steps are usually very few, defined by its timesteps.
-        # We won't override second_pass_config["num_inference_steps"] from the Gradio `steps`
-        # as it's meant for the primary generation.
-
-    # Determine initial generation dimensions (downscaled)
-    # These are the dimensions for the *first pass* of the multi-scale pipeline
-    initial_gen_height = int(height_padded_target * downscale_factor)
-    initial_gen_width = int(width_padded_target * downscale_factor)
-
-    initial_gen_height, initial_gen_width = round_to_nearest_resolution_acceptable_by_vae(
-        initial_gen_height, initial_gen_width, vae_spatial_scale_factor
-    )
+    actual_height = int(height_ui)
+    actual_width = int(width_ui)
+    actual_num_frames = int(num_frames_ui)
+
+    # Padded dimensions for pipeline
+    height_padded = ((actual_height - 1) // 32 + 1) * 32
+    width_padded = ((actual_width - 1) // 32 + 1) * 32
+    num_frames_padded = ((actual_num_frames - 2) // 8 + 1) * 8 + 1
     
-    shared_pipeline_args = {
+    padding_values = calculate_padding(actual_height, actual_width, height_padded, width_padded)
+
+    call_kwargs = {
         "prompt": prompt,
         "negative_prompt": negative_prompt,
-        "num_frames": num_frames_padded_target, # Always generate padded num_frames
-        "frame_rate": 30, # Example, or get from UI if available
-        "guidance_scale": guidance_scale,
-        "generator": generator,
-        "conditioning_items": conditioning_items_list if conditioning_items_list else None,
-        "skip_layer_strategy": SkipLayerStrategy.AttentionValues, # Default or from YAML
-        "offload_to_cpu": False, # Managed by global DEVICE
-        "is_video": True,
-        "vae_per_channel_normalize": True, # Common default
-        "mixed_precision": (PIPELINE_CONFIG_YAML["precision"] == "bfloat16"),
-        "enhance_prompt": False, # Controlled by Gradio app logic if needed for full LTX script
-        "image_cond_noise_scale": 0.025, # from YAML decode_noise_scale, or make it a param
-        "media_items": input_media_for_vid2vid if mode == "video-to-video" else None,
-        # "decode_timestep" and "decode_noise_scale" are part of first_pass/second_pass or direct call
+        "height": height_padded, # Use padded for pipeline
+        "width": width_padded,   # Use padded for pipeline
+        "num_frames": num_frames_padded, # Use padded for pipeline
+        "frame_rate": 30, 
+        "generator": torch.Generator(device=current_device).manual_seed(int(seed_ui)),
+        "output_type": "pt",
+        "conditioning_items": None,
+        "media_items": None,
+        "decode_timestep": PIPELINE_CONFIG_YAML["decode_timestep"],
+        "decode_noise_scale": PIPELINE_CONFIG_YAML["decode_noise_scale"],
+        "stochastic_sampling": PIPELINE_CONFIG_YAML["stochastic_sampling"],
+        "image_cond_noise_scale": 0.15, # from inference.py defaults
+        "is_video": True, # Assume video output
+        "vae_per_channel_normalize": True, # from inference.py defaults
+        "mixed_precision": (PIPELINE_CONFIG_YAML["precision"] == "mixed_precision"),
+        "offload_to_cpu": False, # For Gradio, keep on device
+        "enhance_prompt": False, # Assuming no UI for this yet, stick to YAML or handle separately
     }
 
-    # --- Generation ---
-    if improve_texture:
-        print("Using LTXMultiScalePipeline for generation...")
-        # Ensure first_pass_config and second_pass_config have necessary overrides
-        # The 'steps' from Gradio applies to the first pass's num_inference_steps if timesteps not set
-        if "timesteps" not in first_pass_config:
-             first_pass_config["num_inference_steps"] = steps
+    stg_mode_str = PIPELINE_CONFIG_YAML.get("stg_mode", "attention_values")
+    if stg_mode_str.lower() in ["stg_av", "attention_values"]:
+        call_kwargs["skip_layer_strategy"] = SkipLayerStrategy.AttentionValues
+    elif stg_mode_str.lower() in ["stg_as", "attention_skip"]:
+        call_kwargs["skip_layer_strategy"] = SkipLayerStrategy.AttentionSkip
+    elif stg_mode_str.lower() in ["stg_r", "residual"]:
+        call_kwargs["skip_layer_strategy"] = SkipLayerStrategy.Residual
+    elif stg_mode_str.lower() in ["stg_t", "transformer_block"]:
+        call_kwargs["skip_layer_strategy"] = SkipLayerStrategy.TransformerBlock
+    else:
+        raise ValueError(f"Invalid stg_mode: {stg_mode_str}")
+
+    if mode == "image-to-video" and input_image_filepath:
+        try:
+            # Ensure the input image is loaded with original H/W for correct aspect ratio handling by the function
+            media_tensor = load_image_to_tensor_with_resize_and_crop(
+                input_image_filepath, actual_height, actual_width
+            )
+            media_tensor = torch.nn.functional.pad(media_tensor, padding_values)
+            call_kwargs["conditioning_items"] = [ConditioningItem(media_tensor.to(current_device), 0, 1.0)]
+        except Exception as e:
+            print(f"Error loading image {input_image_filepath}: {e}")
+            raise gr.Error(f"Could not load image: {e}")
+
+
+    elif mode == "video-to-video" and input_video_filepath:
+        try:
+            call_kwargs["media_items"] = load_media_file(
+                media_path=input_video_filepath,
+                height=actual_height, 
+                width=actual_width,
+                max_frames=int(ui_frames_to_use),
+                padding=padding_values
+            ).to(current_device)
+        except Exception as e:
+            print(f"Error loading video {input_video_filepath}: {e}")
+            raise gr.Error(f"Could not load video: {e}")
+    
+    # Multi-scale or single-scale pipeline call
+    if improve_texture_flag:
+        if not latent_upsampler_instance:
+            raise gr.Error("Spatial upscaler model not loaded, cannot use multi-scale.")
         
-        first_pass_config.setdefault("decode_timestep", PIPELINE_CONFIG_YAML.get("decode_timestep", 0.05))
-        first_pass_config.setdefault("decode_noise_scale", PIPELINE_CONFIG_YAML.get("decode_noise_scale", 0.025))
-        second_pass_config.setdefault("decode_timestep", PIPELINE_CONFIG_YAML.get("decode_timestep", 0.05))
-        second_pass_config.setdefault("decode_noise_scale", PIPELINE_CONFIG_YAML.get("decode_noise_scale", 0.025))
-
-        # The multi_scale_pipe's __call__ expects width and height for the *initial* (downscaled) generation
-        result_frames_tensor = multi_scale_pipe(
-            **shared_pipeline_args,
-            width=initial_gen_width,
-            height=initial_gen_height,
-            downscale_factor=downscale_factor, # This might be used internally by multi_scale_pipe
-            first_pass=first_pass_config,
-            second_pass=second_pass_config,
-            output_type="pt" # Get tensor for further processing
-        ).images
+        multi_scale_pipeline_obj = LTXMultiScalePipeline(pipeline_instance, latent_upsampler_instance)
         
-        # LTXMultiScalePipeline should return images at 2x the initial_gen_width/height
-        # So, result_frames_tensor is at initial_gen_width*2, initial_gen_height*2
-
+        # Prepare pass-specific arguments, overriding with UI inputs where appropriate
+        first_pass_args = PIPELINE_CONFIG_YAML.get("first_pass", {}).copy()
+        first_pass_args["guidance_scale"] = float(ui_guidance_scale)
+        if "timesteps" not in first_pass_args: # Only if YAML doesn't define timesteps
+            first_pass_args["num_inference_steps"] = int(ui_steps)
+
+        second_pass_args = PIPELINE_CONFIG_YAML.get("second_pass", {}).copy()
+        second_pass_args["guidance_scale"] = float(ui_guidance_scale)
+        # num_inference_steps for second pass is typically determined by its YAML timesteps
+
+        multi_scale_call_kwargs = call_kwargs.copy()
+        multi_scale_call_kwargs.update({
+            "downscale_factor": PIPELINE_CONFIG_YAML["downscale_factor"],
+            "first_pass": first_pass_args,
+            "second_pass": second_pass_args,
+        })
+        
+        print(f"Calling multi-scale pipeline with effective height={actual_height}, width={actual_width}")
+        result_images_tensor = multi_scale_pipeline_obj(**multi_scale_call_kwargs).images
     else:
-        print("Using LTXVideoPipeline (first pass) + Manual Upsample + Decode...")
-        # 1. First pass generation at downscaled resolution
-        if "timesteps" not in first_pass_config:
-             first_pass_config["num_inference_steps"] = steps
-
-        first_pass_args = {
-            **shared_pipeline_args,
-            **first_pass_config,
-            "width": initial_gen_width,
-            "height": initial_gen_height,
-            "output_type": "latent"
-        }
-        latents = pipe(**first_pass_args).images # .images here is actually latents
-        print("First pass done!")
-        # 2. Upsample latents manually
-        # Need to handle normalization around latent upsampler if it expects unnormalized latents
-        latents_unnorm = un_normalize_latents(latents, pipe.vae, vae_per_channel_normalize=True)
-        upsampled_latents_unnorm = latent_upsampler(latents_unnorm)
-        upsampled_latents = normalize_latents(upsampled_latents_unnorm, pipe.vae, vae_per_channel_normalize=True)
+        # Single pass call (using base pipeline)
+        single_pass_call_kwargs = call_kwargs.copy()
+        single_pass_call_kwargs["guidance_scale"] = float(ui_guidance_scale)
         
-        # 3. Decode upsampled latents
-        # The upsampler typically doubles the spatial dimensions
-        upscaled_height_for_decode = initial_gen_height * 2
-        upscaled_width_for_decode = initial_gen_width * 2
+        # For single pass, if YAML doesn't have top-level timesteps, use ui_steps
+        # The current YAML is multi-scale focused, so it lacks top-level step control.
+        # We'll assume for a base call, num_inference_steps is directly taken from UI.
+        single_pass_call_kwargs["num_inference_steps"] = int(ui_steps)
+        # Remove pass-specific args if they accidentally slipped in
+        single_pass_call_kwargs.pop("first_pass", None)
+        single_pass_call_kwargs.pop("second_pass", None)
+        single_pass_call_kwargs.pop("downscale_factor", None)
         
-        # Prepare target_shape for VAE decoder
-        # batch_size, channels, num_frames, height, width
-        # Latents are (B, C, F_latent, H_latent, W_latent)
-        # Target shape for vae.decode is pixel space
-        # num_video_frames_final = upsampled_latents.shape[2] * pipe.vae.temporal_downscale_factor
-        # if causal, it might be (F_latent - 1) * factor + 1
-        num_video_frames_final = (upsampled_latents.shape[2] -1) * pipe.vae.temporal_downscale_factor + 1
-
-
-        decode_kwargs = {
-            "target_shape": (
-                upsampled_latents.shape[0], # batch
-                3, # out channels
-                num_video_frames_final,
-                upscaled_height_for_decode,
-                upscaled_width_for_decode
-            )
-        }
-        if pipe.vae.decoder.timestep_conditioning:
-            decode_kwargs["timestep"] = torch.tensor([PIPELINE_CONFIG_YAML.get("decode_timestep", 0.05)] * upsampled_latents.shape[0]).to(DEVICE)
-            # Add noise for decode if specified, similar to LTXVideoPipeline's call
-            noise = torch.randn_like(upsampled_latents)
-            decode_noise_val = PIPELINE_CONFIG_YAML.get("decode_noise_scale", 0.025)
-            upsampled_latents = upsampled_latents * (1 - decode_noise_val) + noise * decode_noise_val
-
-        print("before vae decoding")
-        result_frames_tensor = pipe.vae.decode(upsampled_latents, **decode_kwargs).sample
-        print("after vae decoding?")
-        # result_frames_tensor shape: (B, C, F_video, H_video, W_video)
-
-    # --- Post-processing: Cropping and Converting to PIL ---
-    # Crop to original num_frames (before padding)
-    result_frames_tensor = result_frames_tensor[:, :, :num_frames, :, :]
-
-    # Unpad to target height and width
-    _, _, _, current_h, current_w = result_frames_tensor.shape
-    
-    # Calculate crop needed if current dimensions are larger than padded_target
-    # This happens if multi_scale_pipe output is larger than height_padded_target
-    crop_y_start = (current_h - height_padded_target) // 2
-    crop_x_start = (current_w - width_padded_target) // 2
-    
-    result_frames_tensor = result_frames_tensor[
-        :, :, :, 
-        crop_y_start : crop_y_start + height_padded_target, 
-        crop_x_start : crop_x_start + width_padded_target
-    ]
+        print(f"Calling base pipeline with height={height_padded}, width={width_padded}")
+        result_images_tensor = pipeline_instance(**single_pass_call_kwargs).images
+
+    # Crop to original requested dimensions (num_frames, height, width)
+    # Padding: (pad_left, pad_right, pad_top, pad_bottom)
+    pad_left, pad_right, pad_top, pad_bottom = padding_values
     
-    # Now remove the padding added for VAE compatibility
-    pad_left, pad_right, pad_top, pad_bottom = padding_target
-    unpad_bottom = -pad_bottom if pad_bottom > 0 else result_frames_tensor.shape[3]
-    unpad_right = -pad_right if pad_right > 0 else result_frames_tensor.shape[4]
-
-    result_frames_tensor = result_frames_tensor[
-        :, :, :,
-        pad_top : unpad_bottom,
-        pad_left : unpad_right
+    # Calculate slice indices, ensuring they don't go negative if padding was zero
+    slice_h_end = -pad_bottom if pad_bottom > 0 else None
+    slice_w_end = -pad_right if pad_right > 0 else None
+
+    result_images_tensor = result_images_tensor[
+        :, :, :actual_num_frames, pad_top:slice_h_end, pad_left:slice_w_end
     ]
 
+    # Convert tensor to video file
+    video_np = result_images_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy()
+    video_np = np.clip(video_np * 0.5 + 0.5, 0, 1) # from [-1,1] to [0,1]
+    video_np = (video_np * 255).astype(np.uint8)
 
-    # Convert tensor to list of PIL Images
-    video_pil_list = []
-    # result_frames_tensor shape: (B, C, F, H, W)
-    # We expect B=1 from typical generation
-    video_single_batch = result_frames_tensor[0] # Shape: (C, F, H, W)
-    video_single_batch = (video_single_batch / 2 + 0.5).clamp(0, 1) # Normalize to [0,1]
-    video_single_batch = video_single_batch.permute(1, 2, 3, 0).cpu().float().numpy() # F, H, W, C
+    temp_dir = tempfile.mkdtemp()
+    timestamp = random.randint(10000,99999) # Add timestamp to avoid caching issues
+    output_video_path = os.path.join(temp_dir, f"output_{timestamp}.mp4")
     
-    for frame_idx in range(video_single_batch.shape[0]):
-        frame_np = (video_single_batch[frame_idx] * 255).astype(np.uint8)
-        video_pil_list.append(Image.fromarray(frame_np))
-
-    # Save video
-    output_video_path = "output.mp4" # Gradio handles temp files
-    export_to_video(video_pil_list, output_video_path, fps=24) # Assuming fps from original script
+    try:
+        with imageio.get_writer(output_video_path, fps=call_kwargs["frame_rate"], macro_block_size=1) as video_writer:
+            for frame_idx in range(video_np.shape[0]):
+                progress(frame_idx / video_np.shape[0], desc="Saving video")
+                video_writer.append_data(video_np[frame_idx])
+    except Exception as e:
+        print(f"Error saving video: {e}")
+        # Fallback to saving frame by frame if container issue
+        try:
+            with imageio.get_writer(output_video_path, fps=call_kwargs["frame_rate"], format='FFMPEG', codec='libx264', quality=8, macro_block_size=None) as video_writer:
+                 for frame_idx in range(video_np.shape[0]):
+                    progress(frame_idx / video_np.shape[0], desc="Saving video (fallback)")
+                    video_writer.append_data(video_np[frame_idx])
+        except Exception as e2:
+            print(f"Fallback video saving error: {e2}")
+            raise gr.Error(f"Failed to save video: {e2}")
+
+
+    # Clean up temporary image/video files if they were created by Gradio
+    if isinstance(input_image_filepath, tempfile._TemporaryFileWrapper):
+        input_image_filepath.close()
+        if os.path.exists(input_image_filepath.name):
+            os.remove(input_image_filepath.name)
+    if isinstance(input_video_filepath, tempfile._TemporaryFileWrapper):
+        input_video_filepath.close()
+        if os.path.exists(input_video_filepath.name):
+            os.remove(input_video_filepath.name)
+            
     return output_video_path
 
-
+# --- Gradio UI Definition (from user) ---
 css="""
 #col-container {
     margin: 0 auto;
@@ -384,89 +304,82 @@ css="""
 }
 """
 
-with gr.Blocks(css=css, theme=gr.themes.Ocean()) as demo:
+with gr.Blocks(css=css, theme=gr.themes.Glass()) as demo: # Changed theme for variety
     gr.Markdown("# LTX Video 0.9.7 Distilled (using LTX-Video lib)")
+    gr.Markdown("Generates a short video based on text prompt, image, or existing video.")
     with gr.Row():
         with gr.Column():
             with gr.Group():
                 with gr.Tab("text-to-video") as text_tab:
-                    image_n = gr.Image(label="", visible=False, value=None) # Ensure None for path
-                    video_n = gr.Video(label="", visible=False, value=None) # Ensure None for path
-                    t2v_prompt = gr.Textbox(label="prompt", value="A majestic dragon flying over a medieval castle")
-                    t2v_button = gr.Button("Generate Text-to-Video")
+                    # Hidden inputs for consistent generate() signature
+                    image_n_hidden = gr.Textbox(label="image_n", visible=False, value=None)
+                    video_n_hidden = gr.Textbox(label="video_n", visible=False, value=None)
+                    t2v_prompt = gr.Textbox(label="Prompt", value="A majestic dragon flying over a medieval castle", lines=3)
+                    t2v_button = gr.Button("Generate Text-to-Video", variant="primary")
                 with gr.Tab("image-to-video") as image_tab:
-                    video_i = gr.Video(label="", visible=False, value=None)
-                    image_i2v = gr.Image(label="input image", type="filepath")
-                    i2v_prompt = gr.Textbox(label="prompt", value="The creature from the image starts to move")
-                    i2v_button = gr.Button("Generate Image-to-Video")
+                    video_i_hidden = gr.Textbox(label="video_i", visible=False, value=None)
+                    image_i2v = gr.Image(label="Input Image", type="filepath", sources=["upload", "webcam"])
+                    i2v_prompt = gr.Textbox(label="Prompt", value="The creature from the image starts to move", lines=3)
+                    i2v_button = gr.Button("Generate Image-to-Video", variant="primary")
                 with gr.Tab("video-to-video") as video_tab:
-                    image_v = gr.Image(label="", visible=False, value=None)
-                    video_v2v = gr.Video(label="input video")
-                    frames_to_use = gr.Number(label="num frames to use",info="first # of frames to use from the input video for conditioning/transformation", value=9)
-                    v2v_prompt = gr.Textbox(label="prompt", value="Change the style to cinematic anime")
-                    v2v_button = gr.Button("Generate Video-to-Video")
+                    image_v_hidden = gr.Textbox(label="image_v", visible=False, value=None)
+                    video_v2v = gr.Video(label="Input Video", sources=["upload", "webcam"])
+                    frames_to_use = gr.Slider(label="Frames to use from input video", minimum=9, maximum=MAX_NUM_FRAMES, value=9, step=8, info="Number of initial frames to use for conditioning/transformation. Must be N*8+1.")
+                    v2v_prompt = gr.Textbox(label="Prompt", value="Change the style to cinematic anime", lines=3)
+                    v2v_button = gr.Button("Generate Video-to-Video", variant="primary")
 
-                improve_texture = gr.Checkbox(label="improve texture (multi-scale)", value=True, info="Uses a two-pass generation for better quality, but is slower.")
+            improve_texture = gr.Checkbox(label="Improve Texture (multi-scale)", value=True, info="Uses a two-pass generation for better quality, but is slower. Recommended for final output.")
 
         with gr.Column():
-            output = gr.Video(interactive=False)
+            output_video = gr.Video(label="Generated Video", interactive=False)
+            gr.Markdown("Note: Generation can take a few minutes depending on settings and hardware.")
 
     with gr.Accordion("Advanced settings", open=False):
-        negative_prompt_input = gr.Textbox(label="negative prompt", value="worst quality, inconsistent motion, blurry, jittery, distorted")
+        negative_prompt_input = gr.Textbox(label="Negative Prompt", value="worst quality, inconsistent motion, blurry, jittery, distorted", lines=2)
         with gr.Row():
-            seed_input = gr.Number(label="seed", value=42, precision=0)
-            randomize_seed_input = gr.Checkbox(label="randomize seed", value=False)
+            seed_input = gr.Number(label="Seed", value=42, precision=0, minimum=0, maximum=2**32-1)
+            randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=False)
         with gr.Row():
-            guidance_scale_input = gr.Slider(label="guidance scale", minimum=0, maximum=10, value=1.0, step=0.1, info="For distilled models, CFG is often 1.0 (disabled) or very low.") # Distilled model might not need high CFG
-            steps_input = gr.Slider(label="Steps (for first pass if multi-scale)", minimum=1, maximum=30, value=PIPELINE_CONFIG_YAML.get("first_pass", {}).get("timesteps", [1]*8).__len__(), step=1, info="Number of inference steps. If YAML defines timesteps, this is ignored for that pass.") # Default to length of first_pass timesteps
-            num_frames_input = gr.Slider(label="# frames", minimum=9, maximum=121, value=25, step=8, info="Should be N*8+1, e.g., 9, 17, 25...") # Adjusted for LTX structure
+            # For distilled models, CFG is often 1.0 (disabled) or very low.
+            guidance_scale_input = gr.Slider(label="Guidance Scale (CFG)", minimum=1.0, maximum=10.0, value=PIPELINE_CONFIG_YAML.get("first_pass", {}).get("guidance_scale", 1.0), step=0.1, info="Controls how much the prompt influences the output. Higher values = stronger influence.")
+            # Default to length of first_pass timesteps, if available
+            default_steps = len(PIPELINE_CONFIG_YAML.get("first_pass", {}).get("timesteps", [1]*7)) # Fallback to 7 if not defined
+            steps_input = gr.Slider(label="Inference Steps (for first pass if multi-scale)", minimum=1, maximum=30, value=default_steps, step=1, info="Number of denoising steps. More steps can improve quality but increase time. If YAML defines 'timesteps' for a pass, this UI value is ignored for that pass.")
         with gr.Row():
-            height_input = gr.Slider(label="height", value=512, step=8, minimum=256, maximum=MAX_IMAGE_SIZE) # Step by VAE factor
-            width_input = gr.Slider(label="width", value=704, step=8, minimum=256, maximum=MAX_IMAGE_SIZE) # Step by VAE factor
-
-    t2v_button.click(fn=generate,
-                      inputs=[t2v_prompt,
-                              negative_prompt_input,
-                              image_n, # Pass None for image
-                              video_n, # Pass None for video
-                              height_input,
-                              width_input,
-                              gr.State("text-to-video"),
-                              steps_input,
-                              num_frames_input,
-                              gr.State(0), # frames_to_use not relevant for t2v
-                              seed_input,
-                              randomize_seed_input, guidance_scale_input, improve_texture],
-                      outputs=[output])
-
-    i2v_button.click(fn=generate,
-                      inputs=[i2v_prompt,
-                              negative_prompt_input,
-                              image_i2v,
-                              video_i, # Pass None for video
-                              height_input,
-                              width_input,
-                              gr.State("image-to-video"),
-                              steps_input,
-                              num_frames_input,
-                              gr.State(0), # frames_to_use not relevant for i2v initial frame
-                              seed_input,
-                              randomize_seed_input, guidance_scale_input, improve_texture],
-                      outputs=[output])
-
-    v2v_button.click(fn=generate,
-                      inputs=[v2v_prompt,
-                              negative_prompt_input,
-                              image_v, # Pass None for image
-                              video_v2v,
-                              height_input,
-                              width_input,
-                              gr.State("video-to-video"),
-                              steps_input,
-                              num_frames_input,
-                              frames_to_use,
-                              seed_input,
-                              randomize_seed_input, guidance_scale_input, improve_texture],
-                      outputs=[output])
-
-demo.launch()
+            num_frames_input = gr.Slider(label="Number of Frames to Generate", minimum=9, maximum=MAX_NUM_FRAMES, value=25, step=8, info="Total frames in the output video. Should be N*8+1 (e.g., 9, 17, 25...).")
+        with gr.Row():
+            height_input = gr.Slider(label="Height", value=512, step=32, minimum=256, maximum=MAX_IMAGE_SIZE, info="Must be divisible by 32.")
+            width_input = gr.Slider(label="Width", value=704, step=32, minimum=256, maximum=MAX_IMAGE_SIZE, info="Must be divisible by 32.")
+    
+    # Define click actions
+    # Note: gr.State passes the current value of the component without creating a UI element for it.
+    # We use hidden Textbox inputs for image_n, video_n etc. and pass their `value` (which is None)
+    # to ensure the `generate` function always receives these arguments.
+    
+    t2v_inputs = [t2v_prompt, negative_prompt_input, image_n_hidden, video_n_hidden,
+                  height_input, width_input, gr.State("text-to-video"),
+                  steps_input, num_frames_input, gr.State(0), # frames_to_use not relevant for t2v
+                  seed_input, randomize_seed_input, guidance_scale_input, improve_texture]
+    
+    i2v_inputs = [i2v_prompt, negative_prompt_input, image_i2v, video_i_hidden,
+                  height_input, width_input, gr.State("image-to-video"),
+                  steps_input, num_frames_input, gr.State(0), # frames_to_use not relevant for i2v initial frame
+                  seed_input, randomize_seed_input, guidance_scale_input, improve_texture]
+
+    v2v_inputs = [v2v_prompt, negative_prompt_input, image_v_hidden, video_v2v,
+                  height_input, width_input, gr.State("video-to-video"),
+                  steps_input, num_frames_input, frames_to_use,
+                  seed_input, randomize_seed_input, guidance_scale_input, improve_texture]
+
+    t2v_button.click(fn=generate, inputs=t2v_inputs, outputs=[output_video])
+    i2v_button.click(fn=generate, inputs=i2v_inputs, outputs=[output_video])
+    v2v_button.click(fn=generate, inputs=v2v_inputs, outputs=[output_video])
+
+if __name__ == "__main__":
+    # Clean up old model directory if it exists from previous runs
+    if os.path.exists(models_dir) and os.path.isdir(models_dir):
+        print(f"Cleaning up old model directory: {models_dir}")
+        # shutil.rmtree(models_dir) # Optional: uncomment to force re-download on every run
+    Path(models_dir).mkdir(parents=True, exist_ok=True)
+    
+    demo.queue().launch(debug=True, share=False)