LTX-2.5

"""LTX-2.5, generated from a declarative template."""

from __future__ import annotations

import base64
import tempfile

from pydantic import BaseModel, ConfigDict, Field, model_validator

from runware_serverless import serve

MODEL = "Lightricks/LTX-2.5-Diffusers"
REVISION = "4791a1100e0cf97f7ed5c55c12dd71dce13f9dd3"
CACHE = "/weights"
FRAME_RATE = 24
WORKFLOW = "text2video"


class Request(BaseModel):
    # An input this model does not read is a line the caller believed was doing
    # something, so it is refused rather than dropped. Without this, pydantic
    # ignores it: a width sent to a model of fixed size, or a misspelled seed,
    # both come back as a result that silently answered a different request.
    model_config = ConfigDict(extra="forbid")

    prompt: str = Field(description="What to generate, as video with its own audio.")
    negative_prompt: str = Field(default="", description="What to keep out of it.")
    width: int = Field(default=704, ge=256, le=1280, multiple_of=32)
    height: int = Field(default=512, ge=256, le=1280, multiple_of=32)
    min_seconds: float = Field(
        default=1.0,
        ge=1.0,
        le=20.0,
        description="Lower bound on the length the model picks for the prompt.",
    )
    max_seconds: float = Field(
        default=10.0,
        ge=1.0,
        le=20.0,
        description="Upper bound on the length the model picks. The frame count is the model's to choose between these, so it is never one its own stride rejects.",
    )
    seed: int | None = Field(default=None, description="Repeat a result by repeating its seed.")

    @model_validator(mode="after")
    def _open_range(self) -> "Request":
        # The model refuses a collapsed range, since a single length leaves its
        # duration head nothing to choose and generally lands off the VAE's
        # temporal grid. Refused here so the caller is told by the field that is
        # wrong rather than by a 400 from the generation.
        if self.min_seconds >= self.max_seconds:
            raise ValueError(
                "min_seconds must be below max_seconds: the model picks the"
                " length between them and needs a range to pick from"
            )

        return self


class Generated(BaseModel):
    video: str = Field(description="The result as a base64 MP4, with its audio muxed in.")
    frames: int
    seconds: float
    audio_sample_rate: int = Field(description="The rate the muxed track was written at. Reported because writing it at the wrong rate stretches the audio against the picture, and nothing else in the result would show that.")
    seed: int | None


