Update app.py

app.py

@@ -1,137 +1,379 @@
 import gradio as gr
-import spaces 
+import spaces
 import torch
-# from pipeline_ltx_condition import LTXVideoCondition, LTXConditionPipeline
-# from diffusers import LTXLatentUpsamplePipeline
-from diffusers import LTXConditionPipeline, LTXLatentUpsamplePipeline
-from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition
-from diffusers.utils import export_to_video, load_video
 import numpy as np
+import os
+import yaml
+import random
+from PIL import Image
+import imageio # For export_to_video and reading video frames
+from pathlib import Path
+from huggingface_hub import hf_hub_download
 
+# --- 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
+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,
+    seed_everething,
+    calculate_padding,
+)
+from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
+from ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler
+# --- End LTX-Video Imports ---
 
-pipe = LTXConditionPipeline.from_pretrained("linoyts/LTX-Video-0.9.7-distilled-diffusers", torch_dtype=torch.bfloat16)
-pipe_upsample = LTXLatentUpsamplePipeline.from_pretrained("a-r-r-o-w/LTX-Video-0.9.7-Latent-Spatial-Upsampler-diffusers", vae=pipe.vae, torch_dtype=torch.bfloat16)
-pipe.to("cuda")
-pipe_upsample.to("cuda")
-pipe.vae.enable_tiling()
+# --- 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
+# ---
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 2048
+# --- 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 = "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)
 
-def round_to_nearest_resolution_acceptable_by_vae(height, width):
-    print("before rounding",height, width)
-    height = height - (height % pipe.vae_spatial_compression_ratio)
-    width = width - (width % pipe.vae_spatial_compression_ratio)
-    print("after rounding",height, width)
-    return height, width
+# 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
 
-def change_mode_to_text():
-  return gr.update(value="text-to-video")
+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)
 
-def change_mode_to_image():
-  return gr.update(value="image-to-video")
+    # 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
+)
+
+# Create Latent Upsampler
+latent_upsampler = create_latent_upsampler(
+    latent_upsampler_model_path=spatial_upsampler_path,
+    device=DEVICE
+)
+latent_upsampler = latent_upsampler.to(torch.bfloat16 if PIPELINE_CONFIG_YAML["precision"] == "bfloat16" else torch.float32)
+
+
+# Multi-scale pipeline (wrapper)
+multi_scale_pipe = LTXMultiScalePipeline(
+    video_pipeline=pipe,
+    latent_upsampler=latent_upsampler
+)
+# --- 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
 
-def change_mode_to_video():
-  return gr.update(value="video-to-video")
-    
 @spaces.GPU
 def generate(prompt,
              negative_prompt,
-             image,
-             video,
+             image_path, # Gradio gives filepath for Image component
+             video_path, # Gradio gives filepath for Video component
              height,
              width,
              mode,
-             steps,
+             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)):
-    
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-        
-    # Part 1. Generate video at smaller resolution
-    # Text-only conditioning is also supported without the need to pass `conditions`
-    expected_height, expected_width = height, width 
-    downscale_factor = 2 / 3
-    downscaled_height, downscaled_width = int(expected_height * downscale_factor), int(expected_width * downscale_factor)
-    downscaled_height, downscaled_width = round_to_nearest_resolution_acceptable_by_vae(downscaled_height, downscaled_width)
-
-    print(mode)
-    if mode == "text-to-video" and (video is not None):
-        video = load_video(video)[:frames_to_use]
-        condition = True
-    elif mode == "image-to-video" and (image is not None):
-        print("WTFFFFFF 1")
-        video = [image]
-        condition = True
-    else:
-       condition=False
+    seed_everething(seed)
+    
+    generator = torch.Generator(device=DEVICE).manual_seed(seed)
 
-    if condition:
-        print("WTFFFFFF 2")
-        condition1 = LTXVideoCondition(video=video, frame_index=0)
-    else:
-        condition1 = None
+    # --- 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
+
+    # 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
     
-    latents = pipe(
-        conditions=condition1,
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        width=downscaled_width,
-        height=downscaled_height,
-        num_frames=num_frames,
-        num_inference_steps=steps,
-        decode_timestep = 0.05,
-        decode_noise_scale = 0.025,
-        guidance_scale=guidance_scale,
-        generator=torch.Generator(device="cuda").manual_seed(seed),
-        output_type="latent",
-    ).frames
-   
-        
-    # Part 2. Upscale generated video using latent upsampler with fewer inference steps
-    # The available latent upsampler upscales the height/width by 2x
+    # 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))
+    
+    # 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
+    )
+    
+    shared_pipeline_args = {
+        "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
+    }
+
+    # --- Generation ---
     if improve_texture:
-        upscaled_height, upscaled_width = downscaled_height * 2, downscaled_width * 2
-        upscaled_latents = pipe_upsample(
-            latents=latents,
-            output_type="latent"
-        ).frames
+        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
+        
+        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
         
