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)