import os
import subprocess
import sys

# ZeroGPU: torch.compile / dynamo unsupported — disable before any torch import.
os.environ["TORCH_COMPILE_DISABLE"] = "1"
os.environ["TORCHDYNAMO_DISABLE"] = "1"

# (removed runtime xformers install -> would pull torch 2.8 and break the AOTI .pt2; SDPA used)

# --- clone + install the NATIVE LTX-2 codebase at the pinned commit the working ZeroGPU spaces use ---
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 = "ae855f8538843825f9015a419cf4ba5edaf5eec2"
if not os.path.exists(LTX_REPO_DIR):
    subprocess.run(["git", "clone", LTX_REPO_URL, LTX_REPO_DIR], check=True)
    subprocess.run(["git", "-C", LTX_REPO_DIR, "checkout", LTX_COMMIT], check=True)
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

import numpy as np
import imageio.v3 as iio
from PIL import Image, ImageOps

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

import spaces
import gradio as gr
from huggingface_hub import hf_hub_download, snapshot_download

# Import LTX modules in the proven order — importing ltx_core.quantization/loader FIRST hits a
# circular import (fp8_cast <-> loader.fuse_loras). Importing the model modules first forces the
# correct init order (mirrors the working reference Space).
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number, decode_video as _vae_decode_video  # noqa: F401
from ltx_core.model.upsampler import upsample_video as _upsample_video  # noqa: F401
from ltx_core.model.audio_vae import encode_audio as _vae_encode_audio  # noqa: F401
from ltx_core.quantization import QuantizationPolicy
from ltx_core.loader import LoraPathStrengthAndSDOps, LTXV_LORA_COMFY_RENAMING_MAP
from ltx_pipelines.ic_lora import ICLoraPipeline
from ltx_pipelines.utils.media_io import encode_video

# --- ZeroGPU loader patch -------------------------------------------------------------
# The native loader opens safetensors directly on the CUDA device
# (safe_open(path, device="cuda")), doing the host->device copy in safetensors' own C++
# (cudaMemcpy) — bypassing torch.Tensor.to, the call ZeroGPU patches to virtualise + pack
# weights at module scope. Result: "No CUDA GPUs are available" at startup, nothing packs.
# Patch it to open on CPU then move via torch.Tensor.to (ZeroGPU-virtualisable).
import safetensors as _safetensors
import ltx_core.loader.sft_loader as _sft
from ltx_core.loader.primitives import StateDict as _StateDict

def _zerogpu_safe_load(self, path, sd_ops, device=None):
    device = device or torch.device("cpu")
    sd, size, dtype = {}, 0, set()
    model_paths = path if isinstance(path, list) else [path]
    for shard_path in model_paths:
        with _safetensors.safe_open(shard_path, framework="pt", device="cpu") as f:
            for name in f.keys():
                expected = name if sd_ops is None else sd_ops.apply_to_key(name)
                if expected is None:
                    continue
                value = f.get_tensor(name).to(device=device)  # torch path -> ZeroGPU-virtualised
                kvs = ((expected, value),)
                if sd_ops is not None:
                    kvs = sd_ops.apply_to_key_value(expected, value)
                for k, v in kvs:
                    size += v.nbytes
                    dtype.add(v.dtype)
                    sd[k] = v
    return _StateDict(sd=sd, device=device, size=size, dtype=dtype)

_sft.SafetensorsStateDictLoader.load = _zerogpu_safe_load
print("[PATCH] safetensors loader -> CPU-open + torch.to (ZeroGPU-virtualisable)")
# --------------------------------------------------------------------------------------

# --- attention backend patch (FA3 crashes on Blackwell ZeroGPU; use xformers/SDPA) ---
import torch.nn.functional as F
from ltx_core.model.transformer import attention as _attn_mod

def _sdpa_as_mea(query, key, value, attn_bias=None, scale=None, **kwargs):
    q, k, v = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
    return F.scaled_dot_product_attention(q, k, v, scale=scale).transpose(1, 2)

# IMPORTANT (ZeroGPU): never query CUDA at module scope. SDPA works on every GPU (incl.
# Blackwell ZeroGPU, where FA3 crashes), so patch it unconditionally.
_attn_mod.memory_efficient_attention = _sdpa_as_mea
print("[ATTN] SDPA (patched at module scope, no CUDA query)")

logging.getLogger().setLevel(logging.INFO)

# =========================== PER-LORA CONFIG (colorization) ===========================
TITLE = "LTX-2.3 Colorization + Old Film Hints (native LTX-2)"
LORA_REPO = "Lightricks/LTX-2.3-22b-IC-LoRA-Colorization"
LORA_FILE = "ltx-2.3-22b-ic-lora-colorization-0.9.safetensors"
LORA_SCALE = 1.0

