"""LTX-2.3-22B Cinemagraph LoRA — image-to-video demo (ZeroGPU).

Turns a still image into a looping cinemagraph with selective motion, using the
`Lightricks/LTX-2.3-22b-LoRA-Cinemagraph` LoRA on top of the LTX-2.3-22B base.

Design (per maintainer instructions):
- Uses the LATEST LTX-2 codebase (github.com/Lightricks/LTX-2 main, no pinned
  stale commit) so pipeline code matches the model card's inference assumptions.
- Native single-stage `TI2VidOneStagePipeline` with real CFG guidance — matches
  the model card's recommended settings (steps 30, guidance 4.0, STG scale 1.0
  on block 29). bf16 everywhere, NO fp8 quantization, NO LoRA fusion.
- The Cinemagraph LoRA is loaded with the ComfyUI key-renaming map (its keys are
  prefixed with `diffusion_model.`), strength user-tunable in 0.7–3.0.
- Text encoding (Gemma-3 + LTX-2.3 connectors) is offloaded to a remote encoder
  Space (linoyts/ltx23-gemma-encoder-api) via gradio_client, so Gemma never
  loads locally. `pipeline.prompt_encoder` is swapped for a stub that returns the
  precomputed positive/negative embeddings (structure mirrors linoyts/ltx23-dev-api).

At the current HEAD, LTX-2's attention auto-selector already avoids FA3/FA4 on
consumer Blackwell (sm_120) and falls back to torch SDPA — so no attention
monkeypatching is needed (this was the source of the old pinned-commit issues).
"""

import os
import subprocess
import sys

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ["TORCH_COMPILE_DISABLE"] = "1"
os.environ["TORCHDYNAMO_DISABLE"] = "1"

# --- Fetch the LATEST LTX-2 codebase (no stale pinned commit) ---------------
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")

if not os.path.exists(LTX_REPO_DIR):
    subprocess.run(
        ["git", "clone", "--depth", "1", LTX_REPO_URL, LTX_REPO_DIR], check=True
    )