@serve
class GeneratedVideo:
    def load(self) -> None:
        if hasattr(self, "pipeline"):
            return

        # Imported here rather than at the top of the file. The build imports
        # this file with the platform's own dependencies, so nothing above
        # exists until a worker runs.
        import torch
        from diffusers import ModularPipeline

        from diffusers import TorchAoConfig
        from torchao.quantization import Int8WeightOnlyConfig
        from transformers import TorchAoConfig as TransformersTorchAoConfig

        quantization = {
            "transformer": TorchAoConfig(Int8WeightOnlyConfig(version=2)),
            "text_encoder": TransformersTorchAoConfig(
                Int8WeightOnlyConfig(version=2)
            ),
        }

        self.torch = torch

        self.pipeline = ModularPipeline.from_pretrained(
            MODEL, revision=REVISION, cache_dir=CACHE
        )

        # Decode with the convolutional VAE instead of LTX-2.5's own diffusion
        # decoder, which its class docstring prescribes as the way to do this.
        # The diffusion decoder is the better-looking default, and its decode
        # builds a dense seq_len x seq_len FlexAttention mask whose cost cannot
        # be bounded from the request: measured on a 96 GB card it asked for
        # 30.96 GiB at 704x448 and 96.26 GiB at 384x256, so it does not even
        # fall with resolution. An entry cannot publish a size limit against
        # that. The swap trades the decoder for the VAE, one component for one,
        # and a convolution's cost is predictable.
        #
        # Imported by module path because the step has no public export.
        from diffusers.modular_pipelines.ltx2.modular_blocks_ltx2 import (
            LTX2AutoDecoderStep,
        )

        # Assigned through _blocks and not the public blocks property, which
        # returns a deepcopy: a swap through it mutates a throwaway and the
        # pipeline goes on running the decoder it already had, silently. The
        # other public route, LTX25AutoBlocks().init_pipeline(MODEL), takes
        # neither revision nor cache_dir, and an entry does not give up pinning
        # its weights to keep to public attributes.
        #
        # Before load_components, which reads a copy of this: the workflow then
        # resolves against the swapped block and loads the VAE rather than the
        # diffusion decoder.
        self.pipeline._blocks.sub_blocks["decode"] = LTX2AutoDecoderStep()
        # Every component lands on the card as it loads. A worker's host RAM is
        # half its card's VRAM, so a pipeline assembled on the host and moved
        # runs out long before the card does, and the kill arrives as the
        # container's exit 137 rather than as a CUDA error.
        #
        # This workflow and not every declared component, because three of them
        # are 11.7 GB the card then has no room to generate in: the prompt
        # enhancer, which is opt-in and never enabled here, its processor, and
        # the convolutional VAE, which LTX-2.5 replaces with the diffusion
        # decoder. Loading all fifteen leaves 82 GiB resident on a 96 GB card
        # and the first request dies needing 16 GiB more.
        self.pipeline.load_components(
            workflow=WORKFLOW,
            dtype=torch.bfloat16,
            device_map="cuda",
            quantization_config=quantization,
        )

        # load_components reports a component it could not build as a warning
        # and leaves it None, so a wrong argument above does not fail the load:
        # it surfaces one request later as an attribute error on None, far from
        # its cause. Fail here instead, naming what is missing. Scoped to the
        # workflow, because the components it prunes are meant to stay None.
        needed = {
            spec.name
            for spec in self.pipeline.blocks.get_workflow(WORKFLOW).expected_components
        }
        missing = [
            name for name in self.pipeline.null_component_names if name in needed
        ]
        if missing:
            raise RuntimeError(
                "these components did not load, and the warnings above this line"
                " say why: " + ", ".join(missing)
            )

        # The swap above is silent when it fails, so it is checked rather than
        # trusted: the VAE is what this workflow needs only once the decode
        # block is the convolutional one, and a diffusion decoder still loaded
        # means the assignment went to a copy.
        if getattr(self.pipeline, "vae", None) is None:
            raise RuntimeError(
                "the decode block was not swapped: the pipeline loaded"
                " the diffusion decoder, whose decode does not fit this card"
            )

    def predict(self, request: Request) -> Generated:
        from diffusers.utils.export_utils import encode_video

        generator = None
        if request.seed is not None:
            generator = self.torch.Generator().manual_seed(request.seed)

        # num_frames is left out on purpose. The duration head picks a length for
        # the prompt within these bounds, and a count it chose is one its own
        # stride accepts.
        results = self.pipeline(
            prompt=request.prompt,
            negative_prompt=request.negative_prompt or None,
            width=request.width,
            height=request.height,
            frame_rate=FRAME_RATE,
            min_seconds=request.min_seconds,
            max_seconds=request.max_seconds,
            generator=generator,
            output=["videos", "audio"],
        )

        videos, audio = _named(results)

        # Both are indexed, because the pipeline returns a batch and the encoder
        # takes one clip. Dropping the batch dimension on the audio matters more
        # than it looks: the encoder treats any shape carrying a 1 as mono and
        # flattens it, so a stereo [1, 2, N] becomes one interleaved channel of
        # 2N and the track plays at twice the length of the picture.
        clip, track = videos[0], audio[0]

        rate = self.pipeline.vocoder.config.output_sampling_rate

        with tempfile.NamedTemporaryFile(suffix=".mp4") as handle:
            encode_video(
                clip,
                fps=FRAME_RATE,
                output_path=handle.name,
                audio=track,
                # The vocoder's output rate, which is what the waveform is, and
                # not pipeline.audio_sampling_rate: that property reports the
                # audio VAE's mel input rate, 16 kHz against this vocoder's
                # 48 kHz, so the muxed track came out three times the length of
                # the picture. diffusers' own example reads it off the vocoder
                # for the same reason.
                audio_sample_rate=rate,
            )
            handle.seek(0)
            encoded = base64.b64encode(handle.read()).decode("ascii")

        # What came out, not what was asked for: the length was the model's to
        # choose.
        return Generated(
            video=encoded,
            frames=len(clip),
            seconds=len(clip) / FRAME_RATE,
            audio_sample_rate=rate,
            seed=request.seed,
        )


def _named(results: object) -> tuple:
    """The named outputs, however the pipeline hands them back."""
    if isinstance(results, dict):
        return results["videos"], results["audio"]
    if hasattr(results, "videos"):
        return results.videos, results.audio

    return tuple(results)

configVersion: 1
id: ltx-2-5
name: LTX-2.5
source:
  type: code
  entry: model.py
compute:
  gpuType: rtx-pro-6000
build:
  requirements:
    - torch==2.14.1
    - torchvision==0.29.1
    - diffusers==0.42.1
    - transformers==5.21.0
    - accelerate==1.15.0
    - av==16.0.1
    - torchao==0.18.0
volumes:
  - /weights
secrets:
  - name: HF_TOKEN