"""LTX-2.5 workflow — model runs directly on this Space's ZeroGPU allocation. Instead of calling the upstream Lightricks/LTX-2.5 Space via gradio_client (which drops the visitor's X-IP-Token and shares one anonymous rate-limit bucket), the pipeline is loaded here and bound to the canvas as a `fn` operator. Each visitor's run uses their own ZeroGPU quota automatically. Model code adapted from https://huggingface.co/spaces/Lightricks/LTX-2.5 (two-stage distilled recipe: 8 steps @ half res, x2 latent upsample, 3 steps @ full res, then decode). """ import os # ZeroGPU's MIG slice trips an NVML assert in torch's native caching allocator; # cudaMallocAsync is NVML-safe. Must be set before torch initializes. os.environ.setdefault("PYTORCH_ALLOC_CONF", "backend:cudaMallocAsync") os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync") import hashlib import random import tempfile import time from collections import OrderedDict import gradio as gr import PIL.Image import spaces import torch from huggingface_hub import snapshot_download from diffusers import ( LTX2ImageToVideoPipeline, LTX2LatentUpsamplePipeline, LTX2Pipeline, LTX2VideoDiffusionDecodePipeline, LTX2VideoDiffusionDecoderModel, ) from diffusers.models.autoencoders.ltx2_diffusion_decoder import ( LTX2VideoVaeNeighborhoodNattenProcessor, ) from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.utils import ( DEFAULT_NEGATIVE_PROMPT, DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES, ) from diffusers.utils import encode_video HF_TOKEN = os.environ.get("HF_TOKEN") # Gated (auto-approve): the Space's HF_TOKEN secret must have accepted the license. MODEL_ID = os.environ.get("LTX25_MODEL_ID", "Lightricks/LTX-2.5-Diffusers") MAX_SEED = 2**31 - 1 FRAME_RATE = 24.0 STAGE_1_SIGMAS = DISTILLED_SIGMA_VALUES STAGE_2_SIGMAS = STAGE_2_DISTILLED_SIGMA_VALUES GUIDANCE_SCALE = 1.0 AUDIO_GUIDANCE_SCALE = 1.0 AUTO_MAX_SECONDS = 15.0 AUTO_MAX_FRAMES = int(AUTO_MAX_SECONDS * FRAME_RATE) // 8 * 8 + 1 # 361, 8k+1 grid def frames_from_duration(seconds: float) -> int: """Snap seconds to the VAE's 8k+1 frame grid at 24 fps (2 s -> 49 frames).""" return max(25, int(round(float(seconds) * FRAME_RATE)) // 8 * 8 + 1) def load_conditioning_image(path: str, width: int, height: int) -> PIL.Image.Image: """Cover-resize and center-crop to exactly width x height (no aspect distortion).""" img = PIL.Image.open(path).convert("RGB") scale = max(width / img.width, height / img.height) img = img.resize((round(img.width * scale), round(img.height * scale)), PIL.Image.LANCZOS) left, top = (img.width - width) // 2, (img.height - height) // 2 return img.crop((left, top, left + width, top + height)) print("[ltx25] loading LTX-2.5 distilled...", flush=True) # Snapshot without `transformer_full/` (from_pretrained would otherwise pull both DiTs). MODEL_DIR = snapshot_download( MODEL_ID, ignore_patterns=["transformer_full/*"], token=HF_TOKEN, max_workers=8 ) pipe = LTX2Pipeline.from_pretrained(MODEL_DIR, dtype=torch.bfloat16) diffusion_decoder = LTX2VideoDiffusionDecoderModel.from_pretrained( MODEL_DIR, subfolder="diffusion_decoder", dtype=torch.bfloat16 ) diffusion_decoder.set_attn_processor(LTX2VideoVaeNeighborhoodNattenProcessor()) diffusion_decoder.enable_tiling() latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( MODEL_DIR, subfolder="latent_upsampler", dtype=torch.bfloat16 ) # Module scope + .to("cuda") is the ZeroGPU pattern: weights are packed to disk # and streamed into VRAM on the first @spaces.GPU entry. pipe.to("cuda") diffusion_decoder.to("cuda") latent_upsampler.to("cuda") pipe.vae.enable_tiling() AUDIO_SR = pipe.vocoder.config.output_sampling_rate decode_pipe = LTX2VideoDiffusionDecodePipeline( diffusion_decoder=diffusion_decoder, scheduler=pipe.scheduler ) upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler) try: pipe_i2v = LTX2ImageToVideoPipeline( scheduler=pipe.scheduler, vae=pipe.vae, audio_vae=pipe.audio_vae, text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer, connectors=pipe.connectors, transformer=pipe.transformer, vocoder=pipe.vocoder, processor=getattr(pipe, "processor", None), prompt_enhancer=getattr(pipe, "prompt_enhancer", None), duration_head=getattr(pipe, "duration_head", None), ) print(f"[ltx25] ready, t2v + i2v (audio_sr={AUDIO_SR})", flush=True) except Exception as exc: # noqa: BLE001 pipe_i2v = None print(f"[ltx25] ready, t2v ONLY (i2v init failed: {exc!r})", flush=True) def _save_video(path: str) -> dict: """Serialize the mp4 as a JSON pointer the canvas can render (mirrors gradio.workflow._save_tmp; Workflow.launch() allows the tempdir).""" return { "path": path, "url": f"/gradio_api/file={path}", "orig_name": "ltx25.mp4", "mime_type": "video/mp4", } def _estimate_duration(prompt, image_path, height, width, num_frames, seed, decoder, auto_len) -> int: """ZeroGPU checks the *requested* duration against the visitor's remaining quota, so estimate rather than reserve a flat ceiling. Calibrated on the upstream Space's measurements; see https://huggingface.co/docs/hub/en/spaces-zerogpu#duration-management """ frames = AUTO_MAX_FRAMES if auto_len else int(num_frames) px = int(height) * int(width) * frames decode = 20 if decoder == "diffusion" else 8 return int(min(400, max(60, 30 + px / 1_500_000 + decode))) def _friendly_gpu_error(err: Exception) -> str: msg = (str(err) or "").lower() capacity_hints = ( "gpu limit", "reached its gpu limit", "gpu quota", "out of quota", "quota", "no gpu", "could not allocate", "gpu is busy", "too many", "concurrent", ) if any(h in msg for h in capacity_hints): return ( "⛔ The shared ZeroGPU pool is at capacity right now — not a problem " "with your prompt or account. Wait a minute and retry." ) if "out of memory" in msg or "oom" in msg: return "💥 Ran out of GPU memory. Try a smaller resolution or shorter duration." return "⚠️ Video generation failed. Please try again in a moment." @spaces.GPU(duration=_estimate_duration, size="xlarge") def _generate_gpu(prompt, image_path, height, width, num_frames, seed, decoder, auto_len): """Two-stage distilled generation. Runs only under a ZeroGPU allocation.""" if image_path and pipe_i2v is None: raise gr.Error("Image-to-video unavailable; remove the image for text-to-video.") height, width, num_frames = int(height), int(width), int(num_frames) if height % 64 or width % 64: raise gr.Error(f"Two-stage needs height and width divisible by 64 (got {height}x{width}).") active = pipe_i2v if image_path else pipe # ONE generator threaded through both stages, matching the reference. generator = torch.Generator("cuda").manual_seed(int(seed)) requested = None if auto_len else num_frames shared = dict( prompt=prompt, negative_prompt=DEFAULT_NEGATIVE_PROMPT, frame_rate=FRAME_RATE, guidance_scale=GUIDANCE_SCALE, audio_guidance_scale=AUDIO_GUIDANCE_SCALE, # Distilled recipe is a single plain forward; zero every guidance knob. stg_scale=0.0, audio_stg_scale=0.0, modality_scale=1.0, audio_modality_scale=1.0, guidance_rescale=0.0, audio_guidance_rescale=0.0, spatio_temporal_guidance_blocks=None, generator=generator, return_dict=False, ) if image_path: shared["image"] = load_conditioning_image(image_path, width, height) print(f"[gen] stage 1 @ {width // 2}x{height // 2}", flush=True) s1_latents, s1_audio_latents = active( height=height // 2, width=width // 2, num_frames=requested, max_seconds=AUTO_MAX_SECONDS, sigmas=STAGE_1_SIGMAS, output_type="latent", **shared, ) num_frames = (s1_latents.shape[2] - 1) * pipe.vae_temporal_compression_ratio + 1 print("[gen] x2 latent upsample", flush=True) up_latents = upsample_pipe(latents=s1_latents, output_type="latent", return_dict=False)[0] print(f"[gen] stage 2 @ {width}x{height}x{num_frames}", flush=True) want_latents = decoder == "diffusion" s2 = active( num_frames=num_frames, sigmas=STAGE_2_SIGMAS, latents=up_latents, audio_latents=s1_audio_latents, noise_scale=STAGE_2_SIGMAS[0], output_type="latent" if want_latents else "np", **shared, ) if want_latents: video_latents, audio_latents = s2 print("[gen] diffusion decode", flush=True) video = decode_pipe( video_latents, generator=generator, denormalize=False, output_type="np", return_dict=False, )[0] audio_latents = audio_latents.to(pipe.audio_vae.dtype) mel = pipe.audio_vae.decode(audio_latents, return_dict=False)[0] audio = pipe.vocoder(mel) else: video, audio = s2 with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as fh: path = fh.name audio_kwargs = {} if audio is not None: audio_kwargs = dict(audio=audio[0].float().cpu(), audio_sample_rate=AUDIO_SR) encode_video(video[0], fps=FRAME_RATE, output_path=path, **audio_kwargs) return path, num_frames # ----------------------------------------------------------------------------- # Prompt enhancement — same mechanism and Space as the original demo (Gemma-4 on # a separate ZeroGPU Space). Authenticated with this Space's HF_TOKEN, so no # anonymous rate-limit bucket. Client rebuilt past a TTL: the ZeroGPU proxy # token it carries expires. # ----------------------------------------------------------------------------- ENHANCER_SPACE = "diffusers-internal-dev/LTX-2.4-Prompt-Enhancer" _enh = {"client": None, "built_at": 0.0} _ENH_CLIENT_TTL_S = 900 _ENH_CACHE: OrderedDict[tuple[str, str | None], str] = OrderedDict() _ENH_CACHE_MAX = 64 def _image_key(path: str | None) -> str | None: if not path: return None with open(path, "rb") as fh: return hashlib.sha256(fh.read()).hexdigest()[:16] def _enhancer_client(fresh: bool = False): from gradio_client import Client if fresh or _enh["client"] is None or (time.monotonic() - _enh["built_at"]) > _ENH_CLIENT_TTL_S: _enh["client"] = Client( ENHANCER_SPACE, token=HF_TOKEN, httpx_kwargs={"timeout": 300.0} ) _enh["built_at"] = time.monotonic() return _enh["client"] def _extract_path(image_value): """Image ports arrive from the canvas as a {path|url} dict, a path string, or None.""" if isinstance(image_value, dict): return image_value.get("path") or image_value.get("url") or None return image_value or None def enhance_prompt(prompt, image_path=None, do_enhance=True): """CPU step (runs before any GPU is acquired): rewrite the prompt into a detailed LTX-2.5-style caption via the enhancer Space. Falls back to the raw prompt if the enhancer is unavailable. Bound as a `fn` operator.""" if not prompt or not str(prompt).strip(): raise gr.Error("Please enter a prompt.") prompt = str(prompt).strip() image_path = _extract_path(image_path) if not do_enhance: return prompt key = (prompt, _image_key(image_path)) cached = _ENH_CACHE.get(key) if cached is not None: _ENH_CACHE.move_to_end(key) return cached from gradio_client import handle_file args = (prompt, handle_file(image_path) if image_path else None) try: try: enhanced = _enhancer_client().predict(*args, api_name="/enhance") except Exception as first: # noqa: BLE001 # Usually an expired ZeroGPU proxy token; retry once with a fresh client. print(f"[enhance] first attempt failed ({first!r}); fresh client retry", flush=True) enhanced = _enhancer_client(fresh=True).predict(*args, api_name="/enhance") except Exception as e: # noqa: BLE001 print(f"[enhance] failed, using raw prompt: {e!r}", flush=True) gr.Warning(f"Prompt enhancement unavailable ({type(e).__name__}) — using your prompt as written.") return prompt enhanced = enhanced or prompt _ENH_CACHE[key] = enhanced while len(_ENH_CACHE) > _ENH_CACHE_MAX: _ENH_CACHE.popitem(last=False) return enhanced def generate_video(prompt, image_path=None, height=1472, width=832, duration_s=2, seed=42, decoder="conv", auto_len=False, randomize_seed=True): """Workflow-facing wrapper, bound as a `fn` operator. CPU work (seed, frame snapping) happens here so the GPU is not held for it; allocator rejections are reworded into honest user-facing messages. Returns (video_dict, seed_used, length_text). """ if not prompt or not str(prompt).strip(): raise gr.Error("Please enter a prompt.") image_path = _extract_path(image_path) if randomize_seed: seed = random.randint(0, MAX_SEED) num_frames = frames_from_duration(duration_s) try: video_path, frames = _generate_gpu( str(prompt).strip(), image_path, int(height), int(width), num_frames, int(seed), str(decoder), bool(auto_len), ) except gr.Error: raise except Exception as e: # noqa: BLE001 raise gr.Error(_friendly_gpu_error(e)) from e return _save_video(video_path), int(seed), f"{frames} frames ({frames / FRAME_RATE:.2f}s)" demo = gr.Workflow( graph="workflow.json", bind={"enhance_prompt": enhance_prompt, "generate_video": generate_video}, ) if __name__ == "__main__": demo.launch()