app.py

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)