import sys
import os
import random
from pathlib import Path


# Add vendored packages to Python path (upstream Lightricks/LTX-2, supports 2.3)
current_dir = Path(__file__).parent
sys.path.insert(0, str(current_dir / "packages" / "ltx-pipelines" / "src"))
sys.path.insert(0, str(current_dir / "packages" / "ltx-core" / "src"))

import spaces  # MUST come before torch / any CUDA-touching import (ZeroGPU hook)
import gradio as gr
import numpy as np
import torch
from huggingface_hub import hf_hub_download, snapshot_download

from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import OffloadMode
from ltx_core.loader import LoraPathStrengthAndSDOps, LTXV_LORA_COMFY_RENAMING_MAP
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy

MAX_SEED = np.iinfo(np.int32).max
DEFAULT_SEED = 42
APP_VERSION = "2026-07-09-rope-ref-v2"
# Two-stage distilled: stage 1 at height/2 x width/2, stage 2 upsampled 2x to
# height x width. 768 matches the ComfyUI workflow's base resolution (÷64 ok).
DEFAULT_HEIGHT = 768
DEFAULT_WIDTH = 768

DEFAULT_PROMPT = (
    "ref_t2v: A person with distinctive facial features looking directly at the camera with a confident expression. "
    "They blink naturally, turn their head slightly, and smile subtly. Soft natural lighting illuminates "
    "their face. Realistic skin texture, detailed facial features, cinematic quality."
)

# ---------------------------------------------------------------------------
# Model sources — LTX-2.3, the base the Best-Face-ID LoRA was trained on.
# ---------------------------------------------------------------------------
BASE_REPO_ID = "Lightricks/LTX-2.3"
CHECKPOINT_FILENAME = "ltx-2.3-22b-dev.safetensors"  # bundled: transformer+VAE+audio_vae+vocoder (bf16; fp8-cast at load)
DISTILLED_LORA_FILENAME = "ltx-2.3-22b-distilled-lora-384-1.1.safetensors"
SPATIAL_UPSAMPLER_FILENAME = "ltx-2.3-spatial-upscaler-x2-1.0.safetensors"

FACE_ID_REPO_ID = "Alissonerdx/LTX-Best-Face-ID"
FACE_ID_LORA_FILENAME = "Best_FaceID_v1.0_LoRA.safetensors"

# Gemma-3 text encoder lives in the LTX-2 repo (LTX-2.3 ships no text encoder).
# We download ONLY the transformers `model-*` shards (not the diffusers
# `diffusion_pytorch_model-*` copy) so the block's rglob("*.safetensors") over
# text_encoder/ finds exactly one set of weights.
GEMMA_REPO_ID = "Lightricks/LTX-2"

print("=" * 80)
print(f"Loading LTX-2.3 Distilled pipeline with Best-Face-ID LoRA... app_version={APP_VERSION}")
print("=" * 80)

checkpoint_path = hf_hub_download(repo_id=BASE_REPO_ID, filename=CHECKPOINT_FILENAME)
distilled_lora_path = hf_hub_download(repo_id=BASE_REPO_ID, filename=DISTILLED_LORA_FILENAME)
spatial_upsampler_path = hf_hub_download(repo_id=BASE_REPO_ID, filename=SPATIAL_UPSAMPLER_FILENAME)
face_id_lora_path = hf_hub_download(repo_id=FACE_ID_REPO_ID, filename=FACE_ID_LORA_FILENAME)
gemma_root = snapshot_download(
    repo_id=GEMMA_REPO_ID,
    allow_patterns=[
        "text_encoder/model*.safetensors",
        "text_encoder/model.safetensors.index.json",
        "text_encoder/config.json",
        "text_encoder/generation_config.json",
        "tokenizer/*",
    ],
)

print(f"  checkpoint_path={checkpoint_path}")
print(f"  distilled_lora_path={distilled_lora_path}")
print(f"  spatial_upsampler_path={spatial_upsampler_path}")
print(f"  face_id_lora_path={face_id_lora_path}")
print(f"  gemma_root={gemma_root}")

# distilled LoRA + Face-ID LoRA, both fused into the transformer at build time.
# NOTE: the ComfyUI workflow used 0.6 for a *different* community distilled LoRA
# (dynamic_fro09_avg_rank_111). 1.0 is the standard strength for this official
# rank-384 distilled LoRA; revisit if the 8-step schedule looks under/over-cooked.
loras = [
    LoraPathStrengthAndSDOps(path=distilled_lora_path, strength=1.0, sd_ops=LTXV_LORA_COMFY_RENAMING_MAP),
    LoraPathStrengthAndSDOps(path=face_id_lora_path, strength=1.0, sd_ops=LTXV_LORA_COMFY_RENAMING_MAP),
]

