app.py

import json
import os
import random
import shutil
import tempfile
import threading
import time
import zipfile
from dataclasses import dataclass
from pathlib import Path

# Keep persistent files in the bucket and stage them locally before any model is memory-mapped.
BUCKET_ROOT = Path(os.environ.get("LTX_BUCKET_ROOT", "/data"))
PERSISTENT_MODEL_ROOT = Path(os.environ.get("LTX_MODEL_CACHE_ROOT", BUCKET_ROOT / "ltx-model-cache"))
RUNTIME_MODEL_ROOT = Path(os.environ.get("LTX_RUNTIME_MODEL_ROOT", "/tmp/ltx-model"))
RUNTIME_HF_CACHE_ROOT = Path(os.environ.get("LTX_RUNTIME_HF_CACHE_ROOT", "/tmp/huggingface"))
PERSISTENT_MODEL_ROOT.mkdir(parents=True, exist_ok=True)
RUNTIME_HF_CACHE_ROOT.mkdir(parents=True, exist_ok=True)
os.environ["HF_HOME"] = str(RUNTIME_HF_CACHE_ROOT)
os.environ["HF_HUB_CACHE"] = str(RUNTIME_HF_CACHE_ROOT / "hub")
os.environ["HF_ASSETS_CACHE"] = str(RUNTIME_HF_CACHE_ROOT / "assets")
os.environ["HF_XET_CACHE"] = str(RUNTIME_HF_CACHE_ROOT / "xet")
# Xet's optional chunk and shard caches duplicate large files already staged for one-time loading.
os.environ["HF_XET_CHUNK_CACHE_SIZE_BYTES"] = "0"
os.environ["HF_XET_SHARD_CACHE_SIZE_LIMIT"] = "0"

# ZeroGPU does not support torch.compile/dynamo for this workload, so disable it before torch import.
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import gradio as gr
import imageio.v3 as iio
import numpy as np
import spaces
import torch
from diffusers import LTX2InContextPipeline
from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES
from diffusers.utils import encode_video, load_video
from huggingface_hub import hf_hub_download, snapshot_download
from huggingface_hub.utils import GatedRepoError, HfHubHTTPError
from PIL import Image, ImageOps
from safetensors.torch import load_file
from transformers import Gemma3ForConditionalGeneration, GemmaTokenizerFast


# Shared model and storage configuration for the whole workflow.
BASE_MODEL = "diffusers/LTX-2.3-Distilled-Diffusers"
FPS = 24
NUM_STEPS = len(DISTILLED_SIGMA_VALUES)
MAX_SEED = np.iinfo(np.int32).max
RUNS_DIR = BUCKET_ROOT / "ltxvideo-renders"
INDEX_PATH = RUNS_DIR / "index.json"
ZIP_PATH = RUNS_DIR / "all-renders.zip"
TOKEN_ENV_NAMES = ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGINGFACE_TOKEN")
PERSISTENT_BASE_FILES = (
    "text_encoder/config.json",
    "text_encoder/generation_config.json",
    "text_encoder/model.safetensors.index.json",
    *(f"text_encoder/model-{index:05d}-of-00011.safetensors" for index in range(1, 12)),
    "tokenizer/added_tokens.json",
    "tokenizer/chat_template.jinja",
    "tokenizer/special_tokens_map.json",
    "tokenizer/tokenizer.json",
    "tokenizer/tokenizer.model",
    "tokenizer/tokenizer_config.json",
)


@dataclass(frozen=True)
class EffectConfig:
    # The registry keeps each button's LoRA, prompt recipe, and preprocessing rules together.
    key: str
    label: str
    output_label: str
    lora_repo: str
    lora_file: str
    reference_downscale_factor: int
    default_resolution: tuple[int, int]
    grayscale_reference: bool
    duration_scale: float
    prompt_prefix: str
    prompt_suffix: str