# Identity-safe recipe recommended by Lightricks for this Colorization IC-LoRA:
# skip stage 2 and render on a 2x canvas, so the reference stays anchored.
SKIP_STAGE_2 = True

# The model was trained/validated near 960x544x121 @ 24fps. Higher resolutions are useful,
# but can weaken/partially wash out the colorization, so label them experimental.
RES_PRESETS = {
    "960×544 (training bucket / strongest color)": (960, 544),
    "1216×704 (balanced)": (1216, 704),
    "1536×864 (experimental high-res)": (1536, 864),
    "1920×1088 (experimental / can weaken color)": (1920, 1088),
}
DEFAULT_PRESET = "960×544 (training bucket / strongest color)"
FRAME_CHOICES = [49, 73, 97, 121]
DEFAULT_FRAMES = 121

RESTORATION_MODES = [
    "Colorize only",
    "Colorize + gentle old-film cleanup hint",
    "Colorize + archival restoration look",
    "Experimental high-detail / upscale wording",
]
DEFAULT_RESTORATION_MODE = "Colorize only"


def build_prompt(scene_description, color_description, restoration_mode):
    scene_description = (scene_description or "the provided grayscale / monochrome video").strip()
    color_description = (color_description or "natural, historically plausible colors for the same scene").strip()

    restoration_hint = ""
    if restoration_mode == "Colorize + gentle old-film cleanup hint":
        restoration_hint = (
            " The restored version has stable natural color, slightly reduced flicker, cleaner contrast, "
            "less yellowing or archival tint, and no added objects."
        )
    elif restoration_mode == "Colorize + archival restoration look":
        restoration_hint = (
            " The restored version looks like carefully restored archival film: natural skin tones, "
            "period-plausible wardrobe and materials, balanced contrast, preserved fine grain, and no modern recoloring artifacts."
        )
    elif restoration_mode == "Experimental high-detail / upscale wording":
        restoration_hint = (
            " The edited version has crisp natural colors, clear edges, preserved film grain, and a slightly cleaner high-detail finish, "
            "while avoiding plastic smoothing or geometry changes."
        )

    return (
        f"Reference shows {scene_description}, rendered in grayscale, monochrome, or heavily desaturated footage. "
        f"Edited shows the same scene with natural colors restored. COLORIZE {color_description}."
        f"{restoration_hint} "
        "Subject identity, framing, motion, and background geometry are identical to the reference; "
        "only color information differs between reference and edited."
    )


EXAMPLES = [
    [
        "examples/rabbit_bw.mp4",
        "a small wild rabbit sitting among rough textured boulders with a fallen log and dry grass behind it, in soft natural daylight",
        "a young brown cottontail rabbit with warm tan and grey-brown fur, a pale cream underside and soft pink inner ears, perched on weathered grey granite boulders flecked with green and ochre lichen; golden dry grass and muted green vegetation in warm late-afternoon light",
        "960×544 (training bucket / strongest color)", 121, 42, False, False, "Colorize only",
    ],
    [
        "examples/old_city_bw.mp4",
        "an old black-and-white city street scene with pedestrians, stone buildings, shop windows, and soft daylight",
        "natural archival colors: warm beige stone façades, dark wool coats, muted brown and navy clothing, pale skin tones, brass and wood storefront details, and a slightly cool daylight sky",
        "960×544 (training bucket / strongest color)", 121, 42, False, False, "Colorize + archival restoration look",
    ],
]
# ====================================================================================

FPS = 24.0
MAX_SEED = np.iinfo(np.int32).max
HF_TOKEN = os.environ.get("HF_TOKEN")
LTX_MODEL_REPO = "Lightricks/LTX-2.3"
GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized"


def _src_fps(path, default=FPS):
    try:
        return float(iio.immeta(path, plugin="pyav").get("fps", default)) or default
    except Exception:
        return default


def _prep_reference(path, width, height, num_frames, force_grayscale=False):
    """Resample to 24fps, aspect-fit/crop to WxH, NF frames; optionally force grayscale; write temp mp4."""
    vid = iio.imread(path, plugin="pyav")
    src_fps = _src_fps(path)
    n = len(vid)
    out = []
    for i in range(num_frames):
        idx = min(int(round(i / FPS * src_fps)), n - 1)
        im = Image.fromarray(vid[idx]).convert("RGB")
        im = ImageOps.fit(im, (width, height), Image.LANCZOS)
        if force_grayscale:
            im = im.convert("L").convert("RGB")
        out.append(np.array(im))
    tmp = tempfile.mktemp(suffix=".mp4")
    iio.imwrite(tmp, np.stack(out), fps=FPS, plugin="pyav", codec="libx264")
    return tmp


