import os import subprocess import sys # Allocator config for pixel-space video SR memory spikes — set before torch import. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") # Disable torch.compile / dynamo before any torch import (unsupported on ZeroGPU). os.environ["TORCH_COMPILE_DISABLE"] = "1" os.environ["TORCHDYNAMO_DISABLE"] = "1" # ── Install xformers (memory-efficient attention) and video I/O deps ────────── subprocess.run( [sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"], check=False, ) subprocess.run( [sys.executable, "-m", "pip", "install", "imageio[ffmpeg]", "av", "num2words"], check=False, ) # Reinstall torchaudio to match the pre-installed CUDA torch (deps can pull a CPU-only build). _tv = subprocess.run( [sys.executable, "-c", "import torch; print(torch.__version__)"], capture_output=True, text=True, ) if _tv.returncode == 0: _full_ver = _tv.stdout.strip() _cuda_suffix = _full_ver.split("+")[-1] if "+" in _full_ver else "cu124" _base_ver = _full_ver.split("+")[0] print(f"Detected torch {_full_ver}, reinstalling matching torchaudio...") subprocess.run( [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", f"torchaudio=={_base_ver}", "--index-url", f"https://download.pytorch.org/whl/{_cuda_suffix}"], check=False, ) # ── Clone LTX-2 repo at a pinned commit and install ltx-core / ltx-pipelines ── 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 os.path.exists(LTX_REPO_DIR): subprocess.run(["rm", "-rf", LTX_REPO_DIR], check=True) print(f"Cloning {LTX_REPO_URL}...") subprocess.run(["git", "clone", LTX_REPO_URL, LTX_REPO_DIR], check=True) print(f"Checking out commit {LTX_COMMIT}...") subprocess.run(["git", "-C", LTX_REPO_DIR, "checkout", LTX_COMMIT], check=True) print("Installing ltx-core and ltx-pipelines from pinned repo commit...") 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 torch torch._dynamo.config.suppress_errors = True torch._dynamo.config.disable = True import spaces import gradio as gr import numpy as np from huggingface_hub import hf_hub_download, snapshot_download 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 import QuantizationPolicy from ltx_pipelines.ic_lora import ICLoraPipeline from ltx_pipelines.utils.media_io import encode_video # ── Disable xformers FA3 dispatch (Blackwell / RTX PRO 6000 is sm_120; FA3 is Hopper-only) ── try: from xformers.ops.fmha import _set_use_fa3 _set_use_fa3(False) print("[ATTN] xformers FA3 dispatch disabled (Blackwell-incompatible)") except Exception as e: print(f"[ATTN] FA3 disable skipped: {type(e).__name__}: {e}") logging.getLogger().setLevel(logging.INFO) MAX_SEED = np.iinfo(np.int32).max DEFAULT_FRAME_RATE = 24.0 # ───────────────────────────────────────────────────────────────────────────── # Model / checkpoint locations # ───────────────────────────────────────────────────────────────────────────── LTX_MODEL_REPO = "Lightricks/LTX-2.3" DISTILLED_FILENAME = "ltx-2.3-22b-distilled-1.1.safetensors" SPATIAL_UPSAMPLER_FILENAME = "ltx-2.3-spatial-upscaler-x2-1.1.safetensors" GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized" UPSCALER_REPO = "Lightricks/LTX-2.3-22b-IC-LoRA-Pixel-Spatial-Upscaler" UPSCALER_X4_FILENAME = "ltx-2.3-22b-ic-lora-pixel-spatial-upscaler-x4-0.9.safetensors" print("=" * 80) print("Downloading LTX-2.3 distilled base + spatial upsampler + Gemma + upscaler IC-LoRA...") print("=" * 80) checkpoint_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename=DISTILLED_FILENAME) spatial_upsampler_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename=SPATIAL_UPSAMPLER_FILENAME) gemma_root = snapshot_download(repo_id=GEMMA_REPO) upscaler_lora_path = hf_hub_download(repo_id=UPSCALER_REPO, filename=UPSCALER_X4_FILENAME) print(f"Checkpoint: {checkpoint_path}") print(f"Spatial upsampler: {spatial_upsampler_path}") print(f"Gemma root: {gemma_root}") print(f"Upscaler IC-LoRA: {upscaler_lora_path}") # ───────────────────────────────────────────────────────────────────────────── # IC-LoRA pipeline (4x spatial upscaler) — built lazily inside the GPU worker. # # The LTX loader opens safetensors with `safe_open(device="cuda")` and streams # tensors straight onto the GPU, which bypasses the ZeroGPU `.to("cuda")` hijack # and fails at module scope ("No CUDA GPUs are available"). So we construct the # pipeline on the first `@spaces.GPU` call, where a real GPU is attached, and # cache it in a module global for subsequent calls on the same warm worker. # ───────────────────────────────────────────────────────────────────────────── ic_loras = [ LoraPathStrengthAndSDOps(upscaler_lora_path, 1.0, LTXV_LORA_COMFY_RENAMING_MAP) ] _PIPELINE = None REFERENCE_DOWNSCALE_FACTOR = 4 # x4 upscaler LoRA (from its safetensors metadata) def get_pipeline(): """Build + cache the ICLoraPipeline. Must be called inside a GPU worker.""" global _PIPELINE if _PIPELINE is None: print("[Pipeline] Building ICLoraPipeline (first GPU call)...") _PIPELINE = ICLoraPipeline( distilled_checkpoint_path=checkpoint_path, spatial_upsampler_path=spatial_upsampler_path, gemma_root=gemma_root, loras=ic_loras, quantization=QuantizationPolicy.fp8_cast(), ) print(f"[Pipeline] reference_downscale_factor = {_PIPELINE.reference_downscale_factor}") print("[Pipeline] Ready.") return _PIPELINE # ───────────────────────────────────────────────────────────────────────────── # Resolution helpers # ───────────────────────────────────────────────────────────────────────────── def _round_to(x: int, base: int) -> int: return max(base, int(round(x / base)) * base) def _probe_video(video_path: str): """Return (width, height, fps, duration_seconds) of a video.""" import av with av.open(str(video_path)) as container: s = container.streams.video[0] w = s.codec_context.width h = s.codec_context.height fps = float(s.average_rate) if s.average_rate else DEFAULT_FRAME_RATE if s.duration is not None and s.time_base is not None: dur = float(s.duration * s.time_base) elif container.duration is not None: dur = float(container.duration) / 1_000_000.0 else: dur = 3.0 return w, h, fps, dur # VAE spatial compression, read directly from the checkpoint's embedded VAE # config (encoder_blocks): patch_size(4) * compress_space_res(2) * # compress_all_res(2) * compress_all_res(2) = 32. NOT taken from # ltx_core.types.SpatioTemporalScaleFactors.default() -- that "default" is a # hardcoded literal on VideoDecoder unrelated to the actual decoder_blocks it's # built with, so it can't be trusted to reflect any given checkpoint. VAE_SPATIAL_COMPRESSION = 32 # ICLoraPipeline runs stage 1 at half the requested output resolution # (ic_lora.py: `VideoPixelShape(width=width // 2, height=height // 2, ...)`), # and stage 1's IC-LoRA reference-conditioning video is encoded at that # half-resolution divided again by reference_downscale_factor. So the video # actually reaching the VAE encoder is: # height_full / STAGE_1_RESOLUTION_DIVISOR / reference_downscale_factor # and for that to be divisible by VAE_SPATIAL_COMPRESSION, height_full (and # width_full) must be divisible by their product. STAGE_1_RESOLUTION_DIVISOR = 2 def _reference_resolution_multiple(scale: int) -> int: """Full output resolution must be a multiple of this for the IC-LoRA reference-conditioning video (see constants above) to survive VAE encode.""" return STAGE_1_RESOLUTION_DIVISOR * scale * VAE_SPATIAL_COMPRESSION def compute_output_resolution(video_path: str, scale: int = REFERENCE_DOWNSCALE_FACTOR): """Target output = input resolution * scale, snapped to a multiple of _reference_resolution_multiple(scale).""" w, h, _, _ = _probe_video(video_path) res_multiple = _reference_resolution_multiple(scale) out_w = _round_to(w * scale, res_multiple) out_h = _round_to(h * scale, res_multiple) # Cap the longest side to keep runtime + VRAM bounded. max_side = 1536 if max(out_w, out_h) > max_side: ratio = max_side / max(out_w, out_h) out_w = _round_to(out_w * ratio, res_multiple) out_h = _round_to(out_h * ratio, res_multiple) return int(out_w), int(out_h) def on_video_upload(video_path): if video_path is None: return gr.update(), gr.update() try: out_w, out_h = compute_output_resolution(video_path, scale=REFERENCE_DOWNSCALE_FACTOR) return gr.update(value=out_w), gr.update(value=out_h) except Exception as e: print(f"[probe] failed: {e}") return gr.update(), gr.update() # ───────────────────────────────────────────────────────────────────────────── # Generation # ───────────────────────────────────────────────────────────────────────────── @spaces.GPU(duration=300) @torch.inference_mode() def upscale_video( input_video, prompt: str = "", lora_strength: float = 1.0, enhance_prompt: bool = True, seed: int = 42, randomize_seed: bool = True, width: int = 1536, height: int = 768, duration_seconds: float = 3.0, progress=gr.Progress(track_tqdm=True), ): """Creatively upscale a low-resolution video 4x with the LTX-2.3 IC-LoRA Pixel Spatial Upscaler. Args: input_video: Path to the low-resolution source video to upscale. prompt: Text describing the scene (helps the model synthesize detail). lora_strength: IC-LoRA reference strength. Lower stays closer to the source, higher hallucinates more detail. enhance_prompt: Auto-enhance the prompt with the built-in text model. seed: RNG seed. randomize_seed: Pick a fresh random seed each run. width: Target output width in pixels (rounded to a multiple of _reference_resolution_multiple(scale), 256 for the default 4x LoRA). height: Target output height in pixels (same rounding as width). duration_seconds: Length of the output clip in seconds. """ if input_video is None: raise gr.Error("Please provide a low-resolution input video to upscale.") pipeline = get_pipeline() current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) frame_rate = DEFAULT_FRAME_RATE try: _, _, src_fps, src_dur = _probe_video(str(input_video)) if src_fps and 1.0 <= src_fps <= 60.0: frame_rate = float(src_fps) duration_seconds = min(float(duration_seconds), float(src_dur) + 0.1) except Exception as e: print(f"[probe] failed, using defaults: {e}") num_frames = int(duration_seconds * frame_rate) + 1 num_frames = ((num_frames - 1 + 7) // 8) * 8 + 1 # 8k + 1 num_frames = max(9, min(num_frames, 121)) scale = pipeline.reference_downscale_factor # 4, but read live from the pipeline res_multiple = _reference_resolution_multiple(scale) width = _round_to(int(width), res_multiple) height = _round_to(int(height), res_multiple) print(f"[Upscale x{scale}] -> {width}x{height}, {num_frames} frames @ {frame_rate}fps, " f"seed={current_seed}, strength={lora_strength}") tiling_config = TilingConfig.default() video_chunks_number = get_video_chunks_number(num_frames, tiling_config) video, audio = pipeline( prompt=prompt or "", seed=current_seed, height=int(height), width=int(width), num_frames=num_frames, frame_rate=frame_rate, images=[], video_conditioning=[(str(input_video), float(lora_strength))], enhance_prompt=enhance_prompt, tiling_config=tiling_config, conditioning_attention_strength=1.0, ) output_path = tempfile.mktemp(suffix=".mp4") encode_video( video=video, fps=frame_rate, audio=audio, output_path=output_path, video_chunks_number=video_chunks_number, ) return str(output_path), current_seed # ───────────────────────────────────────────────────────────────────────────── # UI # ───────────────────────────────────────────────────────────────────────────── CSS = """ #col-container { max-width: 1150px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } video { object-fit: contain !important; } """ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """ # LTX-2.3 Pixel Spatial Upscaler (4×) 🎞️→✨ Creatively upscale a **low-resolution video 4×** with the [LTX-2.3-22B IC-LoRA Pixel Spatial Upscaler](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Pixel-Spatial-Upscaler) on top of [LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3). This is a **generative** upscaler — it *synthesizes* new fine detail rather than interpolating, so the result is a high-resolution re-render of your clip. Add a prompt describing the scene to guide the detail. """ ) with gr.Row(): with gr.Column(scale=1): input_video = gr.Video(label="Low-resolution input video") prompt = gr.Textbox( label="Prompt (optional but recommended)", placeholder="describe the scene, e.g. 'a jellyfish gliding through deep blue water, bioluminescent detail'", lines=2, ) run_btn = gr.Button("Upscale 4×", variant="primary", size="lg") with gr.Accordion("Advanced settings", open=False): lora_strength = gr.Slider( label="Reference strength", info="Lower = closer to source, higher = more synthesized detail", minimum=0.5, maximum=1.0, value=1.0, step=0.05, ) enhance_prompt = gr.Checkbox(label="Enhance prompt", value=True) duration_seconds = gr.Slider( label="Output duration (s)", minimum=0.5, maximum=5.0, value=3.0, step=0.5, ) seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=42, step=1) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Row(): width = gr.Number(label="Output width", value=1536, precision=0) height = gr.Number(label="Output height", value=768, precision=0) with gr.Column(scale=1): output_video = gr.Video(label="Upscaled result (4×)", autoplay=True) gr.Examples( examples=[ ["jellyfish_glide.mp4", "a jellyfish gliding gracefully through deep blue water, translucent bioluminescent detail"], ["waterfall_forest.mp4", "a waterfall cascading through a lush green forest, misty spray and wet mossy rocks"], ["squirrel_eating.mp4", "a squirrel eating a nut, detailed fur, soft natural daylight"], ], inputs=[input_video, prompt], outputs=[output_video, seed], fn=upscale_video, cache_examples=True, cache_mode="lazy", ) input_video.change(fn=on_video_upload, inputs=[input_video], outputs=[width, height]) run_btn.click( fn=upscale_video, inputs=[input_video, prompt, lora_strength, enhance_prompt, seed, randomize_seed, width, height, duration_seconds], outputs=[output_video, seed], api_name="upscale", ) if __name__ == "__main__": demo.launch(mcp_server=True)