run.py

"""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()