-        # Part 3. Denoise the upscaled video with few steps to improve texture (optional, but recommended)  
-        video = pipe(
-            conditions=condition1,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            width=upscaled_width,
-            height=upscaled_height,
-            num_frames=num_frames,
-            guidance_scale=guidance_scale,
-            denoise_strength=0.6,  # Effectively, 0.6 * 3 inference steps
-            num_inference_steps=3,
-            latents=upscaled_latents,
-            decode_timestep=0.05,
-            image_cond_noise_scale=0.025,
-            generator=torch.Generator().manual_seed(seed),
-            output_type="pil",
-        ).frames[0]
+        # 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
+
     else:
-        upscaled_height, upscaled_width = downscaled_height * 2, downscaled_width * 2
-        video = pipe_upsample(
-            latents=latents,
-            # output_type="latent"
-        ).frames[0]
+        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
+
+        # 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)
+        
+        # 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
+        
+        # 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
+
+
+        result_frames_tensor = pipe.vae.decode(upsampled_latents, **decode_kwargs).sample
+        # 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
+    ]
     
-    # Part 4. Downscale the video to the expected resolution
-    video = [frame.resize((expected_width, expected_height)) for frame in video]
-    export_to_video(video, "output.mp4", fps=24)
-    return "output.mp4"
+    # 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
+    ]
+
+
+    # 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().numpy() # F, H, W, C
+    
+    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
+    return output_video_path
 
 
 css="""
@@ -141,72 +383,89 @@ css="""
 }
 """
 
-js_func = """
-function refresh() {
-    const url = new URL(window.location);
+with gr.Blocks(css=css, theme=gr.themes.Ocean()) as demo:
+    gr.Markdown("# LTX Video 0.9.7 Distilled (using LTX-Video lib)")
+    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")
+                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")
+                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", type="filepath")
+                    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")
 
-    if (url.searchParams.get('__theme') !== 'dark') {
-        url.searchParams.set('__theme', 'dark');
-        window.location.href = url.href;
-    }
-}
-"""
+                improve_texture = gr.Checkbox(label="improve texture (multi-scale)", value=True, info="Uses a two-pass generation for better quality, but is slower.")
 
-with gr.Blocks(css=css, theme=gr.themes.Ocean()) as demo:
+        with gr.Column():
+            output = gr.Video(interactive=False)
 
-  gr.Markdown("# LTX Video 0.9.7 Distilled")
-  mode = gr.State(value="text-to-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)
-        with gr.Tab("image-to-video") as image_tab:
-          image = gr.Image(label="input image")
-        with gr.Tab("video-to-video") as video_tab:
-          video = 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", value=1)
-        prompt = gr.Textbox(label="prompt")
-        improve_texture = gr.Checkbox(label="improve texture", value=False, info="slows down generation")
-      run_button = gr.Button()
-    with gr.Column():
-      output = gr.Video(interactive=False)
-      
-
-  with gr.Accordion("Advanced settings", open=False):
-     negative_prompt = gr.Textbox(label="negative prompt", value="worst quality, inconsistent motion, blurry, jittery, distorted", visible=False)  
-     with gr.Row():
-      seed = gr.Number(label="seed", value=0, precision=0)
-      randomize_seed = gr.Checkbox(label="randomize seed")
-     with gr.Row():
-      guidance_scale= gr.Slider(label="guidance scale", minimum=0, maximum=10, value=3, step=1)
-      steps = gr.Slider(label="Steps", minimum=1, maximum=30, value=8, step=1)
-      num_frames = gr.Slider(label="# frames", minimum=1, maximum=161, value=96, step=1)
-     with gr.Row():
-       height = gr.Slider(label="height", value=512, step=1, maximum=2048)
-       width = gr.Slider(label="width", value=704, step=1, maximum=2048)
-    
+    with gr.Accordion("Advanced settings", open=False):
+        negative_prompt_input = gr.Textbox(label="negative prompt", value="worst quality, inconsistent motion, blurry, jittery, distorted")
+        with gr.Row():
+            seed_input = gr.Number(label="seed", value=42, precision=0)
+            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
+        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
 
-  text_tab.select(fn=change_mode_to_text, inputs=[], outputs=[mode])
-  image_tab.select(fn=change_mode_to_image, inputs=[], outputs=[mode])
-  video_tab.select(fn=change_mode_to_video, inputs=[], outputs=[mode])
-    
-  run_button.click(fn=generate, 
-                   inputs=[prompt,
-             negative_prompt,
-             image,
-             video,
-             height,
-             width,
-             mode,
-             steps,
-             num_frames,
-             frames_to_use,
-             seed,
-             randomize_seed,guidance_scale, improve_texture], 
-                   outputs=[output])
+    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()
+demo.launch()