Spaces:
Sleeping
Sleeping
| """LTX-2.3 dev two-stage T2V backend (API Space for the Pre-Production Playground frontend). | |
| TI2VidTwoStagesPipeline: dev checkpoint stage 1 with CFG/STG guidance, distilled-LoRA | |
| stage 2 + x2 spatial upsampler. Text encoding (positive AND negative, for CFG) is | |
| offloaded to the Gemma encoder Space (TEXT_ENCODER_SPACE) — encode_prompts is | |
| monkeypatched with the precomputed [ctx_p, ctx_n]. Same ZeroGPU setup as the distilled | |
| backend. | |
| """ | |
| import os | |
| import subprocess | |
| import sys | |
| os.environ["TORCH_COMPILE_DISABLE"] = "1" | |
| os.environ["TORCHDYNAMO_DISABLE"] = "1" | |
| subprocess.run([sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", | |
| "--no-build-isolation"], check=False) | |
| 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_SHA = os.environ.get("LTX_COMMIT_SHA", "ae855f8538843825f9015a419cf4ba5edaf5eec2") | |
| if not os.path.exists(LTX_REPO_DIR): | |
| os.makedirs(LTX_REPO_DIR) | |
| subprocess.run(["git", "init", LTX_REPO_DIR], check=True) | |
| subprocess.run(["git", "remote", "add", "origin", LTX_REPO_URL], cwd=LTX_REPO_DIR, check=True) | |
| subprocess.run(["git", "fetch", "--depth", "1", "origin", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, check=True) | |
| subprocess.run(["git", "checkout", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, 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 json | |
| import logging | |
| import random | |
| import struct | |
| import tempfile | |
| import torch | |
| torch._dynamo.config.suppress_errors = True | |
| torch._dynamo.config.disable = True | |
| import gradio as gr | |
| import numpy as np | |
| import spaces | |
| from gradio_client import Client | |
| from huggingface_hub import hf_hub_download | |
| import ltx_pipelines.ti2vid_two_stages as ti2vid_module | |
| from ltx_core.components.guiders import MultiModalGuiderParams | |
| from ltx_core.loader import LoraPathStrengthAndSDOps | |
| from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number | |
| from ltx_core.quantization import QuantizationPolicy | |
| from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput | |
| from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline | |
| from ltx_pipelines.utils.media_io import encode_video | |
| # xformers attention: force-patch + disable Hopper-only FA3 dispatch (Blackwell ZeroGPU). | |
| from ltx_core.model.transformer import attention as _attn_mod | |
| try: | |
| from xformers.ops import memory_efficient_attention as _mea | |
| _attn_mod.memory_efficient_attention = _mea | |
| from xformers.ops.fmha import _set_use_fa3 | |
| _set_use_fa3(False) | |
| print("[ATTN] xformers patched, FA3 disabled") | |
| except Exception as e: | |
| print(f"[ATTN] xformers patch failed: {type(e).__name__}: {e}") | |
| # Chunked-read safetensors loader: safe_open mmap deadlocks on FUSE-backed storage. | |
| from ltx_core.loader.primitives import StateDict | |
| from ltx_core.loader.sft_loader import SafetensorsStateDictLoader | |
| _SAFETENSORS_DTYPE_MAP = { | |
| "F64": torch.float64, "F32": torch.float32, "F16": torch.float16, | |
| "BF16": torch.bfloat16, "F8_E5M2": torch.float8_e5m2, "F8_E4M3": torch.float8_e4m3fn, | |
| "I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8, | |
| "U8": torch.uint8, "BOOL": torch.bool, | |
| } | |
| def _patched_load(self, path, sd_ops, device=None): | |
| sd, size, dtype = {}, 0, set() | |
| device = device or torch.device("cpu") | |
| for shard_path in (path if isinstance(path, list) else [path]): | |
| with open(shard_path, "rb") as f: | |
| header_len = struct.unpack("<Q", f.read(8))[0] | |
| header = json.loads(f.read(header_len).decode("utf-8")) | |
| data_base = 8 + header_len | |
| for name, meta in header.items(): | |
| if name == "__metadata__": | |
| continue | |
| expected_name = name if sd_ops is None else sd_ops.apply_to_key(name) | |
| if expected_name is None: | |
| continue | |
| start, end = meta["data_offsets"] | |
| f.seek(data_base + start) | |
| buf = f.read(end - start) | |
| t = torch.frombuffer(bytearray(buf), dtype=_SAFETENSORS_DTYPE_MAP[meta["dtype"]] | |
| ).reshape(meta["shape"]) | |
| t = t.to(device=device, non_blocking=True, copy=False) | |
| kvs = (((expected_name, t),) if sd_ops is None | |
| else sd_ops.apply_to_key_value(expected_name, t)) | |
| for key, v in kvs: | |
| size += v.nbytes | |
| dtype.add(v.dtype) | |
| sd[key] = v | |
| return StateDict(sd=sd, device=device, size=size, dtype=dtype) | |
| SafetensorsStateDictLoader.load = _patched_load | |
| logging.getLogger().setLevel(logging.INFO) | |
| MAX_SEED = np.iinfo(np.int32).max | |
| FRAME_RATE = 24.0 | |
| DEFAULT_NEGATIVE = ( | |
| "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, " | |
| "motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, " | |
| "transition, static" | |
| ) | |
| # LTX-2.3 guider defaults | |
| VIDEO_STG, VIDEO_RESCALE, A2V_SCALE, STG_BLOCKS = 1.0, 0.7, 3.0, [28] | |
| AUDIO_CFG, AUDIO_STG, AUDIO_RESCALE, V2A_SCALE = 7.0, 1.0, 0.7, 3.0 | |
| LTX_REPO = os.environ.get("LTX_REPO", "Lightricks/LTX-2.3") | |
| DEV_FILE = "ltx-2.3-22b-dev.safetensors" | |
| DISTILLED_LORA_FILE = "ltx-2.3-22b-distilled-lora-384-1.1.safetensors" | |
| UPSCALER_FILE = "ltx-2.3-spatial-upscaler-x2-1.1.safetensors" | |
| TOKEN = os.environ.get("HF_TOKEN") | |
| TEXT_ENCODER_SPACE = os.environ.get("TEXT_ENCODER_SPACE", "linoyts/ltx23-gemma-encoder-api") | |
| print("Downloading checkpoints…") | |
| checkpoint_path = hf_hub_download(LTX_REPO, DEV_FILE, token=TOKEN) | |
| distilled_lora_path = hf_hub_download(LTX_REPO, DISTILLED_LORA_FILE, token=TOKEN) | |
| upsampler_path = hf_hub_download(LTX_REPO, UPSCALER_FILE, token=TOKEN) | |
| pipeline = TI2VidTwoStagesPipeline( | |
| checkpoint_path=checkpoint_path, | |
| distilled_lora=[LoraPathStrengthAndSDOps(path=distilled_lora_path, strength=1.0, sd_ops=None)], | |
| spatial_upsampler_path=upsampler_path, | |
| gemma_root=None, # text encoding happens on TEXT_ENCODER_SPACE | |
| loras=[], | |
| quantization=QuantizationPolicy.fp8_cast(), | |
| ) | |
| print("Preloading models (ZeroGPU tensor packing)…") | |
| s1, s2 = pipeline.stage_1_model_ledger, pipeline.stage_2_model_ledger | |
| _cached_s1 = {n: getattr(s1, n)() for n in ("transformer", "video_encoder")} | |
| _cached_s2 = {n: getattr(s2, n)() for n in ( | |
| "transformer", "video_decoder", "audio_decoder", "vocoder", "spatial_upsampler")} | |
| for name, model in _cached_s1.items(): | |
| setattr(s1, name, (lambda m: (lambda: m))(model)) | |
| for name, model in _cached_s2.items(): | |
| setattr(s2, name, (lambda m: (lambda: m))(model)) | |
| print("Pipeline ready.") | |
| # ---- remote text encoding (positive + negative, needed for CFG) ---- | |
| _emb_cache = {} | |
| def fetch_embeddings(prompt, negative_prompt, enhance_prompt, seed): | |
| """Call the Gemma encoder Space; returns {'positive', 'negative', 'final_prompt'}.""" | |
| key = (prompt, negative_prompt, bool(enhance_prompt), int(seed) if enhance_prompt else 0) | |
| if key in _emb_cache: | |
| return _emb_cache[key] | |
| client = Client(TEXT_ENCODER_SPACE, token=TOKEN) | |
| emb_file, final_prompt, status = client.predict( | |
| prompt=prompt, negative_prompt=negative_prompt or "", encode_negative=True, | |
| enhance_prompt=enhance_prompt, seed=int(seed), api_name="/encode", | |
| ) | |
| print(f"[encoder] {status} | final prompt: {final_prompt[:80]}…") | |
| data = torch.load(emb_file, map_location="cpu", weights_only=True) | |
| if len(_emb_cache) > 32: | |
| _emb_cache.clear() | |
| _emb_cache[key] = data | |
| return data | |
| def _to_output(pack, device): | |
| return EmbeddingsProcessorOutput( | |
| video_encoding=pack["video"].to(device), | |
| audio_encoding=pack["audio"].to(device) if pack["audio"] is not None else None, | |
| attention_mask=pack["mask"].to(device), | |
| ) | |
| def gpu_duration(embeddings, prompt, negative_prompt="", duration=3.0, *args, **kwargs): | |
| return min(420, 120 + int(duration) * 30) | |
| def _generate_gpu(embeddings, prompt, negative_prompt, duration, used_seed, | |
| num_inference_steps, video_cfg_scale, width, height, | |
| progress=gr.Progress(track_tqdm=True)): | |
| num_frames = ((int(duration * FRAME_RATE) + 1 - 1 + 7) // 8) * 8 + 1 | |
| tiling_config = TilingConfig.default() | |
| video_chunks_number = get_video_chunks_number(num_frames, tiling_config) | |
| precomputed = [_to_output(embeddings["positive"], "cuda"), | |
| _to_output(embeddings["negative"], "cuda")] | |
| original = ti2vid_module.encode_prompts | |
| ti2vid_module.encode_prompts = lambda *a, **kw: precomputed | |
| try: | |
| video, audio = pipeline( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| seed=used_seed, | |
| height=int(height), | |
| width=int(width), | |
| num_frames=num_frames, | |
| frame_rate=FRAME_RATE, | |
| num_inference_steps=int(num_inference_steps), | |
| video_guider_params=MultiModalGuiderParams( | |
| cfg_scale=video_cfg_scale, stg_scale=VIDEO_STG, rescale_scale=VIDEO_RESCALE, | |
| modality_scale=A2V_SCALE, skip_step=0, stg_blocks=STG_BLOCKS), | |
| audio_guider_params=MultiModalGuiderParams( | |
| cfg_scale=AUDIO_CFG, stg_scale=AUDIO_STG, rescale_scale=AUDIO_RESCALE, | |
| modality_scale=V2A_SCALE, skip_step=0, stg_blocks=STG_BLOCKS), | |
| images=[], | |
| tiling_config=tiling_config, | |
| enhance_prompt=False, # already enhanced on the encoder Space | |
| ) | |
| finally: | |
| ti2vid_module.encode_prompts = original | |
| 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 output_path | |
| def generate( | |
| prompt: str, | |
| negative_prompt: str = DEFAULT_NEGATIVE, | |
| duration: float = 5.0, | |
| seed: int = 42, | |
| randomize_seed: bool = False, | |
| num_inference_steps: int = 30, | |
| video_cfg_scale: float = 3.0, | |
| width: int = 1536, | |
| height: int = 864, | |
| enhance_prompt: bool = True, | |
| progress=gr.Progress(track_tqdm=True), | |
| ): | |
| """Dev two-stage T2V: returns (video_path, used_seed). Encoder call runs outside the GPU slot.""" | |
| used_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) | |
| progress(0.05, desc="encoding prompts (Gemma Space)…") | |
| embeddings = fetch_embeddings(prompt, negative_prompt, enhance_prompt, used_seed) | |
| output_path = _generate_gpu(embeddings, prompt, negative_prompt, duration, used_seed, | |
| num_inference_steps, video_cfg_scale, width, height) | |
| return output_path, used_seed | |
| with gr.Blocks(title="LTX-2.3 dev API") as demo: | |
| gr.Markdown("# 🎬 LTX-2.3 dev two-stage — hero backend\n" | |
| "API backend for the LTX-2.3 Pre-Production Playground.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt = gr.Textbox(label="Prompt", lines=3) | |
| negative_prompt = gr.Textbox(label="Negative prompt", value=DEFAULT_NEGATIVE, lines=2) | |
| duration = gr.Slider(1, 10, value=5, step=0.5, label="Duration (s)") | |
| with gr.Row(): | |
| seed = gr.Number(label="Seed", value=42, precision=0) | |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=False) | |
| with gr.Row(): | |
| num_inference_steps = gr.Slider(10, 50, value=30, step=1, label="Steps") | |
| video_cfg_scale = gr.Slider(1, 8, value=3.0, step=0.1, label="CFG") | |
| with gr.Row(): | |
| width = gr.Number(label="Width", value=1536, precision=0) | |
| height = gr.Number(label="Height", value=864, precision=0) | |
| enhance_prompt = gr.Checkbox(label="Enhance prompt", value=True) | |
| btn = gr.Button("Generate", variant="primary") | |
| with gr.Column(): | |
| video = gr.Video(label="Video", autoplay=True) | |
| used_seed = gr.Number(label="Used seed", precision=0) | |
| btn.click(generate, | |
| [prompt, negative_prompt, duration, seed, randomize_seed, num_inference_steps, | |
| video_cfg_scale, width, height, enhance_prompt], | |
| [video, used_seed], api_name="generate") | |
| if __name__ == "__main__": | |
| demo.launch() | |