import glob import os import subprocess import sys # Disable torch.compile / dynamo before any torch import os.environ["TORCH_COMPILE_DISABLE"] = "1" os.environ["TORCHDYNAMO_DISABLE"] = "1" # --- PULL THE HF TOKEN SECURELY FROM THE SPACE'S SECRETS --- HF_TOKEN = os.environ.get("HF_TOKEN") if not HF_TOKEN: print("*** WARNING: HF_TOKEN environment variable not found! Gated models like Gemma will fail to download.") # Clone LTX-2 repo and install packages LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git" LTX_REPO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "LTX-2") LTX_COMMIT_SHA = "ae855f8538843825f9015a419cf4ba5edaf5eec2" if not os.path.exists(LTX_REPO_DIR): print(f"Cloning {LTX_REPO_URL}... ") os.makedirs(LTX_REPO_DIR) subprocess.run(["git", "init", LTX_REPO_DIR], check=True) subprocess.run(["git", "remote", "add", "origin", LTX_REPO_URL], cwd=LTX_REPO_DIR, check=True) subprocess.run(["git", "fetch", "--depth", "1", "origin", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, check=True) subprocess.run(["git", "checkout", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, check=True) freshly_cloned = True else: freshly_cloned = False if freshly_cloned: print("Installing ltx-core and ltx-pipelines from cloned repo... ") subprocess.run( [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-core"), "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines")], check=True, ) sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines", "src")) sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-core", "src")) import logging import random import tempfile from pathlib import Path import torch import torch._dynamo import spaces import gradio as gr import numpy as np import transformers from huggingface_hub import hf_hub_download, snapshot_download from ltx_core.model.video_vae import SpatialTilingConfig, TemporalTilingConfig, TilingConfig, get_video_chunks_number from ltx_core.quantization import QuantizationPolicy from ltx_pipelines.distilled import DistilledPipeline from ltx_pipelines.utils.args import ImageConditioningInput from ltx_pipelines.utils.media_io import encode_video # Force-patch xformers attention into the LTX attention module. from ltx_core.model.transformer import attention as _attn_mod print(f"[ATTN] Before patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}") try: from xformers.ops import memory_efficient_attention as _mea _attn_mod.memory_efficient_attention = _mea print(f"[ATTN] After patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}") except Exception as e: print(f"[ATTN] xformers patch FAILED: {type(e).__name__}: {e}") try: from xformers.ops.fmha import _set_use_fa3 _set_use_fa3(False) print("[ATTN] xformers FA3 dispatch disabled (Blackwell-incompatible)") except Exception as e: print(f"[ATTN] FA3 disable FAILED: {type(e).__name__}: {e}") # --- PATCH 1: TRANSFORMERS 5.x COMPATIBILITY (SiglipVisionModel) --- try: from transformers.models.siglip import modeling_siglip if hasattr(modeling_siglip, 'SiglipVisionModel'): SiglipVisionModel = modeling_siglip.SiglipVisionModel if not hasattr(SiglipVisionModel, 'vision_model'): print("[PATCH 1] Adding missing 'vision_model' property to SiglipVisionModel for transformers 5.x...") SiglipVisionModel.vision_model = property(lambda self: self) except Exception as e: print(f"[PATCH 1] Could not patch SiglipVisionModel: {e}") # --- PATCH 2: TRANSFORMERS 5.x COMPATIBILITY (Gemma3TextConfig) --- try: from transformers.models.gemma3.configuration_gemma3 import Gemma3TextConfig if not hasattr(Gemma3TextConfig, '_patched_rope_local'): original_gemma3_init = Gemma3TextConfig.__init__ def patched_gemma3_init(self, *args, **kwargs): original_gemma3_init(self, *args, **kwargs) if not hasattr(self, 'rope_local_base_freq'): self.rope_local_base_freq = 10000.0 if hasattr(self, 'rope_scaling') and isinstance(self.rope_scaling, dict): if 'rope_type' not in self.rope_scaling: self.rope_scaling['rope_type'] = self.rope_scaling.get('type', 'default') Gemma3TextConfig.__init__ = patched_gemma3_init Gemma3TextConfig._patched_rope_local = True print("[PATCH 2] Patched Gemma3TextConfig for transformers 5.x (rope_local_base_freq + rope_type)...") except Exception as e: print(f"[PATCH 2] Could not patch Gemma3TextConfig: {e}") # --- PATCH 3: TRANSFORMERS 5.x COMPATIBILITY (ROPE_INIT_FUNCTIONS) --- try: from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS if "default" not in ROPE_INIT_FUNCTIONS: print("[PATCH 3] Injecting 'default' into ROPE_INIT_FUNCTIONS for transformers 5.x... ") def _compute_default_rope_parameters(config=None, device=None, seq_len=None, layer_type=None, **rope_kwargs): if len(rope_kwargs) > 0: base = rope_kwargs["base"] dim = rope_kwargs["dim"] elif config is not None: base = getattr(config, "rope_theta", getattr(config, "rope_local_base_freq", 10000.0)) partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0) head_dim = getattr(config, "head_dim", getattr(config, "hidden_size", 256) // getattr(config, "num_attention_heads", 8)) dim = int(head_dim * partial_rotary_factor) else: base = 10000.0 dim = 256 attention_factor = 1.0 inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim)) return inv_freq, attention_factor ROPE_INIT_FUNCTIONS["default"] = _compute_default_rope_parameters except Exception as e: print(f"[PATCH 3] Could not patch ROPE_INIT_FUNCTIONS: {e}") # --- PATCH 4: TRANSFORMERS 5.x COMPATIBILITY (Gemma3TextModel) --- try: import torch.nn as nn from transformers.models.gemma3.modeling_gemma3 import Gemma3TextModel if not hasattr(Gemma3TextModel, '_patched_rotary_emb_local'): original_gemma3_text_init = Gemma3TextModel.__init__ def patched_gemma3_text_init(self, *args, **kwargs): original_gemma3_text_init(self, *args, **kwargs) if not hasattr(self, 'rotary_emb_local'): self.rotary_emb_local = nn.Module() Gemma3TextModel.__init__ = patched_gemma3_text_init Gemma3TextModel._patched_rotary_emb_local = True print("[PATCH 4] Added dummy 'rotary_emb_local' to Gemma3TextModel for transformers 5.x...") except Exception as e: print(f"[PATCH 4] Could not patch Gemma3TextModel: {e}") # --- PATCH 5: SURGICAL KEY MAPPING FOR GEMMA 3 IN LTX-2 --- # LTX-2 expects 'model.model.vision_tower' but HF safetensors provides 'model.vision_tower' try: from ltx_core.loader.sft_loader import SafetensorsStateDictLoader from ltx_core.loader.primitives import StateDict _original_load = SafetensorsStateDictLoader.load def _patched_load(self, path, sd_ops, device=None): # Let the original loader handle all the complex dtype/device/quantization logic safely state_dict = _original_load(self, path, sd_ops, device) new_sd = {} for key, tensor in state_dict.sd.items(): if key.startswith("model.vision_tower."): new_key = "model.model." + key[6:] elif key.startswith("model.language_model."): new_key = "model.model." + key[6:] else: new_key = key new_sd[new_key] = tensor return StateDict( sd=new_sd, device=state_dict.device, size=state_dict.size, dtype=state_dict.dtype ) SafetensorsStateDictLoader.load = _patched_load print("[PATCH 5] SafetensorsStateDictLoader.load patched to fix Gemma 3 key mapping (model.* -> model.model.*)") except Exception as e: print(f"[PATCH 5] Could not apply key mapping patch: {e}") # --- PATCH 7: MATERIALIZE LEFTOVER META TENSORS (unused vision_tower / rotary buffers) --- # Also manually restores the cached_property caching contract that ModelLedger.text_encoder # originally had, since replacing the class attribute with a plain function loses it. try: from ltx_pipelines.utils import model_ledger as _ml def _materialize_meta_(module: torch.nn.Module, device): """Replace any still-meta params/buffers with real (zero) tensors in place, without touching tensors that already loaded correctly.""" for name, p in list(module.named_parameters(recurse=True)): if p.is_meta: *path, leaf = name.split(".") parent = module for part in path: parent = getattr(parent, part) new_param = torch.nn.Parameter( torch.zeros(p.shape, dtype=p.dtype, device=device), requires_grad=p.requires_grad, ) setattr(parent, leaf, new_param) for name, b in list(module.named_buffers(recurse=True)): if b.is_meta: *path, leaf = name.split(".") parent = module for part in path: parent = getattr(parent, part) new_buf = torch.zeros(b.shape, dtype=b.dtype, device=device) parent.register_buffer(leaf, new_buf, persistent=False) def _patched_text_encoder(self): model = self.text_encoder_builder.build(device=self._target_device(), dtype=self.dtype) target_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") _materialize_meta_(model, target_device) return model.to(target_device).eval() _ml.ModelLedger.text_encoder = _patched_text_encoder print("[PATCH 7] Patched ModelLedger.text_encoder to materialize leftover meta tensors " "AND restore caching (instance __dict__) so it isn't rebuilt on every call.") except Exception as e: print(f"[PATCH 7] Could not apply ModelLedger.text_encoder patch: {e}") logging.getLogger().setLevel(logging.INFO) MAX_SEED = np.iinfo(np.int32).max DEFAULT_FRAME_RATE = 24.0 RESOLUTIONS = { "high": {"16:9": (1536, 1024), "9:16": (1024, 1536), "1:1": (1024, 1024)}, "low": {"16:9": (768, 512), "9:16": (512, 768), "1:1": (768, 768)}, } BASE_DIR = os.path.dirname(os.path.abspath(__file__)) LTX_MOUNT = os.path.join(BASE_DIR, "models_ltx") GEMMA_MOUNT = "/home/user/app/models_gemma" DISTILLED_FILENAME = "ltx-2.3-22b-distilled-1.1.safetensors" UPSCALER_FILENAME = "ltx-2.3-spatial-upscaler-x2-1.1.safetensors" def download_assets_safeguard(): os.makedirs(GEMMA_MOUNT, exist_ok=True) os.makedirs(LTX_MOUNT, exist_ok=True) weights_present = bool(glob.glob(os.path.join(GEMMA_MOUNT, "model-*.safetensors"))) if not weights_present: print("Downloading Gemma 3 12B weights (safetensors)...") try: snapshot_download( repo_id="google/gemma-3-12b-it", local_dir=GEMMA_MOUNT, ignore_patterns=["*.bin", "*.pt", "original/*"], token=HF_TOKEN ) print("Gemma weights cached cleanly!") except Exception as e: print(f"Failed loading Gemma weights: {e}") distilled_path = os.path.join(LTX_MOUNT, DISTILLED_FILENAME) upscaler_path = os.path.join(LTX_MOUNT, UPSCALER_FILENAME) if not os.path.exists(distilled_path): print(f"Downloading {DISTILLED_FILENAME} into stable local mount...") hf_hub_download(repo_id="Lightricks/LTX-2.3", filename=DISTILLED_FILENAME, local_dir=LTX_MOUNT, token=HF_TOKEN) if not os.path.exists(upscaler_path): print(f"Downloading {UPSCALER_FILENAME} into stable local mount...") hf_hub_download(repo_id="Lightricks/LTX-2.3", filename=UPSCALER_FILENAME, local_dir=LTX_MOUNT, token=HF_TOKEN) print("[STARTUP] Warming up local asset registers...") download_assets_safeguard() pipeline = None def log_memory(tag: str): if torch.cuda.is_available(): allocated = torch.cuda.memory_allocated() / 1024**3 peak = torch.cuda.max_memory_allocated() / 1024**3 free, total = torch.cuda.mem_get_info() print(f"[VRAM {tag}] allocated={allocated:.2f}GB peak={peak:.2f}GB free={free / 1024**3:.2f}GB total={total / 1024**3:.2f}GB") def detect_aspect_ratio(image) -> str: if image is None: return "16:9" if hasattr(image, "size"): w, h = image.size elif hasattr(image, "shape"): h, w = image.shape[:2] else: return "16:9" ratio = w / h candidates = {"16:9": 16 / 9, "9:16": 9 / 16, "1:1": 1.0} return min(candidates, key=lambda k: abs(ratio - candidates[k])) def safe_image_upload(image, high_res): if image is None: return gr.update(), gr.update() try: aspect = detect_aspect_ratio(image) tier = "high" if high_res else "low" w, h = RESOLUTIONS[tier][aspect] return gr.update(value=w), gr.update(value=h) except Exception as e: print(f"[Warning] Failed to preprocess image upload: {e}") return gr.update(value=1536), gr.update(value=1024) TILING_PRESETS = { "default": TilingConfig( spatial_config=SpatialTilingConfig(tile_size_in_pixels=768, tile_overlap_in_pixels=64), temporal_config=TemporalTilingConfig(tile_size_in_frames=80, tile_overlap_in_frames=24), ), "low-memory": TilingConfig( spatial_config=SpatialTilingConfig(tile_size_in_pixels=512, tile_overlap_in_pixels=64), temporal_config=TemporalTilingConfig(tile_size_in_frames=48, tile_overlap_in_frames=16), ), "high-quality": TilingConfig( spatial_config=SpatialTilingConfig(tile_size_in_pixels=1024, tile_overlap_in_pixels=128), temporal_config=TemporalTilingConfig(tile_size_in_frames=128, tile_overlap_in_frames=32), ), } @spaces.GPU(duration=180) @torch.inference_mode() def generate_video( input_image, prompt: str, duration: float, enhance_prompt: bool, seamless_loop: bool, generate_audio: bool, custom_audio, seed: int, randomize_seed: bool, height: int, width: int, tiling_preset: str, progress=gr.Progress(track_tqdm=True), ): global pipeline if input_image is None: raise gr.Error("An Input Image is required for Image-to-Video generation!") try: torch.cuda.reset_peak_memory_stats() log_memory("start") if pipeline is None: print("Initializing DistilledPipeline on active ZeroGPU worker...") distilled_checkpoint_path = os.path.join(LTX_MOUNT, DISTILLED_FILENAME) spatial_upsampler_path = os.path.join(LTX_MOUNT, UPSCALER_FILENAME) pipeline = DistilledPipeline( distilled_checkpoint_path=distilled_checkpoint_path, spatial_upsampler_path=spatial_upsampler_path, gemma_root=GEMMA_MOUNT, loras=[], quantization=QuantizationPolicy.fp8_cast(), ) print("Pipeline successfully configured!") # CRITICAL FIX: Preload all models to avoid meta tensor initialization errors print("Preloading all models (including Gemma text encoder)...") ledger = pipeline.model_ledger _ = ledger.transformer() _ = ledger.video_encoder() _ = ledger.video_decoder() _ = ledger.audio_decoder() _ = ledger.vocoder() _ = ledger.spatial_upsampler() print("Core models preloaded successfully!") current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) frame_rate = DEFAULT_FRAME_RATE num_frames = int(duration * frame_rate) + 1 num_frames = ((num_frames - 1 + 7) // 8) * 8 + 1 print(f"Generating I2V: {height}x{width}, {num_frames} frames ({duration}s), seed={current_seed}") images = [] owned_temp_image = None output_dir = Path("outputs") output_dir.mkdir(exist_ok=True) temp_image_path = output_dir / f"temp_input_{current_seed}.jpg" if hasattr(input_image, "save"): input_image.save(temp_image_path) owned_temp_image = temp_image_path else: temp_image_path = Path(input_image) images = [ImageConditioningInput(path=str(temp_image_path), frame_idx=0, strength=1.0)] if seamless_loop: images.append( ImageConditioningInput(path=str(temp_image_path), frame_idx=num_frames - 1, strength=1.0) ) print(f"[LOOP] Seamless loop enabled: conditioning frame 0 and frame {num_frames - 1}") try: tiling_config = TILING_PRESETS.get(tiling_preset, TILING_PRESETS["default"]) video_chunks_number = get_video_chunks_number(num_frames, tiling_config) log_memory("before pipeline call") video, audio = pipeline( prompt=prompt, seed=current_seed, height=int(height), width=int(width), num_frames=num_frames, frame_rate=frame_rate, images=images, tiling_config=tiling_config, enhance_prompt=enhance_prompt, ) log_memory("after pipeline call") output_path = tempfile.mktemp(suffix=".mp4") encode_video( video=video, fps=frame_rate, audio=audio, output_path=output_path, video_chunks_number=video_chunks_number, ) log_memory("after encode_video") # --- AUDIO POST-PROCESSING --- if not generate_audio: silent_path = tempfile.mktemp(suffix=".mp4") try: subprocess.run( ["ffmpeg", "-y", "-i", output_path, "-an", "-c:v", "copy", silent_path], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL ) os.remove(output_path) output_path = silent_path except Exception as e: print(f"Failed to strip audio: {e}") if custom_audio is not None: final_path = tempfile.mktemp(suffix=".mp4") try: subprocess.run( [ "ffmpeg", "-y", "-i", output_path, "-i", custom_audio, "-c:v", "copy", "-c:a", "aac", "-map", "0:v:0", "-map", "1:a:0", "-shortest", final_path ], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL ) os.remove(output_path) output_path = final_path except Exception as e: print(f"Failed to mux custom audio: {e}") # ----------------------------- torch.cuda.empty_cache() return str(output_path), current_seed finally: if owned_temp_image is not None: Path(owned_temp_image).unlink(missing_ok=True) except Exception as e: import traceback log_memory("on error") torch.cuda.empty_cache() print(f"Error: {str(e)}\n{traceback.format_exc()}") raise gr.Error(f"Generation failed: {str(e)}") with gr.Blocks(title="LTX-2.3 I2V & Seamless Loops") as demo: gr.Markdown("# LTX-2.3 Distilled I2V: Image-to-Video & Seamless Loops") gr.Markdown( "High quality Image-to-Video generation with native Seamless Looping (First-Last Frame conditioning). \n" "[[model]](https://huggingface.co/Lightricks/LTX-2.3) " "[[code]](https://github.com/Lightricks/LTX-2)" ) with gr.Row(): with gr.Column(): input_image = gr.Image(label="Input Image (Required)", type="pil") prompt = gr.Textbox( label="Motion Prompt", info="Describe how the image should move. For loops, mention 'seamless cyclic motion'.", value="Make this image come alive with cinematic motion, smooth animation", lines=3, placeholder="Describe the motion and animation you want...", ) with gr.Row(): duration = gr.Slider(label="Duration (seconds)", minimum=1.0, maximum=10.0, value=3.0, step=0.1) with gr.Column(): enhance_prompt = gr.Checkbox(label="Enhance Prompt", value=False) high_res = gr.Checkbox(label="High Resolution", value=True) seamless_loop = gr.Checkbox( label="Seamless Loop", value=False, info="Forces the last frame to match the first frame for perfect looping." ) generate_audio = gr.Checkbox( label="Generate Native Audio", value=True, info="Uncheck to disable native AI audio (saves VRAM/time)." ) custom_audio = gr.Audio( label="Upload Custom Audio (Optional - Replaces native AI audio)", type="filepath" ) generate_btn = gr.Button("Generate I2V", variant="primary", size="lg") with gr.Accordion("Advanced Settings", open=False): seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=10, step=1) randomize_seed = gr.Checkbox(label="Randomize Seed", value=True) with gr.Row(): width = gr.Number(label="Width", value=1536, precision=0) height = gr.Number(label="Height", value=1024, precision=0) tiling_preset = gr.Dropdown( label="Tiling Preset", choices=list(TILING_PRESETS.keys()), value="default", info="default: balanced · low-memory: smaller tiles, less VRAM · high-quality: larger tiles, fewer seam artefacts", ) with gr.Column(): output_video = gr.Video(label="Generated Video", autoplay=True, loop=True) input_image.change( fn=safe_image_upload, inputs=[input_image, high_res], outputs=[width, height], ) high_res.change( fn=safe_image_upload, inputs=[input_image, high_res], outputs=[width, height], ) generate_btn.click( fn=generate_video, inputs=[ input_image, prompt, duration, enhance_prompt, seamless_loop, generate_audio, custom_audio, seed, randomize_seed, height, width, tiling_preset, ], outputs=[output_video, seed], ) css = """ .fillable{max-width: 1200px !important} .progress-text {color: white} """ if __name__ == "__main__": demo.launch(theme=gr.themes.Citrus(), css=css, ssr_mode=False)