EFFECTS = {
    "decompress": EffectConfig(
        key="decompress",
        label="Decompress",
        output_label="Restored",
        lora_repo="Lightricks/LTX-2.3-22b-IC-LoRA-Decompression",
        lora_file="ltx-2.3-22b-ic-lora-decompression-0.9.safetensors",
        reference_downscale_factor=1,
        default_resolution=(960, 544),
        grayscale_reference=False,
        duration_scale=1.4,
        prompt_prefix="Reference shows the same scene with compression artifacts, macroblocking, ringing, banding, and chroma bleed.",
        prompt_suffix="Remove compression artifacts while preserving subject identity, framing, motion, geometry, and natural audio.",
    ),
    "deblur": EffectConfig(
        key="deblur",
        label="Deblur",
        output_label="Sharpened",
        lora_repo="Lightricks/LTX-2.3-22b-IC-LoRA-Deblur",
        lora_file="ltx-2.3-22b-ic-lora-deblur-0.9.safetensors",
        reference_downscale_factor=1,
        default_resolution=(960, 544),
        grayscale_reference=False,
        duration_scale=1.4,
        prompt_prefix="Reference shows the same scene heavily out of focus with soft defocused blur and no fine detail.",
        prompt_suffix="Restore sharp focus with crisp detail and clean edges while preserving subject identity, framing, motion, geometry, and natural audio.",
    ),
    "colorize": EffectConfig(
        key="colorize",
        label="Colorize",
        output_label="Colorized",
        lora_repo="Lightricks/LTX-2.3-22b-IC-LoRA-Colorization",
        lora_file="ltx-2.3-22b-ic-lora-colorization-0.9.safetensors",
        reference_downscale_factor=1,
        default_resolution=(960, 544),
        grayscale_reference=True,
        duration_scale=1.4,
        prompt_prefix="Reference shows the same scene in grayscale with no usable color information.",
        prompt_suffix="Restore natural, coherent color while preserving subject identity, framing, motion, geometry, and natural audio.",
    ),
}


def _get_hf_token() -> str | None:
    # Accept the common Hugging Face secret names so a valid token is not missed by spelling.
    for name in TOKEN_ENV_NAMES:
        token = os.environ.get(name)
        if token and token.strip():
            return token.strip()
    return None


def _token_status() -> str:
    # Log only the secret name that is present; never print token contents.
    present = [name for name in TOKEN_ENV_NAMES if os.environ.get(name)]
    return present[0] if present else "none"


def _ensure_persistent_base_files() -> None:
    # Download only missing persistent components; existing large shards are never fetched again.
    for filename in PERSISTENT_BASE_FILES:
        destination = PERSISTENT_MODEL_ROOT / filename
        if destination.is_file() and destination.stat().st_size > 0:
            continue
        print(f"[CACHE] Downloading persistent component: {filename}", flush=True)
        hf_hub_download(
            BASE_MODEL,
            filename,
            token=_get_hf_token(),
            local_dir=PERSISTENT_MODEL_ROOT,
        )


def _stage_persistent_components() -> None:
    # Copy bucket-backed safetensors sequentially so PyTorch later memory-maps only local files.
    _ensure_persistent_base_files()
    if RUNTIME_MODEL_ROOT.exists():
        shutil.rmtree(RUNTIME_MODEL_ROOT)
    RUNTIME_MODEL_ROOT.mkdir(parents=True, exist_ok=True)
    for component in ("text_encoder", "tokenizer"):
        print(f"[CACHE] Staging persistent {component} into local runtime storage", flush=True)
        shutil.copytree(
            PERSISTENT_MODEL_ROOT / component,
            RUNTIME_MODEL_ROOT / component,
            copy_function=shutil.copyfile,
        )


def _download_runtime_components() -> None:
    # Fetch only components not already loaded into RAM after their local staging files are removed.
    print("[CACHE] Downloading non-persistent model components into runtime storage", flush=True)
    snapshot_download(
        BASE_MODEL,
        token=_get_hf_token(),
        local_dir=RUNTIME_MODEL_ROOT,
        ignore_patterns=["text_encoder/*", "tokenizer/*"],
    )


_stage_persistent_components()

# Load the largest component first, then release its local files before downloading the rest.
print("[CACHE] Loading staged text encoder into RAM", flush=True)
text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
    RUNTIME_MODEL_ROOT / "text_encoder",
    dtype=torch.bfloat16,
    local_files_only=True,
)
tokenizer = GemmaTokenizerFast.from_pretrained(
    RUNTIME_MODEL_ROOT / "tokenizer",
    local_files_only=True,
)
shutil.rmtree(RUNTIME_MODEL_ROOT / "text_encoder")
shutil.rmtree(RUNTIME_MODEL_ROOT / "tokenizer")
print("[CACHE] Released staged text encoder files from runtime storage", flush=True)
_download_runtime_components()

