Spaces:
Sleeping
Sleeping
| """LTX-2.3 Gemma text-encoder backend (API Space for the Pre-Production Playground). | |
| Runs Gemma-3-12B + the LTX-2.3 embeddings connectors and returns final | |
| EmbeddingsProcessorOutput tensors (video/audio encodings + attention mask) as a .pt file, | |
| for positive and (optionally) negative prompts — so the generation backends don't need | |
| to hold Gemma in VRAM. Mirrors helpers.encode_prompts at the pinned LTX-2 commit. | |
| Adapted from linoyts/gemma-text-encoder (LTX-2) to LTX-2.3. | |
| """ | |
| import os | |
| import subprocess | |
| import sys | |
| os.environ["TORCH_COMPILE_DISABLE"] = "1" | |
| os.environ["TORCHDYNAMO_DISABLE"] = "1" | |
| 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 struct | |
| import tempfile | |
| import time | |
| import torch | |
| torch._dynamo.config.suppress_errors = True | |
| torch._dynamo.config.disable = True | |
| import gradio as gr | |
| import spaces | |
| from huggingface_hub import hf_hub_download, snapshot_download | |
| from ltx_pipelines.utils.helpers import generate_enhanced_prompt | |
| from ltx_pipelines.utils.model_ledger import ModelLedger | |
| # 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 | |
| LTX_REPO = os.environ.get("LTX_REPO", "Lightricks/LTX-2.3") | |
| CHECKPOINT_FILE = "ltx-2.3-22b-distilled-1.1.safetensors" # connector weights only | |
| GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized" | |
| TOKEN = os.environ.get("HF_TOKEN") | |
| print("Downloading checkpoint (connectors) + Gemma…") | |
| checkpoint_path = hf_hub_download(LTX_REPO, CHECKPOINT_FILE, token=TOKEN) | |
| gemma_root = snapshot_download(GEMMA_REPO, token=TOKEN) | |
| ledger = ModelLedger( | |
| dtype=torch.bfloat16, | |
| device="cuda", | |
| checkpoint_path=checkpoint_path, | |
| gemma_root_path=gemma_root, | |
| ) | |
| print("Preloading text encoder + embeddings processor…") | |
| _text_encoder = ledger.text_encoder() | |
| _embeddings_processor = ledger.gemma_embeddings_processor() | |
| print("Encoder ready.") | |
| def _pack(output): | |
| return {"video": output.video_encoding.cpu(), | |
| "audio": output.audio_encoding.cpu() if output.audio_encoding is not None else None, | |
| "mask": output.attention_mask.cpu()} | |
| def encode( | |
| prompt: str, | |
| negative_prompt: str = "", | |
| encode_negative: bool = False, | |
| enhance_prompt: bool = True, | |
| seed: int = 42, | |
| ): | |
| """Encode prompt (and optionally negative prompt) → .pt with final LTX-2.3 embeddings. | |
| Returns (embeddings_file, final_prompt, status). | |
| """ | |
| start = time.time() | |
| final_prompt = prompt | |
| if enhance_prompt: | |
| final_prompt = generate_enhanced_prompt(_text_encoder, prompt, None, seed=int(seed)) | |
| prompts = [final_prompt] + ([negative_prompt or ""] if encode_negative else []) | |
| raw = [_text_encoder.encode(p) for p in prompts] | |
| outputs = [_embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw] | |
| data = {"positive": _pack(outputs[0]), | |
| "negative": _pack(outputs[1]) if encode_negative else None, | |
| "final_prompt": final_prompt} | |
| out_path = tempfile.mktemp(suffix=".pt") | |
| torch.save(data, out_path) | |
| status = f"encoded in {time.time() - start:.1f}s" | |
| return out_path, final_prompt, status | |
| with gr.Blocks(title="LTX-2.3 Gemma encoder API") as demo: | |
| gr.Markdown("# 🧠 LTX-2.3 Gemma text encoder\n" | |
| "API backend for the LTX-2.3 Pre-Production Playground — returns final " | |
| "embeddings (video/audio encodings + mask) for positive and negative prompts.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt = gr.Textbox(label="Prompt", lines=4) | |
| negative_prompt = gr.Textbox(label="Negative prompt", lines=2) | |
| encode_negative = gr.Checkbox(label="Encode negative (CFG / two-stage)", value=False) | |
| enhance_prompt = gr.Checkbox(label="Enhance prompt", value=True) | |
| seed = gr.Number(label="Enhancement seed", value=42, precision=0) | |
| btn = gr.Button("Encode", variant="primary") | |
| with gr.Column(): | |
| emb_file = gr.File(label="Embeddings (.pt)") | |
| final_prompt = gr.Textbox(label="Final prompt used", lines=4) | |
| status = gr.Textbox(label="Status") | |
| btn.click(encode, [prompt, negative_prompt, encode_negative, enhance_prompt, seed], | |
| [emb_file, final_prompt, status], api_name="encode") | |
| if __name__ == "__main__": | |
| demo.launch() | |