subprocess.run(
    [sys.executable, "-m", "pip", "install", "--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 json
import logging
import random
import struct
import tempfile

import spaces
import torch

torch._dynamo.config.suppress_errors = True
torch._dynamo.config.disable = True

import gradio as gr
import numpy as np
from gradio_client import Client, handle_file
from huggingface_hub import hf_hub_download

from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.components.schedulers import LTX2Scheduler
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
from ltx_pipelines.utils import get_device
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.blocks import (
    AudioDecoder,
    DiffusionStage,
    ImageConditioner,
    VideoDecoder,
)
from ltx_pipelines.utils.media_io import encode_video

# --- Chunked-read safetensors loader --------------------------------------
# safe_open's mmap deadlocks on the FUSE-backed HF cache; read shards manually.
from ltx_core.loader.primitives import StateDict
from ltx_core.loader.sft_loader import SafetensorsStateDictLoader

_SAFETENSORS_DTYPE_MAP = {
    "F64": torch.float64, "F32": torch.float32, "F16": torch.float16,
    "BF16": torch.bfloat16, "F8_E5M2": torch.float8_e5m2, "F8_E4M3": torch.float8_e4m3fn,
    "I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
    "U8": torch.uint8, "BOOL": torch.bool,
}


def _patched_load(self, path, sd_ops, device=None):
    sd, size, dtype = {}, 0, set()
    device = device or torch.device("cpu")
    for shard_path in (path if isinstance(path, list) else [path]):
        with open(shard_path, "rb") as f:
            header_len = struct.unpack("<Q", f.read(8))[0]
            header = json.loads(f.read(header_len).decode("utf-8"))
            data_base = 8 + header_len
            for name, meta in header.items():
                if name == "__metadata__":
                    continue
                expected_name = name if sd_ops is None else sd_ops.apply_to_key(name)
                if expected_name is None:
                    continue
                start, end = meta["data_offsets"]
                f.seek(data_base + start)
                buf = f.read(end - start)
                t = torch.frombuffer(bytearray(buf), dtype=_SAFETENSORS_DTYPE_MAP[meta["dtype"]]
                                     ).reshape(meta["shape"])
                t = t.to(device=device, non_blocking=True, copy=False)
                kvs = (((expected_name, t),) if sd_ops is None
                       else sd_ops.apply_to_key_value(expected_name, t))
                for key, v in kvs:
                    size += v.nbytes
                    dtype.add(v.dtype)
                    sd[key] = v
    return StateDict(sd=sd, device=device, size=size, dtype=dtype)


SafetensorsStateDictLoader.load = _patched_load

logging.getLogger().setLevel(logging.INFO)
MAX_SEED = np.iinfo(np.int32).max
FRAME_RATE = 24.0
TRIGGER = "CINEMAGRAPH_MOTION"

# Model card recommended settings for the Cinemagraph LoRA.
DEFAULT_STEPS = 30
DEFAULT_CFG = 4.0            # guidance scale 4.0 (non-distilled path)
DEFAULT_LORA_STRENGTH = 1.0  # recommended range 0.7-3.0
VIDEO_STG_SCALE = 1.0        # STG scale 1.0
VIDEO_STG_BLOCKS = [29]      # target block 29
VIDEO_RESCALE = 0.7          # LTX-2.3 default
A2V_SCALE = 3.0              # LTX-2.3 default modality scale
# Audio guider defaults (LTX-2.3).
AUDIO_CFG, AUDIO_STG, AUDIO_RESCALE, V2A_SCALE = 7.0, 1.0, 0.7, 3.0

# Best results at 512x704 (training resolution). Single-stage needs /32.
DEFAULT_WIDTH = 704
DEFAULT_HEIGHT = 512
DEFAULT_NUM_FRAMES = 121  # ~5s at 24fps; must be 8k+1

DEFAULT_NEGATIVE = (
    "cars moving, camera movement, pan, tilt, zoom, parallax, whole image moving, "
    "background sliding, person moving, flicker on entire image, noisy texture, "
    "crawling texture, distorted, blurry, low quality"
)

LTX_REPO = os.environ.get("LTX_REPO", "Lightricks/LTX-2.3")
DEV_FILE = "ltx-2.3-22b-dev.safetensors"
LORA_REPO = "Lightricks/LTX-2.3-22b-LoRA-Cinemagraph"
LORA_FILE = "ltx-2.3-22b-lora-cinemagraph-0.9.safetensors"
TOKEN = os.environ.get("HF_TOKEN")
TEXT_ENCODER_SPACE = os.environ.get("TEXT_ENCODER_SPACE", "linoyts/ltx23-gemma-encoder-api")

print("Downloading base checkpoint + Cinemagraph LoRA…")
checkpoint_path = hf_hub_download(LTX_REPO, DEV_FILE, token=TOKEN)
lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=TOKEN)


# --- Remote Gemma text encoder (positive + negative, for CFG) --------------
# The remote Space returns a .pt: {"positive": {video, audio, mask},
# "negative": {...} | None, "final_prompt": str}. We reconstruct
# EmbeddingsProcessorOutput and hand them to the pipeline via a stub encoder.

_emb_cache: dict = {}


def _to_output(pack, device) -> EmbeddingsProcessorOutput:
    return EmbeddingsProcessorOutput(
        video_encoding=pack["video"].to(device),
        audio_encoding=pack["audio"].to(device) if pack["audio"] is not None else None,
        attention_mask=pack["mask"].to(device),
    )


def fetch_embeddings(prompt, negative_prompt, enhance_prompt, seed):
    """Call the Gemma encoder Space; returns the loaded .pt dict."""
    key = (prompt, negative_prompt, bool(enhance_prompt), int(seed) if enhance_prompt else 0)
    if key in _emb_cache:
        return _emb_cache[key]
    client = Client(TEXT_ENCODER_SPACE, token=TOKEN, httpx_kwargs={"timeout": 300})
    emb_file, final_prompt, status = client.predict(
        prompt=prompt, negative_prompt=negative_prompt or "", encode_negative=True,
        enhance_prompt=enhance_prompt, seed=int(seed), api_name="/encode",
    )
    print(f"[encoder] {status} | final prompt: {final_prompt[:80]}…")
    data = torch.load(emb_file, map_location="cpu", weights_only=True)
    if len(_emb_cache) > 32:
        _emb_cache.clear()
    _emb_cache[key] = data
    return data


class _RemotePromptEncoder:
    """Drop-in replacement for the pipeline's PromptEncoder.

    Its __call__ signature matches the native PromptEncoder
    (prompts, *, enhance_first_prompt, enhance_prompt_image, enhance_prompt_seed)
    but returns precomputed [positive, negative] EmbeddingsProcessorOutput from
    the remote Gemma Space, so Gemma is never loaded on this worker.
    """

    def __init__(self):
        self._pending = None

    def set(self, embeddings, device):
        self._pending = (embeddings, device)

    def __call__(self, prompts, *, enhance_first_prompt=False,
                 enhance_prompt_image=None, enhance_prompt_seed=42):
        embeddings, device = self._pending
        return [
            _to_output(embeddings["positive"], device),
            _to_output(embeddings["negative"], device),
        ]


print("Building single-stage pipeline (bf16, Cinemagraph LoRA, no quant)…")
# We bypass TI2VidOneStagePipeline.__init__ because it eagerly constructs the
# native PromptEncoder, which calls find_matching_file(gemma_root, ...) and
# crashes when gemma_root is None. Instead we build every component EXCEPT the
# prompt encoder by hand (identical arguments to the constructor), and plug in
# the remote Gemma stub. Text encoding therefore never loads Gemma here.
_DEVICE = get_device()
_DTYPE = torch.bfloat16
_lora = LoraPathStrengthAndSDOps(
    path=lora_path, strength=DEFAULT_LORA_STRENGTH, sd_ops=LTXV_LORA_COMFY_RENAMING_MAP)

pipeline = object.__new__(TI2VidOneStagePipeline)
pipeline.dtype = _DTYPE
pipeline.device = _DEVICE
pipeline._scheduler = LTX2Scheduler()
_remote_encoder = _RemotePromptEncoder()
pipeline.prompt_encoder = _remote_encoder
pipeline.image_conditioner = ImageConditioner(
    checkpoint_path=checkpoint_path, dtype=_DTYPE, device=_DEVICE)
pipeline.stage = DiffusionStage.from_checkpoint(
    checkpoint_path=checkpoint_path, dtype=_DTYPE, device=_DEVICE,
    loras=(_lora,), quantization=None)  # plain bf16 — NO fp8, NO fusion
pipeline.video_decoder = VideoDecoder(
    checkpoint_path=checkpoint_path, dtype=_DTYPE, device=_DEVICE)
pipeline.audio_decoder = AudioDecoder(
    checkpoint_path=checkpoint_path, dtype=_DTYPE, device=_DEVICE)
print("Pipeline ready.")


def _round_frames(duration_s: float) -> int:
    n = int(round(duration_s * FRAME_RATE))
    return ((n - 1 + 7) // 8) * 8 + 1  # nearest valid 8k+1, >= 9


def gpu_duration(embeddings, image_path, prompt, negative_prompt, num_frames, used_seed,
                 num_inference_steps=DEFAULT_STEPS, *args, **kwargs):
    # Per-call cost = 46GB bf16 transformer build (from FUSE cache) + denoise.
    # Scales mildly with steps and frame count.
    base = 300
    return min(1500, int(base + int(num_inference_steps) * 8 + int(num_frames) * 0.6))


@spaces.GPU(duration=gpu_duration, size="xlarge")
@torch.inference_mode()
def _generate_gpu(embeddings, image_path, prompt, negative_prompt, num_frames, used_seed,
                  num_inference_steps, cfg_scale, lora_strength, width, height,
                  progress=gr.Progress(track_tqdm=True)):
    tiling_config = TilingConfig.default()
    video_chunks_number = get_video_chunks_number(num_frames, tiling_config)

    # Rebuild the diffusion stage with the requested LoRA strength (functional swap).
    stage = pipeline.stage
    if abs(lora_strength - DEFAULT_LORA_STRENGTH) > 1e-6:
        stage = pipeline.stage.with_loras(
            (LoraPathStrengthAndSDOps(
                path=lora_path, strength=float(lora_strength), sd_ops=LTXV_LORA_COMFY_RENAMING_MAP),)
        )

    _remote_encoder.set(embeddings, pipeline.device)

    images = [ImageConditioningInput(path=image_path, frame_idx=0, strength=1.0)]

    orig_stage = pipeline.stage
    pipeline.stage = stage
    try:
        video, audio = pipeline(
            prompt=prompt,
            negative_prompt=negative_prompt,
            seed=used_seed,
            height=int(height),
            width=int(width),
            num_frames=num_frames,
            frame_rate=FRAME_RATE,
            num_inference_steps=int(num_inference_steps),
            video_guider_params=MultiModalGuiderParams(
                cfg_scale=float(cfg_scale), stg_scale=VIDEO_STG_SCALE,
                rescale_scale=VIDEO_RESCALE, modality_scale=A2V_SCALE,
                skip_step=0, stg_blocks=VIDEO_STG_BLOCKS),
            audio_guider_params=MultiModalGuiderParams(
                cfg_scale=AUDIO_CFG, stg_scale=AUDIO_STG,
                rescale_scale=AUDIO_RESCALE, modality_scale=V2A_SCALE,
                skip_step=0, stg_blocks=VIDEO_STG_BLOCKS),
            images=images,
            tiling_config=tiling_config,
            enhance_prompt=False,  # enhancement already done on the encoder Space
        )
    finally:
        pipeline.stage = orig_stage

    output_path = tempfile.mktemp(suffix=".mp4")
    encode_video(video=video, fps=FRAME_RATE, audio=None, output_path=output_path,
                 video_chunks_number=video_chunks_number)
    return output_path


def _ensure_trigger(prompt: str) -> str:
    p = (prompt or "").strip()
    if TRIGGER.lower() not in p.lower():
        p = f"{TRIGGER}, {p}" if p else TRIGGER
    return p


def generate(
    image_path: str,
    prompt: str,
    negative_prompt: str = DEFAULT_NEGATIVE,
    duration: float = 5.0,
    seed: int = 42,
    randomize_seed: bool = True,
    num_inference_steps: int = DEFAULT_STEPS,
    cfg_scale: float = DEFAULT_CFG,
    lora_strength: float = DEFAULT_LORA_STRENGTH,
    width: int = DEFAULT_WIDTH,
    height: int = DEFAULT_HEIGHT,
    enhance_prompt: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a looping cinemagraph video from a still image.

    Args:
        image_path: input still image (filepath). Its first frame is preserved
            and only the described element is animated.
        prompt: what should move; be explicit about what stays frozen. The
            `CINEMAGRAPH_MOTION` trigger is auto-prepended if missing.
        negative_prompt: things to avoid (camera motion, people moving, etc.).
        duration: video length in seconds (~24 fps).
        seed: RNG seed.
        randomize_seed: pick a random seed each run.
        num_inference_steps: denoising steps (30 recommended).
        cfg_scale: classifier-free guidance scale (4.0 recommended).
        lora_strength: Cinemagraph LoRA weight (0.7-3.0; higher = stronger motion).
        width: output width (multiple of 32).
        height: output height (multiple of 32).
        enhance_prompt: let Gemma rewrite the prompt before encoding.

    Returns:
        (video_path, used_seed)
    """
    if not image_path:
        raise gr.Error("Please upload an input image.")
    used_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    prompt = _ensure_trigger(prompt)
    width = int(width) // 32 * 32
    height = int(height) // 32 * 32
    num_frames = _round_frames(duration)
    progress(0.05, desc="encoding prompt (Gemma Space)…")
    embeddings = fetch_embeddings(prompt, negative_prompt, enhance_prompt, used_seed)
    output_path = _generate_gpu(
        embeddings, image_path, prompt, negative_prompt, num_frames, used_seed,
        num_inference_steps, cfg_scale, lora_strength, width, height)
    return output_path, used_seed


CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    ["examples/Glasses.jpeg",
     "CINEMAGRAPH_MOTION, tripod locked-off static camera, zero camera movement, the man, "
     "face, hair, clothing, beach, sky, and background remain completely frozen, only the "
     "beach reflection inside his sunglasses moves, reflected time-lapse of ocean waves roll "
     "and shimmer in the lenses, everything outside the glasses stays still, seamless natural loop"],
    ["examples/Motel.jpeg",
     "CINEMAGRAPH_MOTION, tripod locked-off static camera, zero camera movement, only the neon "
     "sign flickers, the starburst, Desert, MOTEL, and NO VACANCY neon tubes subtly pulse and "
     "flicker like a real vintage neon sign, everything else stays perfectly still, seamless natural loop"],
    ["examples/Jump.jpeg",
     "CINEMAGRAPH_MOTION, tripod locked-off static camera, zero camera movement, the man and cow "
     "remain completely frozen, only the clouds in the sky move in fast time-lapse, clouds drift "
     "and roll across the sky, grass and foreground remain still, seamless natural loop."],
    ["examples/Tears.jpeg",
     "CINEMAGRAPH_MOTION, tripod locked-off static camera, zero camera movement, the woman is "
     "frozen, her face, eyes, hair, clothing, seat, and airplane interior remain completely frozen, "
     "only the illustrated tear drops move downward along her cheeks like a cutout collage "
     "animation, seamless natural loop."],
]

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="LTX-2.3 Cinemagraph") as demo:
    gr.Markdown(
        "# 🎞️ LTX-2.3-22B Cinemagraph\n"
        "Turn a **still image** into a looping **cinemagraph** — one element moves "
        "(water, neon, clouds, reflections) while the rest of the scene stays frozen.\n\n"
        "LoRA: [`Lightricks/LTX-2.3-22b-LoRA-Cinemagraph`]"
        "(https://huggingface.co/Lightricks/LTX-2.3-22b-LoRA-Cinemagraph) on "
        "[LTX-2.3-22B](https://huggingface.co/Lightricks/LTX-2.3). Be explicit about what moves "
        "and what stays still. Avoid camera language. `CINEMAGRAPH_MOTION` is added automatically."
    )
    with gr.Row():
        with gr.Column():
            image = gr.Image(label="Input image", type="filepath")
            prompt = gr.Textbox(
                label="Prompt (what moves, what stays frozen)", lines=4,
                placeholder="only the water in the fountain flows; everything else frozen; seamless natural loop")
            with gr.Row():
                duration = gr.Slider(1, 8, value=5, step=0.5, label="Duration (s)")
                run = gr.Button("Generate", variant="primary", scale=1)
            with gr.Accordion("Advanced settings", open=False):
                negative_prompt = gr.Textbox(label="Negative prompt", value=DEFAULT_NEGATIVE, lines=3)
                with gr.Row():
                    seed = gr.Number(label="Seed", value=42, precision=0)
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                with gr.Row():
                    num_inference_steps = gr.Slider(15, 40, value=DEFAULT_STEPS, step=1, label="Steps")
                    cfg_scale = gr.Slider(1.0, 8.0, value=DEFAULT_CFG, step=0.1, label="Guidance (CFG)")
                lora_strength = gr.Slider(0.7, 3.0, value=DEFAULT_LORA_STRENGTH, step=0.1,
                                          label="Cinemagraph LoRA strength")
                with gr.Row():
                    width = gr.Slider(320, 1024, value=DEFAULT_WIDTH, step=32, label="Width")
                    height = gr.Slider(320, 1024, value=DEFAULT_HEIGHT, step=32, label="Height")
                enhance_prompt = gr.Checkbox(label="Enhance prompt (Gemma rewrite)", value=False)
        with gr.Column():
            video = gr.Video(label="Cinemagraph", autoplay=True, loop=True)
            used_seed = gr.Number(label="Used seed", precision=0)

    inputs = [image, prompt, negative_prompt, duration, seed, randomize_seed,
              num_inference_steps, cfg_scale, lora_strength, width, height, enhance_prompt]
    run.click(generate, inputs, [video, used_seed], api_name="generate")

    gr.Examples(
        examples=EXAMPLES,
        inputs=[image, prompt],
        outputs=[video, used_seed],
        fn=generate,
        cache_examples=False,   # ZeroGPU + heavy per-call build: run on click
        run_on_click=True,
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True)