# The shared Diffusers engine is loaded once; individual effect LoRAs are fused in and out per click.
pipe = LTX2InContextPipeline.from_pretrained(
    RUNTIME_MODEL_ROOT,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    torch_dtype=torch.bfloat16,
    local_files_only=True,
)
# Keep inactive pipeline components in system RAM so the 22B model fits within a 16 GB T4.
pipe.enable_model_cpu_offload()
pipe.vae.enable_tiling()

# LoRA files are downloaded lazily so one gated adapter cannot stop the whole UI from loading.
LORA_PATHS = {}

# A lock protects global LoRA fusion state because the one shared engine is mutable.
PIPE_LOCK = threading.Lock()
FUSED_EFFECT = None


def _get_lora_path(config: EffectConfig) -> str:
    # Download the requested adapter only when its button is used, with a clear gated-token error.
    if config.key in LORA_PATHS:
        return LORA_PATHS[config.key]
    token = _get_hf_token()
    if token is None:
        raise gr.Error(
            f"{config.label} needs gated model access, but the Space runtime has no Hugging Face token. "
            "Add a Space secret named HF_TOKEN with a read token that has accepted this Lightricks model."
        )
    try:
        # Persist each LoRA as an ordinary file, then stage it locally before safetensors reads it.
        persistent_dir = PERSISTENT_MODEL_ROOT / "loras" / config.key
        persistent_path = persistent_dir / config.lora_file
        if not persistent_path.is_file():
            persistent_path = Path(
                hf_hub_download(
                    config.lora_repo,
                    config.lora_file,
                    token=token,
                    local_dir=persistent_dir,
                )
            )
        runtime_dir = Path("/tmp/ltx-loras") / config.key
        runtime_dir.mkdir(parents=True, exist_ok=True)
        path = runtime_dir / config.lora_file
        if not path.is_file() or path.stat().st_size != persistent_path.stat().st_size:
            shutil.copyfile(persistent_path, path)
    except GatedRepoError as exc:
        raise gr.Error(
            f"{config.label} could not access {config.lora_repo}. "
            f"The Space sees token secret '{_token_status()}', but that token is not authorized for this gated repo. "
            "Use the same account/token that shows 'access granted' on the model page, and make sure the token has read access."
        ) from exc
    except HfHubHTTPError as exc:
        raise gr.Error(f"{config.label} LoRA download failed from {config.lora_repo}: {exc}") from exc
    LORA_PATHS[config.key] = str(path)
    return str(path)


def _ensure_storage() -> None:
    # The mounted bucket is the durable place for render history and zip artifacts.
    RUNS_DIR.mkdir(parents=True, exist_ok=True)
    if not INDEX_PATH.exists():
        INDEX_PATH.write_text("[]", encoding="utf-8")


def _read_index() -> list[dict]:
    # Invalid index JSON should not strand existing videos; start a clean list instead.
    _ensure_storage()
    try:
        return json.loads(INDEX_PATH.read_text(encoding="utf-8"))
    except Exception:
        return []


def _write_index(records: list[dict]) -> None:
    # Atomic replace keeps the render list readable if the Space is interrupted mid-write.
    _ensure_storage()
    tmp_path = INDEX_PATH.with_suffix(".tmp")
    tmp_path.write_text(json.dumps(records, indent=2), encoding="utf-8")
    tmp_path.replace(INDEX_PATH)


def _render_table(records: list[dict]) -> list[list[str]]:
    # Gradio Dataframe expects simple rows, so keep the persistent JSON richer than the UI.
    return [
        [
            str(i + 1),
            record.get("effect", ""),
            record.get("frames", ""),
            record.get("seed", ""),
            record.get("created_at", ""),
            record.get("file", ""),
        ]
        for i, record in enumerate(records)
    ]


def refresh_history() -> tuple[list[list[str]], str | None]:
    # Refresh is used by the manual button and after every completed render.
    records = _read_index()
    zip_path = str(ZIP_PATH) if ZIP_PATH.exists() else None
    return _render_table(records), zip_path


def _src_fps(path: str, default: int = FPS) -> float:
    # Read source FPS when available so long clips sample frames at the right cadence.
    try:
        return float(iio.immeta(path, plugin="pyav").get("fps", default)) or default
    except Exception:
        return float(default)


def _pick_resolution(path: str, config: EffectConfig) -> tuple[int, int]:
    # Preserve portrait orientation by swapping the configured landscape dimensions.
    width, height = config.default_resolution
    try:
        first = iio.imread(path, plugin="pyav", index=0)
        if first.shape[0] > first.shape[1]:
            width, height = height, width
    except Exception:
        pass
    return width, height


