import spaces import sys import os import subprocess @spaces.GPU(duration=120) def install_flashattn(): subprocess.run(['sh', './flashattn.sh']) install_flashattn() # --- PyTorch Environment Setup --- os.environ['PYTORCH_NVML_BASED_CUDA_CHECK'] = '1' os.environ['TORCH_LINALG_PREFER_CUSOLVER'] = '1' os.environ['PYTORCH_ALLOC_CONF'] = 'expandable_segments:True,pinned_use_background_threads:True' os.environ["SAFETENSORS_FAST_GPU"] = "1" os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = '1' import torch # Set precision settings for reproducibility and performance torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False torch.backends.cudnn.deterministic = False torch.backends.cudnn.benchmark = False # Set to True for potential speedup if input sizes are static, False for dynamic torch.backends.cuda.preferred_blas_library="cublas" torch.backends.cuda.preferred_linalg_library="cusolver" torch.set_float32_matmul_precision("highest") FTP_HOST = os.getenv("FTP_HOST") FTP_USER = os.getenv("FTP_USER") FTP_PASS = os.getenv("FTP_PASS") FTP_DIR = os.getenv("FTP_DIR") import cv2 import gc import paramiko #from image_gen_aux import UpscaleWithModel # REMOVED: UpscaleWithModel import import numpy as np import gradio as gr import random import yaml from pathlib import Path import imageio import tempfile from PIL import Image from huggingface_hub import hf_hub_download import shutil from diffusers import AutoencoderKL from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem, LTXMultiScalePipeline from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy from inference import ( create_ltx_video_pipeline, create_latent_upsampler, load_image_to_tensor_with_resize_and_crop, seed_everething, get_device, calculate_padding, load_media_file ) from moviepy.editor import VideoFileClip, concatenate_videoclips from typing import Any, Dict, Optional, Tuple # Imports for TeaCache from ltx_video.models.transformers.transformer3d import Transformer3DModel, Transformer3DModelOutput from diffusers.utils import logging import re logger = logging.get_logger(__name__) # --- Start TeaCache Integration --- # 1. Store the original, unbound forward method from the class definition original_transformer_forward = Transformer3DModel.forward # 2. Define our new, robust wrapper function def teacache_wrapper_forward(self, hidden_states: torch.Tensor, **kwargs): if not hasattr(self, "enable_teacache") or not self.enable_teacache: # Call the original method if TeaCache is disabled return original_transformer_forward(self, hidden_states=hidden_states, **kwargs) # Determine if we should calculate or skip should_calc = True if self.cnt > 0 and self.cnt < self.num_steps - 1: if (hasattr(self, "previous_hidden_states") and self.previous_hidden_states is not None and self.previous_hidden_states.shape == hidden_states.shape): rel_l1_dist = ((hidden_states - self.previous_hidden_states).abs().mean() / self.previous_hidden_states.abs().mean()).cpu().item() self.accumulated_rel_l1_distance += rel_l1_dist if self.accumulated_rel_l1_distance < self.rel_l1_thresh: should_calc = False else: self.accumulated_rel_l1_distance = 0 else: # Force calculation if shapes mismatch or it's the first time in a new pass self.accumulated_rel_l1_distance = 0 self.cnt += 1 if not should_calc and hasattr(self, "previous_residual") and self.previous_residual is not None and self.previous_residual.shape == hidden_states.shape: # SKIP: Use the cached result # The pipeline expects a Transformer3DModelOutput object. return Transformer3DModelOutput(sample=self.previous_residual + hidden_states) else: # COMPUTE: Call the original, stored method, passing 'self' explicitly self.previous_hidden_states = hidden_states.clone() output = original_transformer_forward(self, hidden_states=hidden_states, **kwargs) # Handle both tuple and object return types from the original function if isinstance(output, tuple): output_tensor = output[0] else: output_tensor = output.sample self.previous_residual = output_tensor - hidden_states return output # 3. Apply the patch Transformer3DModel.forward = teacache_wrapper_forward print("✅ Transformer3DModel patched with robust TeaCache Wrapper.") # --- End TeaCache Integration --- MAX_SEED = np.iinfo(np.int32).max # REMOVED: Upscaler pipeline initialization #upscaler = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cuda:0")) # REMOVED: SDXL Image-to-Image enhancer pipeline initialization # print("Loading SDXL Image-to-Image pipeline...") # enhancer_pipeline = StableDiffusionXLImg2ImgPipeline.from_pretrained( # "ford442/stable-diffusion-xl-refiner-1.0-bf16", # use_safetensors=True, # requires_aesthetics_score=True, # ) # enhancer_pipeline.vae.set_default_attn_processor() # enhancer_pipeline.to("cpu") # print("SDXL Image-to-Image pipeline loaded successfully.") config_file_path = "configs/ltxv-13b-0.9.8-distilled.yaml" with open(config_file_path, "r") as file: PIPELINE_CONFIG_YAML = yaml.safe_load(file) LTX_REPO = "Lightricks/LTX-Video" MAX_IMAGE_SIZE = PIPELINE_CONFIG_YAML.get("max_resolution", 1280) MAX_NUM_FRAMES = 900 models_dir = "downloaded_models_gradio_cpu_init" Path(models_dir).mkdir(parents=True, exist_ok=True) pipeline_instance = None latent_upsampler_instance = None temporal_upsampler_instance = None print("Downloading models (if not present)...") distilled_model_actual_path = hf_hub_download(repo_id=LTX_REPO, filename=PIPELINE_CONFIG_YAML["checkpoint_path"], local_dir=models_dir, local_dir_use_symlinks=False) PIPELINE_CONFIG_YAML["checkpoint_path"] = distilled_model_actual_path SPATIAL_UPSCALER_FILENAME = PIPELINE_CONFIG_YAML["spatial_upscaler_model_path"] spatial_upscaler_actual_path = hf_hub_download(repo_id=LTX_REPO, filename=SPATIAL_UPSCALER_FILENAME, local_dir=models_dir, local_dir_use_symlinks=False) PIPELINE_CONFIG_YAML["spatial_upscaler_model_path"] = spatial_upscaler_actual_path # --- Download Temporal Upscaler (Corrected Location) --- TEMPORAL_UPSCALER_FILENAME = "ltxv-temporal-upscaler-0.9.8.safetensors" try: print(f"Downloading temporal upscaler model: {TEMPORAL_UPSCALER_FILENAME}") temporal_upscaler_actual_path = hf_hub_download( repo_id=LTX_REPO, filename=TEMPORAL_UPSCALER_FILENAME, local_dir=models_dir, local_dir_use_symlinks=False ) PIPELINE_CONFIG_YAML["temporal_upscaler_model_path"] = temporal_upscaler_actual_path except Exception as e: print(f"Warning: Could not download temporal upscaler ({TEMPORAL_UPSCALER_FILENAME}). Proceeding without it. Error: {e}") PIPELINE_CONFIG_YAML["temporal_upscaler_model_path"] = None # --- Create Pipeline Instances --- print("Creating LTX Video pipeline on CPU...") 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="cpu", enhance_prompt=False, 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"] ) if PIPELINE_CONFIG_YAML.get("spatial_upscaler_model_path"): print("Creating latent upsampler on CPU...") latent_upsampler_instance = create_latent_upsampler(PIPELINE_CONFIG_YAML["spatial_upscaler_model_path"], device="cpu") if PIPELINE_CONFIG_YAML.get("temporal_upscaler_model_path"): print("Creating temporal upsampler on CPU...") temporal_upsampler_instance = create_latent_upsampler(PIPELINE_CONFIG_YAML["temporal_upscaler_model_path"], device="cpu") target_inference_device = "cuda" print(f"Target inference device: {target_inference_device}") pipeline_instance.to(target_inference_device) if latent_upsampler_instance: latent_upsampler_instance.to(target_inference_device) if temporal_upsampler_instance: temporal_upsampler_instance.to(target_inference_device) dynamic_shapes = { "hidden_states": { 0: torch.export.Dim("batch_size"), 1: torch.export.Dim("num_frames"), # <-- CRUCIAL for video 2: torch.export.Dim("sequence_length"), }, "encoder_hidden_states": { 0: torch.export.Dim("batch_size"), 1: torch.export.Dim("text_sequence_length"), }, # ... add other inputs as needed, just like we did before } def get_duration(*args, **kwargs): duration_ui = kwargs.get('duration_ui', 5.0) if duration_ui > 7.0: return 110 if duration_ui > 5.0: return 100 if duration_ui > 4.0: return 90 if duration_ui > 3.0: return 70 if duration_ui > 2.0: return 60 if duration_ui > 1.5: return 50 if duration_ui > 1.0: return 45 if duration_ui > 0.5: return 30 return 90 def upload_to_sftp(local_filepath): if not all([FTP_HOST, FTP_USER, FTP_PASS, FTP_DIR]): print("SFTP credentials not set. Skipping upload.") return try: transport = paramiko.Transport((FTP_HOST, 22)) transport.connect(username=FTP_USER, password=FTP_PASS) sftp = paramiko.SFTPClient.from_transport(transport) remote_filename = os.path.basename(local_filepath) remote_filepath = os.path.join(FTP_DIR, remote_filename) print(f"Uploading {local_filepath} to {remote_filepath}...") sftp.put(local_filepath, remote_filepath) print("Upload successful.") sftp.close() transport.close() except Exception as e: print(f"SFTP upload failed: {e}") gr.Warning(f"SFTP upload failed: {e}") def calculate_new_dimensions(orig_w, orig_h): if orig_w == 0 or orig_h == 0: return int(1024), int(1024) if orig_w >= orig_h: new_h, new_w = 1024, round((1024 * (orig_w / orig_h)) / 32) * 32 else: new_w, new_h = 1024, round((1024 * (orig_h / orig_w)) / 32) * 32 return int(max(256, min(new_h, MAX_IMAGE_SIZE))), int(max(256, min(new_w, MAX_IMAGE_SIZE))) # REMOVED: superres_image function # REMOVED: enhance_frame function # MODIFIED: Removed calls to superres_image and enhance_frame def use_last_frame_as_input(video_filepath): if not video_filepath or not os.path.exists(video_filepath): gr.Warning("No video clip available.") return None, gr.update() cap = None try: cap = cv2.VideoCapture(video_filepath) frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count - 1) ret, frame = cap.read() if not ret: raise ValueError("Failed to read frame.") pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) print("Displaying last frame and switching tab...") yield pil_image, gr.update(selected="i2v_tab") except Exception as e: gr.Error(f"Failed to extract frame: {e}") return None, gr.update() finally: if cap: cap.release() def stitch_videos(clips_list): if not clips_list or len(clips_list) < 2: raise gr.Error("You need at least two clips to stitch them together!") print(f"Stitching {len(clips_list)} clips...") try: video_clips = [VideoFileClip(clip_path) for clip_path in clips_list] final_clip = concatenate_videoclips(video_clips, method="compose") final_output_path = os.path.join(tempfile.mkdtemp(), f"stitched_video_{random.randint(10000,99999)}.mp4") high_quality_params = ['-crf', '0', '-preset', 'veryslow'] final_clip.write_videofile( final_output_path, codec="libx264", audio=False, threads=4, ffmpeg_params=high_quality_params # <-- USE PARAMS HERE ) for clip in video_clips: clip.close() return final_output_path except Exception as e: raise gr.Error(f"Failed to stitch videos: {e}") def clear_clips(): # state, counter, video1, video2, toggle_visible, tensor_state, randomize_seed return [], "Clips created: 0", None, None, gr.update(visible=False, value=False), None, gr.update(value=True) @spaces.GPU(duration=get_duration) def generate(prompt, negative_prompt, clips_list, input_image_filepath, input_video_filepath, last_frame_tensor_from_state, height_ui, width_ui, mode, duration_ui, ui_frames_to_use, seed_ui, randomize_seed, ui_guidance_scale, improve_texture_flag, use_temporal_upscaler_flag, use_last_tensor_flag, num_steps, fps, enable_teacache, teacache_threshold, text_encoder_max_tokens_ui, image_cond_noise_scale_ui, # <--- NEW PARAMETER ADDED HERE progress=gr.Progress(track_tqdm=True)): # --- FIX: TeaCache + Multi-Scale Incompatibility --- # Force disable TeaCache if multi-scale is on, as the state will not be # reset between the first and second pass, corrupting the output. if improve_texture_flag and enable_teacache: gr.Warning("TeaCache is incompatible with Multi-Scale mode. Disabling TeaCache for this run.") enable_teacache = False # Configure TeaCache state on the transformer instance for this run try: pipeline_instance.transformer.enable_teacache = enable_teacache if enable_teacache: print(f"✅ TeaCache is ENABLED with threshold: {teacache_threshold}") pipeline_instance.transformer.rel_l1_thresh = teacache_threshold else: print("❌ TeaCache is DISABLED.") except AttributeError: print("⚠️ Could not configure TeaCache on transformer.") # Set highest precision for the main generation pipeline torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False torch.backends.cudnn.allow_tf32 = False torch.backends.cudnn.deterministic = False torch.set_float32_matmul_precision("highest") torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() if mode not in ["text-to-video", "image-to-video", "video-to-video"]: raise gr.Error(f"Invalid mode: {mode}.") # Input validation if (mode == "image-to-video" and not input_image_filepath and not (use_last_tensor_flag and last_frame_tensor_from_state is not None)): raise gr.Error("Input image is required for image-to-video mode (or 'Use Last Frame' must be checked).") elif mode == "video-to-video" and not input_video_filepath: raise gr.Error("input_video_filepath is required for video-to-video mode") if randomize_seed: seed_ui = random.randint(0, 2**32 - 1) seed_everething(int(seed_ui)) actual_num_frames = max(9, min(MAX_NUM_FRAMES, int(round((max(1, round(duration_ui * fps)) - 1.0) / 8.0) * 8 + 1))) actual_height, actual_width = int(height_ui), int(width_ui) height_padded, width_padded = ((actual_height - 1) // 32 + 1) * 32, ((actual_width - 1) // 32 + 1) * 32 padding_values = calculate_padding(actual_height, actual_width, height_padded, width_padded) num_frames_padded = max(9, ((actual_num_frames - 2) // 8 + 1) * 8 + 1) call_kwargs = { "prompt": prompt, "negative_prompt": negative_prompt, "height": height_padded, "width": width_padded, "num_frames": num_frames_padded, "num_inference_steps": num_steps, "frame_rate": int(fps), "generator": torch.Generator(device=target_inference_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": image_cond_noise_scale_ui, "is_video": True, "vae_per_channel_normalize": True, "mixed_precision": (PIPELINE_CONFIG_YAML["precision"] == "mixed_precision"), "text_encoder_max_tokens": text_encoder_max_tokens_ui, # <-- PASSED HERE "offload_to_cpu": False, "enhance_prompt": False } stg_mode_str = PIPELINE_CONFIG_YAML.get("stg_mode", "attention_values").lower() stg_map = { "stg_av": SkipLayerStrategy.AttentionValues, "attention_values": SkipLayerStrategy.AttentionValues, "stg_as": SkipLayerStrategy.AttentionSkip, "attention_skip": SkipLayerStrategy.AttentionSkip, "stg_r": SkipLayerStrategy.Residual, "residual": SkipLayerStrategy.Residual, "stg_t": SkipLayerStrategy.TransformerBlock, "transformer_block": SkipLayerStrategy.TransformerBlock } call_kwargs["skip_layer_strategy"] = stg_map.get(stg_mode_str, SkipLayerStrategy.AttentionValues) # --- INPUT LOGIC: Priority = Direct Tensor > Image File > Video File --- if use_last_tensor_flag and last_frame_tensor_from_state is not None: print("Using last frame tensor as input (Direct Tensor).") media_tensor = last_frame_tensor_from_state.to(target_inference_device) b, c, n, h, w = media_tensor.shape # CRITICAL FIX 2: Better interpolation for Latents # Latents are compressed data. 'bilinear' can smear features. # 'nearest-exact' preserves distinct latent features better if resizing is strictly necessary. # However, avoiding resizing entirely is best. if h != actual_height or w != actual_width: media_tensor_4d = media_tensor.view(b * n, c, h, w) resized_tensor_4d = torch.nn.functional.interpolate( media_tensor_4d, size=(actual_height, actual_width), mode='bilinear', # Changed from bilinear align_corners=False, antialias=False ) media_tensor_5d = resized_tensor_4d.view(b, c, n, actual_height, actual_width) else: media_tensor_5d = media_tensor # Pad and set call_kwargs["conditioning_items"] = [ConditioningItem( torch.nn.functional.pad(media_tensor, padding_values).to(target_inference_device), 0, 1.0 )] elif mode == "image-to-video": print("Using image file as input.") # Standard noise scale is fine for I2V from file call_kwargs["image_cond_noise_scale"] = 0.05 media_tensor = load_image_to_tensor_with_resize_and_crop(input_image_filepath, actual_height, actual_width) call_kwargs["conditioning_items"] = [ConditioningItem(torch.nn.functional.pad(media_tensor, padding_values).to(target_inference_device), 0, 1.0)] elif mode == "video-to-video": print("Using video file as input.") 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(target_inference_device) # --- END INPUT LOGIC --- if improve_texture_flag and latent_upsampler_instance: # --- Temporal Upscaler Logic --- temporal_upsampler_to_use = None if use_temporal_upscaler_flag and temporal_upsampler_instance: print("Temporal upscaler is ENABLED.") temporal_upsampler_to_use = temporal_upsampler_instance elif use_temporal_upscaler_flag: print("Warning: Temporal upscaler toggled ON, but model 'temporal_upsampler_instance' was not found.") else: print("Temporal upscaler is DISABLED by user toggle.") # --- FIX: Pass temporal_upsampler_to_use as the 3RD POSITIONAL ARG --- multi_scale_pipeline = LTXMultiScalePipeline( pipeline_instance, latent_upsampler_instance ) pass_args = {"guidance_scale": float(ui_guidance_scale)} multi_scale_kwargs = { **call_kwargs, "downscale_factor": PIPELINE_CONFIG_YAML["downscale_factor"], "first_pass": {**PIPELINE_CONFIG_YAML.get("first_pass", {}), **pass_args}, "second_pass": {**PIPELINE_CONFIG_YAML.get("second_pass", {}), **pass_args}, "temporal_upsampler": temporal_upsampler_to_use } # Configure TeaCache for the *first pass* # (Note: This is imperfect, as the second pass will run without a reset) first_pass_steps = multi_scale_kwargs.get("first_pass", {}).get("num_inference_steps", num_steps) pipeline_instance.transformer.num_steps = first_pass_steps pipeline_instance.transformer.cnt = 0 pipeline_instance.transformer.previous_hidden_states = None pipeline_instance.transformer.previous_residual = None pipeline_instance.transformer.accumulated_rel_l1_distance = 0 result_images_tensor = multi_scale_pipeline(**multi_scale_kwargs).images else: # --- Configure TeaCache for a single pass --- pipeline_instance.transformer.num_steps = num_steps pipeline_instance.transformer.cnt = 0 pipeline_instance.transformer.previous_hidden_states = None pipeline_instance.transformer.previous_residual = None pipeline_instance.transformer.accumulated_rel_l1_distance = 0 single_pass_kwargs = {**call_kwargs, "guidance_scale": float(ui_guidance_scale), **PIPELINE_CONFIG_YAML.get("first_pass", {})} result_images_tensor = pipeline_instance(**single_pass_kwargs).images if result_images_tensor is None: raise gr.Error("Generation failed.") pad_l, pad_r, pad_t, pad_b = padding_values result_images_tensor = result_images_tensor[:, :, :actual_num_frames, pad_t:(-pad_b or None), pad_l:(-pad_r or None)] video_np = (np.clip(result_images_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy(), 0, 1) * 255).astype(np.uint8) output_video_path = os.path.join(tempfile.mkdtemp(), f"output_{random.randint(10000,99999)}.mp4") with imageio.get_writer(output_video_path, format='FFMPEG', fps=call_kwargs["frame_rate"], codec='libx264', quality=10, pixelformat='yuv420p') as video_writer: for idx, frame in enumerate(video_np): progress(idx / len(video_np), desc="Saving video clip...") video_writer.append_data(frame) #upload_to_sftp(output_video_path) # --- Extract Last Frame Tensor for Chaining --- # Get the last frame tensor, but KEEP the 'N' dimension (shape [b, c, 1, h, w]) last_frame_tensor_unnormalized = result_images_tensor[:, :, [-1], :, :].clone().cpu() # Normalize it to -1 to 1 range, which the conditioning input expects last_frame_tensor = (last_frame_tensor_unnormalized * 2.0) - 1.0 updated_clips_list = clips_list + [output_video_path] counter_text = f"Clips created: {len(updated_clips_list)}" # Return updates for new UI elements return output_video_path, seed_ui, gr.update(visible=True), updated_clips_list, counter_text, gr.update(visible=True, value=True), last_frame_tensor, gr.update(value=False) def update_task_image(): return "image-to-video" def update_task_text(): return "text-to-video" def update_task_video(): return "video-to-video" css=""" #col-container{margin:0 auto;max-width:900px;} """ with gr.Blocks(css=css) as demo: clips_state = gr.State([]) last_frame_tensor_state = gr.State(value=None) # <-- For Direct Tensor Chaining gr.Markdown("# LTX Video Clip Stitcher") gr.Markdown("Generate short video clips and stitch them together to create a longer animation.") with gr.Row(): with gr.Column(): with gr.Tabs() as tabs: with gr.Tab("image-to-video", id="i2v_tab") as image_tab: video_i_hidden = gr.Textbox(visible=False); image_i2v = gr.Image(label="Input Image", type="filepath", sources=["upload", "webcam", "clipboard"]); i2v_prompt = gr.Textbox(label="Prompt", value="The character from the image starts to move.", lines=3); i2v_button = gr.Button("Generate Image-to-Video Clip", variant="primary") with gr.Tab("text-to-video", id="t2v_tab") as text_tab: image_n_hidden = gr.Textbox(visible=False); video_n_hidden = gr.Textbox(visible=False); 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 Clip", variant="primary") with gr.Tab("video-to-video", id="v2v_tab") as video_tab: image_v_hidden = gr.Textbox(visible=False); 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=120, value=9, step=8, info="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 Clip", variant="primary") duration_input = gr.Slider(label="Clip Duration (seconds)", minimum=1.0, maximum=10.0, value=2.0, step=0.1) improve_texture = gr.Checkbox(label="Improve Texture (multi-scale)", value=True) use_temporal_upscaler = gr.Checkbox(label="Use Temporal Upscaler (for smoothness)", value=True) use_last_tensor_toggle = gr.Checkbox(label="Use Last Frame (Direct Tensor)", value=False, visible=False) # <-- For Chaining # REMOVED: enhance_checkbox # REMOVED: superres_checkbox with gr.Column(): output_video = gr.Video(label="Last Generated Clip", interactive=False) use_last_frame_button = gr.Button("Use Last Frame as Input Image", visible=False) with gr.Accordion("Stitching Controls", open=True): clip_counter_display = gr.Markdown("Clips created: 0") with gr.Row(): stitch_button = gr.Button("🎬 Stitch All Clips"); clear_button = gr.Button("🗑️ Clear All Clips") final_video_output = gr.Video(label="Final Stitched Video", interactive=False) with gr.Accordion("Advanced settings", open=False): mode = gr.Dropdown(["text-to-video", "image-to-video", "video-to-video"], label="task", value="image-to-video", visible=False); negative_prompt_input = gr.Textbox(label="Negative Prompt", value="worst quality, inconsistent motion, blurry, jittery, distorted, low clarity, low resolution, grainy, pixelated, oversaturated, glitchy, noisy", lines=2) with gr.Row(): teacache_checkbox = gr.Checkbox(label="Enable TeaCache Acceleration", value=False) teacache_slider = gr.Slider( minimum=0.01, maximum=0.1, step=0.01, value=0.05, label="TeaCache Threshold (Higher = Faster)" ) with gr.Row(): seed_input = gr.Number(label="Seed", value=42, precision=0); randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=True) with gr.Row(visible=True): # <-- MODIFIED 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) with gr.Row(): # --- NEW UI COMPONENT --- text_encoder_max_tokens_input = gr.Slider( label="Text Encoder Max Tokens (Affects Speed/Fidelity)", minimum=16, maximum=300, value=300, step=1 ) with gr.Row(): height_input = gr.Slider(label="Height", value=1024, step=32, minimum=32, maximum=MAX_IMAGE_SIZE); width_input = gr.Slider(label="Width", value=1024, step=32, minimum=32, maximum=MAX_IMAGE_SIZE); with gr.Row(): image_cond_noise_scale_input = gr.Slider( label="Image Conditioning Noise Scale", minimum=0.0, maximum=1.0, value=0.00, step=0.01, info="Controls how much noise is added to the input image/frame. Lower values preserve the image better, higher values allow more creative freedom (0.0 for near-perfect preservation)." ) num_steps = gr.Slider(label="Steps", value=35, step=1, minimum=1, maximum=420); fps = gr.Slider(label="FPS", value=24.0, step=1.0, minimum=4.0, maximum=60.0) def handle_image_upload_for_dims(f, h, w): if not f: return gr.update(value=h), gr.update(value=w) img = Image.open(f); new_h, new_w = calculate_new_dimensions(img.width, img.height); return gr.update(value=new_h), gr.update(value=new_w) def handle_video_upload_for_dims(f, h, w): if not f or not os.path.exists(str(f)): return gr.update(value=h), gr.update(value=w) with imageio.get_reader(str(f)) as reader: meta = reader.get_meta_data(); orig_w, orig_h = meta.get('size', (reader.get_data(0).shape[1], reader.get_data(0).shape[0])); new_h, new_w = calculate_new_dimensions(orig_w, orig_h); return gr.update(value=new_h), gr.update(value=new_w) image_i2v.upload(handle_image_upload_for_dims, [image_i2v, height_input, width_input], [height_input, width_input]); video_v2v.upload(handle_video_upload_for_dims, [video_v2v, height_input, width_input], [height_input, width_input]); image_tab.select(update_task_image, outputs=[mode]); text_tab.select(update_task_text, outputs=[mode]); video_tab.select(update_task_video, outputs=[mode]) # --- FIX: Add new UI toggles to common_params --- common_params = [ height_input, width_input, mode, duration_input, frames_to_use, seed_input, randomize_seed_input, guidance_scale_input, improve_texture, use_temporal_upscaler, use_last_tensor_toggle, num_steps, fps, teacache_checkbox, teacache_slider, text_encoder_max_tokens_input, # Already there image_cond_noise_scale_input # <--- NEW SLIDER ADDED HERE ] t2v_inputs = [t2v_prompt, negative_prompt_input, clips_state, image_n_hidden, video_n_hidden, last_frame_tensor_state] + common_params; i2v_inputs = [i2v_prompt, negative_prompt_input, clips_state, image_i2v, video_i_hidden, last_frame_tensor_state] + common_params; v2v_inputs = [v2v_prompt, negative_prompt_input, clips_state, image_v_hidden, video_v2v, last_frame_tensor_state] + common_params # --- FIX: Add new UI elements to outputs --- gen_outputs = [ output_video, seed_input, use_last_frame_button, clips_state, clip_counter_display, use_last_tensor_toggle, last_frame_tensor_state, randomize_seed_input # <-- ADDED THIS ] # This function now needs to hide both buttons hide_btn = lambda: (gr.update(visible=False), gr.update(visible=False)) t2v_button.click( hide_btn, outputs=[use_last_frame_button, use_last_tensor_toggle], queue=False ).then( fn=generate, inputs=t2v_inputs, outputs=gen_outputs, api_name="text_to_video" ) i2v_button.click( hide_btn, outputs=[use_last_frame_button, use_last_tensor_toggle], queue=False ).then( fn=generate, inputs=i2v_inputs, outputs=gen_outputs, api_name="image_to_video" ) v2v_button.click( hide_btn, outputs=[use_last_frame_button, use_last_tensor_toggle], queue=False ).then( fn=generate, inputs=v2v_inputs, outputs=gen_outputs, api_name="video_to_video" ) # MODIFIED: Removed enhance_checkbox and superres_checkbox from inputs use_last_frame_button.click(fn=use_last_frame_as_input, inputs=[output_video], outputs=[image_i2v, tabs]) stitch_button.click(fn=stitch_videos, inputs=[clips_state], outputs=[final_video_output]) # Clear button also needs to reset the tensor state and hide the toggle clear_button.click(fn=clear_clips, outputs=[clips_state, clip_counter_display, output_video, final_video_output, use_last_tensor_toggle, last_frame_tensor_state, randomize_seed_input]) if __name__ == "__main__": if os.path.exists(models_dir): print(f"Model directory: {Path(models_dir).resolve()}") demo.queue().launch(debug=True, share=True, mcp_server=True)