# Run the 22B transformer in BF16 (no quantization). fp8-CAST upcasts each weight
# to bf16 on every forward (memory churn on the MIG slice); fp8-SCALED-MM would
# avoid it but needs a pre-quantized fp8 checkpoint (2.3-dev is bf16). bf16 = ~44 GB
# resident, no per-forward churn.
quantization_policy = None

# OffloadMode.NONE = native ZeroGPU pack: build all models at module scope, let the
# `.to("cuda")` hook pack them, and stream them RESIDENT into each fork (fast — warm
# gens ~20-28 s, no per-generation weight loading).
#
# The ZeroGPU-MIG allocator gotcha this navigates: the native CUDA allocator throws
# an NVML assert only when a cudaMalloc *fails* (its failure handler's cudaMemGetInfo
# dies on MIG). ZeroGPU caps the fork's memory fraction low (~56 GB) by default, so a
# ~48 GB resident + activations overflows it -> failed malloc -> assert. The fix is
# NOT offload — it's raising the fraction to 1.0 (95 GB) in generate_video so mallocs
# succeed. The native allocator REUSES freed blocks, so real usage stays ~48-68 GB and
# the full 1-5 s range fits. (Do NOT use cudaMallocAsync: it hoards freed blocks to the
# cap; do NOT use expandable_segments: its cuMemMap pokes NVML on MIG -> assert.)
OFFLOAD_MODE = OffloadMode.NONE
pipeline = DistilledPipeline(
    distilled_checkpoint_path=checkpoint_path,
    gemma_root=gemma_root,
    spatial_upsampler_path=spatial_upsampler_path,
    loras=loras,
    quantization=quantization_policy,
    offload_mode=OFFLOAD_MODE,
)

print("=" * 80)
print("Pipeline constructed (weights not loaded yet)")
print("=" * 80)


def preload_models_for_zerogpu() -> None:
    """Build all block models at MODULE SCOPE so ZeroGPU can preload them.

    Must run in the main process, NOT inside a ``@spaces.GPU`` function.
    ``import spaces`` installs a torch hook that intercepts ``.to("cuda")`` in the
    main process: it records each model's weights, packs them to disk at launch,
    and streams them into VRAM inside every ``@spaces.GPU`` fork. Our block-caching
    patch keeps the built models resident, so ``pipeline.warmup()`` here builds +
    hijacks them exactly once (rather than the upstream build-use-free-per-call).
    """
    if OFFLOAD_MODE != OffloadMode.NONE:
        # CPU/DISK offload: do NOT pack at module scope. The blocks all read from
        # the SAME bundled checkpoint file; packing the small VAE/vocoder blocks
        # makes `spaces` delete that checkpoint blob after packing, which then
        # breaks the transformer's layer-by-layer streaming (FileNotFoundError).
        # With offload, the blocks build lazily inside the @spaces.GPU fork and
        # the checkpoint stays on disk for streaming.
        print("CPU offload: skipping module-scope pack (blocks build in-fork; checkpoint kept on disk)")
        return
    print("Preloading models at module scope for ZeroGPU (.to('cuda') hijack)...")
    pipeline.warmup()
    print("Models preloaded: weights registered for ZeroGPU streaming.")


# Aspect ratios offered in "new scene" mode -> exact (width, height), each a
# multiple of 64 (so stage-1 at half resolution stays a multiple of 32 for the
# patchifier), ~0.44-0.59 MP.
ASPECT_DIMS = {
    "1:1": (768, 768),
    "4:3": (768, 576),
    "3:4": (576, 768),
    "16:9": (1024, 576),
    "9:16": (576, 1024),
}
DEFAULT_ASPECT = "4:3"


def _dims_for_ratio(ratio: float, area: int = 768 * 768) -> tuple[int, int]:
    """Return (width, height) for a target aspect ratio, ~`area` pixels, multiple of 64."""
    import math

    h = math.sqrt(area / ratio)
    w = h * ratio

    def r64(x: float) -> int:
        return max(384, min(1024, int(round(x / 64.0) * 64)))

    return r64(w), r64(h)