def _pick_resolution(path, preset):
    w, h = RES_PRESETS[preset]
    try:
        f0 = iio.imread(path, plugin="pyav", index=0)
        if f0.shape[0] > f0.shape[1]:  # portrait
            w, h = h, w
    except Exception:
        pass
    return w, h


# --- Load native pipeline + IC-LoRA once at module scope (ZeroGPU packs weights here) ---
print("Downloading checkpoints…")
checkpoint_path = hf_hub_download(LTX_MODEL_REPO, "ltx-2.3-22b-distilled-1.1.safetensors", token=HF_TOKEN)
spatial_upsampler_path = hf_hub_download(LTX_MODEL_REPO, "ltx-2.3-spatial-upscaler-x2-1.1.safetensors", token=HF_TOKEN)
gemma_root = snapshot_download(GEMMA_REPO, token=HF_TOKEN)
lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=HF_TOKEN)

print("Building ICLoraPipeline…")
pipeline = ICLoraPipeline(
    distilled_checkpoint_path=checkpoint_path,
    spatial_upsampler_path=spatial_upsampler_path,
    gemma_root=gemma_root,
    loras=[LoraPathStrengthAndSDOps(lora_path, LORA_SCALE, LTXV_LORA_COMFY_RENAMING_MAP)],
    # bf16 (NOT fp8): the IC-LoRA is fused into the transformer at MODULE SCOPE (the GPU
    # worker can't re-open the checkpoint file). fp8_cast()'s fusion runs a custom CUDA kernel
    # that can't be ZeroGPU-virtualised; the bf16 fuse rule is pure torch -> virtualisable.
    quantization=None,
)


def _preload_pin(ledger, tag):
    if ledger is None:
        return
    for name in ["transformer", "video_encoder", "video_decoder", "audio_encoder",
                 "audio_decoder", "vocoder", "spatial_upsampler", "text_encoder",
                 "gemma_embeddings_processor"]:
        fn = getattr(ledger, name, None)
        if callable(fn):
            try:
                obj = fn()
                setattr(ledger, name, (lambda o=obj: o))
                print(f"[preload {tag}] {name} ✓")
            except Exception as e:
                print(f"[preload {tag}] {name} skipped: {e}")


# Preload stage 1 always; preload stage 2 only when two-stage is used (skip_stage_2=False).
# Eagerly pinning both ledgers materializes TWO ~46GB transformers — too big for the ZeroGPU pack.
_preload_pin(getattr(pipeline, "stage_1_model_ledger", None), "stage1")
if not SKIP_STAGE_2:
    _preload_pin(getattr(pipeline, "stage_2_model_ledger", None), "stage2")
print("Pipeline ready.")

# ============================ AOTI (native bf16 transformer graph) ============================
AOTI_REPO = os.environ.get("AOTI_REPO", "linoyts/LTX-2.3-Native-Transformer-GroupA-sm120-cu130-r20")
import types as _types
from dataclasses import replace as _dc_replace
from ltx_core.model.transformer.transformer_args import TransformerArgs as _TA
_TA_FIELDS = list(_TA.__dataclass_fields__.keys())
def _flatten_ta(ta):
    out = []
    for f in _TA_FIELDS:
        v = getattr(ta, f)
        if torch.is_tensor(v):
            out.append(v)
        elif isinstance(v, tuple) and len(v) > 0 and all(torch.is_tensor(x) for x in v):
            out.extend(v)
    return out
def _install_aoti():
    velocity = pipeline.stage_1_model_ledger.transformer().velocity_model
    spaces.aoti_load(module=velocity, repo_id=AOTI_REPO)
    def _proc(self, video, audio, perturbations):
        for blk in self.transformer_blocks:
            o = blk(*(_flatten_ta(video) + _flatten_ta(audio)))
            video = _dc_replace(video, x=o[0]); audio = _dc_replace(audio, x=o[1])
        return video, audio
    velocity._process_transformer_blocks = _types.MethodType(_proc, velocity)
    print(f"[AOTI] loaded {AOTI_REPO} + patched block loop", flush=True)
# For IC-LoRA task adapters, keep AOTI OFF by default. AOTI can replace the
# transformer block execution with a precompiled graph that may not include the
# currently fused IC-LoRA weights. Symptom: the output follows the reference
# weakly / changes motion but stays grayscale, i.e. the Colorization LoRA is
# effectively bypassed. Set USE_AOTI=1 only after verifying the LoRA effect.
USE_AOTI = os.environ.get("USE_AOTI", "0") == "1"
print(f"[AOTI] base torch={torch.__version__} cuda={torch.version.cuda} enabled={USE_AOTI}", flush=True)
if USE_AOTI:
    try:
        _install_aoti(); print("[AOTI] OK", flush=True)
    except Exception as _e:
        import traceback; traceback.print_exc(); print(f"[AOTI] FAILED ({_e!r}) -> EAGER", flush=True)