def _ltx_frame_count(source_frames: int, source_fps: float) -> int:
    # Preserve the full source duration at 24 fps while satisfying LTX's required 8k+1 frame shape.
    frames_at_output_fps = max(1, round(source_frames / source_fps * FPS))
    return max(1, round((frames_at_output_fps - 1) / 8) * 8 + 1)


def _prepare_reference(path: str, width: int, height: int, grayscale: bool) -> tuple[list[Image.Image], int]:
    # Load the complete video and resample its full duration for the Diffusers in-context condition.
    frames = load_video(path)
    if not frames:
        raise gr.Error("Could not read any frames from that video.")
    src_fps = _src_fps(path)
    num_frames = _ltx_frame_count(len(frames), src_fps)
    prepared = []
    for frame_index in range(num_frames):
        source_index = min(int(round(frame_index / FPS * src_fps)), len(frames) - 1)
        frame = ImageOps.fit(frames[source_index].convert("RGB"), (width, height), Image.LANCZOS)
        if grayscale:
            frame = frame.convert("L").convert("RGB")
        prepared.append(frame)
    return prepared, num_frames


def _build_prompt(config: EffectConfig, prompt: str) -> str:
    # The prompt keeps the LoRA task explicit while still honoring the user's scene description.
    scene = prompt.strip() or "the same scene"
    return f"{config.prompt_prefix} {config.label.upper()} {scene}. {config.prompt_suffix}"


def _prepare_effect(config: EffectConfig) -> None:
    # Swap the active LoRA by subtracting the old fused adapter and adding the requested adapter.
    global FUSED_EFFECT
    if FUSED_EFFECT == config.key:
        return
    if FUSED_EFFECT is not None:
        old_path = _get_lora_path(EFFECTS[FUSED_EFFECT])
        pipe.load_lora_weights(load_file(old_path), adapter_name="effect")
        pipe.fuse_lora(lora_scale=-1.0)
        pipe.unload_lora_weights()
    pipe.load_lora_weights(load_file(_get_lora_path(config)), adapter_name="effect")
    pipe.fuse_lora(lora_scale=1.0)
    pipe.unload_lora_weights()
    FUSED_EFFECT = config.key


def _save_iteration(temp_video: str, config: EffectConfig, seed: int, num_frames: int, prompt: str) -> str:
    # Each successful render is copied into the bucket using a stable sequential filename.
    records = _read_index()
    sequence = len(records) + 1
    filename = f"{sequence:04d}-{config.key}.mp4"
    final_path = RUNS_DIR / filename
    shutil.copy2(temp_video, final_path)
    records.append(
        {
            "sequence": sequence,
            "effect": config.label,
            "frames": int(num_frames),
            "seed": int(seed),
            "prompt": prompt,
            "file": str(final_path),
            "created_at": time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime()),
        }
    )
    _write_index(records)
    _make_zip(records)
    return str(final_path)