# size="xlarge" = full 96 GB card (2x quota). Needed: the bf16 22B transformer
# (~44 GB) + VAEs/vocoder are held resident, and with the memory fraction raised to
# 1.0 the diffusion activations for long (121-frame) clips need the full slice.
@spaces.GPU(duration=300, size="xlarge")
def generate_video(
    face_image,
    prompt: str,
    duration: float,
    aspect_ratio: str,
    seed: int,
    randomize_seed: bool,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a short identity-preserving video from a reference face photo.

    Runs the official LTX-2.3 distilled pipeline with the Best-Face-ID LoRA. The
    reference face is injected as a source-phase-tagged overlap reference (the
    ComfyUI BFS mechanism) so the generated video keeps that person's identity.
    Returns a path to an .mp4 file.

    Args:
        face_image (str): Reference face photo (filepath, URL, or base64 image).
            A clean, frontal, well-lit close-up cropped to the face / upper body
            gives the best identity transfer. Required.
        prompt (str): What the person is doing and the scene, described in the
            present tense. Including identity attributes (hair, skin tone, eyes,
            facial hair, glasses, face shape) noticeably improves fidelity.
        duration (float): Video length in seconds, from 1.0 to 8.0.
        seed (int): Random seed, 0 to 2147483647. Ignored when randomize_seed=True.
        randomize_seed (bool): When true, a fresh random seed is used each run.
        aspect_ratio (str): Output aspect ratio — one of "1:1", "4:3", "3:4",
            "16:9", "9:16".

    Returns:
        tuple[str, int]: Filepath of the generated .mp4 video, and the seed used.
    """
    current_seed = int(seed)
    try:
        if face_image is None:
            raise ValueError("Please provide a reference face image.")

        current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

        frame_rate = 24.0
        # LTX needs num_frames == 8k + 1 (one latent temporal patch per 8 frames).
        num_frames = (int(duration * frame_rate) // 8) * 8 + 1

        output_dir = Path("outputs")
        output_dir.mkdir(exist_ok=True)
        output_path = output_dir / f"video_{current_seed}.mp4"

        # Persist the uploaded image to a path (pipeline loads reference from disk).
        temp_image_path = output_dir / f"temp_input_{current_seed}.jpg"
        if hasattr(face_image, "save"):
            face_image.save(temp_image_path)
        else:
            temp_image_path = Path(face_image)

        width, height = ASPECT_DIMS.get(aspect_ratio, ASPECT_DIMS[DEFAULT_ASPECT])

        # Best-Face-ID was not trained as I2V. The face image must not replace
        # frame 0; it is injected below as source-phase-tagged overlap tokens in
        # a separate RoPE source segment.
        images = []

        # `ref_t2v:` is the LoRA's trained caption trigger.
        identity_prompt = prompt if prompt.strip().startswith("ref_t2v:") else f"ref_t2v: {prompt}"
        print(f"Generating (reference-overlap, {int(width)}x{int(height)}, {num_frames}f): {identity_prompt}")

        # --- ZeroGPU diagnostics (safe counters only; mem_get_info() asserts on MIG) ---
        try:
            _p = torch.cuda.get_device_properties(0)
            _backend = getattr(torch.cuda.memory, "get_allocator_backend", lambda: "?")()
            print(
                f"[GPU] {_p.name} total={_p.total_memory/1e9:.1f}GB "
                f"reserved={torch.cuda.memory_reserved()/1e9:.1f}GB "
                f"allocated={torch.cuda.memory_allocated()/1e9:.1f}GB "
                f"backend={_backend} alloc_conf={os.environ.get('PYTORCH_CUDA_ALLOC_CONF')}"
            )
        except Exception as _e:
            print(f"[GPU] diag failed: {_e}")

        # CRITICAL: ZeroGPU caps the fork's per-process memory fraction well below the
        # full MIG slice (~56 GB). Raising it to 100% (~95 GB) is what lets the native
        # allocator satisfy diffusion activations without a failed malloc — and a
        # failed malloc is what triggers the NVML assert on this MIG slice. With this,
        # the plain native allocator (reuses freed blocks, no hoarding) fits the whole
        # 1-5 s range on the resident-packed transformer. Do NOT remove this.
        try:
            torch.cuda.set_per_process_memory_fraction(1.0, torch.cuda.current_device())
        except Exception as _e:
            print(f"[GPU] set_per_process_memory_fraction failed: {_e}")

        # Run under inference_mode (matches upstream's @torch.inference_mode() main).
        # Without it, the block-streaming path (which creates inference tensors) mixes
        # with autograd tracking -> "Inference tensors cannot be saved for backward".
        tiling_config = TilingConfig.default()
        with torch.inference_mode():
            video, audio = pipeline(
                prompt=identity_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=True,
                enhance_prompt_image_path=str(temp_image_path),
                reference_image_path=str(temp_image_path),
                reference_source_id=2.0,
                reference_phase_scale=1.0,
            )

            encode_video(
                video=video,
                fps=int(frame_rate),
                audio=audio,
                output_path=str(output_path),
                video_chunks_number=get_video_chunks_number(num_frames, tiling_config),
            )

        if temp_image_path.exists() and temp_image_path.parent == output_dir:
            temp_image_path.unlink()

        return str(output_path), current_seed

    except Exception as e:
        import traceback

        print(f"Error: {e}\n{traceback.format_exc()}")
        return None, current_seed


with gr.Blocks(title=f"LTX-2.3 Face ID 🎥🧑 | {APP_VERSION}") as demo:
    gr.Markdown("# LTX-2.3 Face ID 🎥🧑: Identity-Preserving Video Generation")
    gr.Markdown(
        f"**Version:** `{APP_VERSION}` · **Trigger:** `ref_t2v:` · **Reference path:** RoPE source-phase overlap, not I2V.\n\n"
        "Generate identity-preserving videos using [Lightricks LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3) "
        "with the [BFS Best-Face-ID LoRA](https://huggingface.co/Alissonerdx/LTX-Best-Face-ID). "
        "Upload a clear, frontal face photo and describe the desired video action.\n\n"
        "> This demo uses the **distilled** model for speed. For maximum quality, run the "
        "**full model** locally with the [ComfyUI BFS Nodes](https://github.com/alisson-anjos/ComfyUI-BFSNodes)."
    )

    with gr.Row():
        with gr.Column():
            face_image = gr.Image(label="Reference Face Image", type="pil")

            prompt = gr.Textbox(
                label="Prompt",
                value=DEFAULT_PROMPT,
                lines=3,
                placeholder="ref_t2v: A person ...",
            )

            duration = gr.Slider(label="Duration (seconds)", minimum=1.0, maximum=8.0, value=3.0, step=0.5)

            generate_btn = gr.Button("Generate Video", variant="primary", size="lg")

            with gr.Accordion("Advanced Settings", open=False):
                aspect_ratio = gr.Dropdown(
                    label="Aspect ratio",
                    choices=list(ASPECT_DIMS.keys()),
                    value=DEFAULT_ASPECT,
                )
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=DEFAULT_SEED, step=1)
                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)

        with gr.Column():
            output_video = gr.Video(label="Generated Video", autoplay=True)

    generate_btn.click(
        fn=generate_video,
        inputs=[face_image, prompt, duration, aspect_ratio, seed, randomize_seed],
        outputs=[output_video, seed],
    )

    gr.Examples(
        examples=[
            [
                str(current_dir / "face_ex1.jpg"),
                "ref_t2v: A man sits in a sunlit cafe, sipping coffee and smiling warmly at the camera. Soft window light, shallow depth of field, cinematic quality, realistic skin texture.",
                3.0,
                DEFAULT_ASPECT,
                DEFAULT_SEED,
                True,
            ],
            [
                str(current_dir / "face_ex2.jpg"),
                "ref_t2v: A woman walks through a vibrant open-air market, glancing at the camera with a gentle smile as stalls of colorful fabrics pass behind her. Natural daylight, shallow depth of field, cinematic.",
                3.0,
                DEFAULT_ASPECT,
                DEFAULT_SEED,
                True,
            ],
            [
                str(current_dir / "face_ex3.jpg"),
                "ref_t2v: A woman stands on a rooftop at golden hour, the city skyline behind her, turning to look at the camera with a calm expression. Warm cinematic lighting.",
                3.0,
                DEFAULT_ASPECT,
                DEFAULT_SEED,
                True,
            ],
        ],
        fn=generate_video,
        inputs=[face_image, prompt, duration, aspect_ratio, seed, randomize_seed],
        outputs=[output_video, seed],
        label="Example",
        cache_examples=False,
    )


css = """
.gradio-container .contain{max-width: 1200px !important; margin: 0 auto !important}
"""

# Build the models now, in the main process, so ZeroGPU packs + streams them.
# Direct module-scope call (NOT @spaces.GPU, NOT demo.load) so the `.to("cuda")`
# hijack fires once for the whole app, before demo.launch() triggers packing.
preload_models_for_zerogpu()

if __name__ == "__main__":
    # mcp_server=True exposes generate_video as an MCP tool (docstring Args become
    # the tool schema) at /gradio_api/mcp/, alongside the normal UI.
    demo.launch(theme=gr.themes.Citrus(), css=css, mcp_server=True)