else:
    print("[AOTI] disabled for Colorization IC-LoRA; using eager transformer so fused LoRA weights are active", flush=True)
# ==============================================================================================


def _duration(*args, **kwargs):
    nf = next((a for a in args if isinstance(a, int) and a in FRAME_CHOICES), DEFAULT_FRAMES)
    return int(60 + nf * 1.2)


@spaces.GPU(duration=_duration)
@torch.inference_mode()
def colorize_video(video, scene_description, color_description, preset, num_frames, seed, randomize,
                   force_grayscale, restoration_mode, progress=gr.Progress(track_tqdm=True)):
    if video is None:
        raise gr.Error("Please upload a grayscale / black-and-white / desaturated video.")
    if not color_description.strip():
        raise gr.Error("Describe the desired natural colors for the scene.")

    seed = random.randint(0, MAX_SEED) if randomize else int(seed)
    num_frames = int(num_frames)
    width, height = _pick_resolution(video, preset)
    ref_path = _prep_reference(video, width, height, num_frames, force_grayscale=force_grayscale)
    tiling = TilingConfig.default()

    # skip_stage_2 outputs at half the passed dims -> pass 2x so output matches the preset.
    # This is also the identity-safe colorization recipe recommended for this LoRA.
    gen_w, gen_h = (width * 2, height * 2) if SKIP_STAGE_2 else (width, height)

    final_prompt = build_prompt(scene_description, color_description, restoration_mode)
    print("[PROMPT]", final_prompt, flush=True)

    video_out, audio_out = pipeline(
        prompt=final_prompt,
        seed=seed, height=gen_h, width=gen_w,
        num_frames=num_frames, frame_rate=FPS,
        images=[], video_conditioning=[(ref_path, 1.0)],
        skip_stage_2=SKIP_STAGE_2, tiling_config=tiling,
    )
    out_path = tempfile.mktemp(suffix=".mp4")
    encode_video(video=video_out, fps=FPS, audio=audio_out, output_path=out_path,
                 video_chunks_number=get_video_chunks_number(num_frames, tiling))
    return out_path, seed


with gr.Blocks(title=TITLE) as demo:
    gr.Markdown(
        "# 🎞️ LTX-2.3 Colorization\n"
        "Colorize black-and-white, monochrome, or desaturated clips while preserving subject identity, "
        "framing, motion, and scene geometry. Uses [LTX-2.3 Distilled](https://huggingface.co/Lightricks/LTX-2.3) "
        "with the [Colorization IC-LoRA](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Colorization)."
    )
    gr.Markdown(
        "⚙️ **AOTI is disabled by default for this Colorization version** so the fused IC-LoRA weights are not bypassed. Set `USE_AOTI=1` only after testing.  \n"
        "Note: the Colorization LoRA is not a true deblur/denoise/decompression/upscaling model. The old-film and high-detail options below are prompt hints only; "
        "for heavy restoration, use a dedicated deblur/decompression pass before or after colorization."
    )

    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Black-and-white / desaturated input video")
            scene_description = gr.Textbox(
                label="Reference scene description",
                lines=3,
                placeholder="e.g. an old black-and-white city street scene with pedestrians, stone buildings, shop windows, and soft daylight",
            )
            color_description = gr.Textbox(
                label="Desired natural colors",
                lines=4,
                placeholder="e.g. warm beige stone façades, dark wool coats, muted brown and navy clothing, natural skin tones, brass storefront details, and a cool daylight sky",
            )
            with gr.Accordion("Settings", open=False):
                preset = gr.Dropdown(list(RES_PRESETS), value=DEFAULT_PRESET, label="Resolution")
                num_frames = gr.Dropdown(FRAME_CHOICES, value=DEFAULT_FRAMES, label="Frames (24fps; best: 121)")
                restoration_mode = gr.Dropdown(RESTORATION_MODES, value=DEFAULT_RESTORATION_MODE, label="Optional restoration / upscale hint")
                force_grayscale = gr.Checkbox(False, label="Force input to grayscale before conditioning")
                randomize = gr.Checkbox(True, label="Randomize seed")
                seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
            run = gr.Button("Colorize", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Colorized result")

    run.click(
        colorize_video,
        inputs=[video_in, scene_description, color_description, preset, num_frames, seed, randomize, force_grayscale, restoration_mode],
        outputs=[video_out, seed],
    )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[video_in, scene_description, color_description, preset, num_frames, seed, randomize, force_grayscale, restoration_mode],
        outputs=[video_out, seed],
        fn=colorize_video,
        cache_examples=True,
        cache_mode="lazy",
    )


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