def _make_zip(records: list[dict]) -> None:
    # Rebuild the download-all archive after each render so the File control is immediately useful.
    _ensure_storage()
    tmp_zip = ZIP_PATH.with_suffix(".tmp")
    with zipfile.ZipFile(tmp_zip, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        zf.writestr("index.json", json.dumps(records, indent=2))
        for record in records:
            video_path = Path(record.get("file", ""))
            if video_path.exists():
                zf.write(video_path, arcname=video_path.name)
    tmp_zip.replace(ZIP_PATH)


def _duration(effect_key: str, video, prompt, seed, randomize) -> int:
    # ZeroGPU needs an estimate, so derive it from the uploaded video's complete duration.
    config = EFFECTS.get(effect_key, EFFECTS["deblur"])
    try:
        metadata = iio.immeta(video, plugin="pyav")
        duration_seconds = float(metadata.get("duration", 0))
        frames = max(1, round(duration_seconds * FPS))
    except Exception:
        frames = 121
    return int(90 + frames * config.duration_scale)


@spaces.GPU(duration=_duration)
@torch.inference_mode()
def run_effect(effect_key: str, video, prompt: str, seed, randomize, progress=gr.Progress(track_tqdm=True)):
    # This is the single execution path behind all effect buttons.
    if video is None:
        raise gr.Error("Please upload a video.")
    config = EFFECTS[effect_key]
    actual_seed = random.randint(0, MAX_SEED) if randomize else int(seed)
    out_width, out_height = _pick_resolution(video, config)
    ref_width = out_width // config.reference_downscale_factor
    ref_height = out_height // config.reference_downscale_factor
    reference_frames, frame_count = _prepare_reference(
        video, ref_width, ref_height, config.grayscale_reference
    )
    full_prompt = _build_prompt(config, prompt)

    def _callback(pipe_obj, step_index, timestep, callback_kwargs):
        # Progress is tied to the distilled sigma schedule length.
        progress((step_index + 1) / NUM_STEPS, desc=f"{config.label} step {step_index + 1}/{NUM_STEPS}")
        return callback_kwargs

    with PIPE_LOCK:
        _prepare_effect(config)
        video_out, audio_out = pipe(
            prompt=full_prompt,
            negative_prompt="",
            reference_conditions=[LTX2ReferenceCondition(frames=reference_frames, strength=1.0)],
            reference_downscale_factor=config.reference_downscale_factor,
            width=out_width,
            height=out_height,
            num_frames=frame_count,
            frame_rate=FPS,
            num_inference_steps=NUM_STEPS,
            sigmas=DISTILLED_SIGMA_VALUES,
            guidance_scale=1.0,
            stg_scale=0.0,
            audio_guidance_scale=1.0,
            audio_stg_scale=0.0,
            generator=torch.Generator(device="cuda").manual_seed(actual_seed),
            callback_on_step_end=_callback,
            output_type="np",
            return_dict=False,
        )

    temp_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    audio_kwargs = {}
    if audio_out is not None:
        audio_kwargs = {"audio": audio_out[0].float().cpu(), "audio_sample_rate": pipe.vocoder.config.output_sampling_rate}
    encode_video(video_out[0], fps=FPS, output_path=temp_path, **audio_kwargs)
    final_path = _save_iteration(temp_path, config, actual_seed, frame_count, prompt)
    history, zip_path = refresh_history()
    status = f"{config.output_label} render saved as {Path(final_path).name}"
    return final_path, final_path, actual_seed, history, zip_path, status


def clear_history() -> tuple[list[list[str]], None, str]:
    # Clearing history removes only this app's bucket render folder, not any other mounted data.
    if RUNS_DIR.exists():
        shutil.rmtree(RUNS_DIR)
    _ensure_storage()
    return [], None, "Render history cleared."


def _click(effect_key: str):
    # Each button binds a fixed effect key while sharing the same visible controls.
    return lambda video, prompt, seed, randomize: run_effect(effect_key, video, prompt, seed, randomize)


with gr.Blocks(title="LTX-2.3 Multi-Effect Video Workflow") as demo:
    gr.Markdown(
        "# LTX-2.3 Multi-Effect Video Workflow\n"
        "Upload a video, apply one effect at a time, then continue from the latest render. "
        "Completed iterations are stored in the mounted bucket and can be downloaded together."
    )
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Current video")
            prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                placeholder="Describe the scene and any sound you want preserved or restored.",
            )
            with gr.Accordion("Settings", open=False):
                randomize = gr.Checkbox(True, label="Randomize seed")
                seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
            with gr.Row():
                decompress_btn = gr.Button("Decompress", variant="primary")
                deblur_btn = gr.Button("Deblur", variant="primary")
                colorize_btn = gr.Button("Colorize", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Latest render")
            status = gr.Markdown()
            history = gr.Dataframe(
                headers=["#", "Effect", "Frames", "Seed", "Created", "File"],
                datatype=["str", "str", "str", "str", "str", "str"],
                label="Available renderings",
                interactive=False,
            )
            with gr.Row():
                refresh_btn = gr.Button("Refresh")
                clear_btn = gr.Button("Clear history")
            download_all = gr.File(label="Download all renders")

    button_inputs = [video_in, prompt, seed, randomize]
    button_outputs = [video_out, video_in, seed, history, download_all, status]
    decompress_btn.click(_click("decompress"), inputs=button_inputs, outputs=button_outputs)
    deblur_btn.click(_click("deblur"), inputs=button_inputs, outputs=button_outputs)
    colorize_btn.click(_click("colorize"), inputs=button_inputs, outputs=button_outputs)
    refresh_btn.click(refresh_history, inputs=[], outputs=[history, download_all])
    clear_btn.click(clear_history, inputs=[], outputs=[history, download_all, status])
    demo.load(refresh_history, inputs=[], outputs=[history, download_all])


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