1	"""Pipeline blocks — each block owns its model lifecycle.
2	Blocks build a model on each ``__call__``, use it, then free GPU memory.
3	This eliminates manual ``del model; cleanup_memory()`` in pipelines and
4	removes the need for :class:`ModelLedger`.
5	"""
6
7	from __future__ import annotations
8
9	import logging
10	from collections.abc import Iterator
11	from contextlib import AbstractContextManager, contextmanager
12	from dataclasses import replace
13	from typing import Callable, TypeVar
14
15	import torch
16
17	from ltx_core.batch_split import BatchSplitAdapter
18	from ltx_core.components.diffusion_steps import EulerDiffusionStep
19	from ltx_core.components.noisers import Noiser
20	from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
21	from ltx_core.components.protocols import DiffusionStepProtocol
22	from ltx_core.layer_streaming import LayerStreamingWrapper
23	from ltx_core.loader import SDOps
24	from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
25	from ltx_core.loader.registry import DummyRegistry, Registry
26	from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
27	from ltx_core.model.audio_vae import (
28	    AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
29	    AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
30	    VOCODER_COMFY_KEYS_FILTER,
31	    AudioDecoderConfigurator,
32	    AudioEncoderConfigurator,
33	    VocoderConfigurator,
34	)
35	from ltx_core.model.audio_vae import (
36	    decode_audio as vae_decode_audio,
37	)
38	from ltx_core.model.transformer import (
39	    LTXV_MODEL_COMFY_RENAMING_MAP,
40	    LTXModelConfigurator,
41	    X0Model,
42	)
43	from ltx_core.model.transformer.compiling import COMPILE_TRANSFORMER, modify_sd_ops_for_compilation
44	from ltx_core.model.upsampler import LatentUpsamplerConfigurator, upsample_video
45	from ltx_core.model.video_vae import (
46	    VAE_DECODER_COMFY_KEYS_FILTER,
47	    VAE_ENCODER_COMFY_KEYS_FILTER,
48	    TilingConfig,
49	    VideoDecoderConfigurator,
50	    VideoEncoder,
51	    VideoEncoderConfigurator,
52	)
53	from ltx_core.quantization import QuantizationPolicy
54	from ltx_core.text_encoders.gemma import (
55	    EMBEDDINGS_PROCESSOR_KEY_OPS,
56	    GEMMA_LLM_KEY_OPS,
57	    GEMMA_MODEL_OPS,
58	    EmbeddingsProcessorConfigurator,
59	    GemmaTextEncoderConfigurator,
60	    module_ops_from_gemma_root,
61	)
62	from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
63	from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
64	from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
65	from ltx_core.utils import find_matching_file
66	from ltx_pipelines.utils.gpu_model import gpu_model
67	from ltx_pipelines.utils.helpers import (
68	    cleanup_memory,
69	    create_noised_state,
70	    generate_enhanced_prompt,
71	)
72	from ltx_pipelines.utils.samplers import euler_denoising_loop
73	from ltx_pipelines.utils.types import Denoiser, ModalitySpec
74
75	logger = logging.getLogger(__name__)
76
77	T = TypeVar("T")
78	_M = TypeVar("_M", bound=torch.nn.Module)
79
80
81	# ---------------------------------------------------------------------------
82	# Internal helpers
83	# ---------------------------------------------------------------------------
84
85
86	@contextmanager
87	def _streaming_model(
88	    model: _M,
89	    layers_attr: str,
90	    target_device: torch.device,
91	    prefetch_count: int,
92	) -> Iterator[_M]:
93	    """Wrap *model* with :class:`LayerStreamingWrapper`, yield it, then tear down."""
94	    wrapped = LayerStreamingWrapper(
95	        model,
96	        layers_attr=layers_attr,
97	        target_device=target_device,
98	        prefetch_count=prefetch_count,
99	    )
100	    try:
101	        yield wrapped  # type: ignore[misc]
102	    finally:
103	        wrapped.teardown()
104	        wrapped.to("meta")
105	        cleanup_memory()
106	        # Flush the host (pinned) memory cache so that freed pinned pages are
107	        # returned to the OS.  Without this, sequential streaming models
108	        # (e.g. text encoder then transformer) exhaust host memory because the
109	        # CachingHostAllocator keeps freed blocks cached indefinitely.
110	        torch.cuda.synchronize(device=target_device)
111	        try:
112	            if hasattr(torch._C, "_host_emptyCache"):
113	                torch._C._host_emptyCache()
114	        except Exception:
115	            logger.warning("Host empty cache cleanup failed; ignoring.", exc_info=True)
116
117
118	def _build_state(
119	    spec: ModalitySpec,
120	    tools: LatentTools,
121	    noiser: Noiser,
122	    dtype: torch.dtype,
123	    device: torch.device,
124	) -> LatentState:
125	    """Create a noised latent state from a modality spec and tools."""
126	    state = create_noised_state(
127	        tools=tools,
128	        conditionings=spec.conditionings,
129	        noiser=noiser,
130	        dtype=dtype,
131	        device=device,
132	        noise_scale=spec.noise_scale,
133	        initial_latent=spec.initial_latent,
134	    )
135	    if spec.frozen:
136	        state = replace(state, denoise_mask=torch.zeros_like(state.denoise_mask))
137	    return state
138
139
140	def _cleanup_iter(it: Iterator[torch.Tensor], model: torch.nn.Module) -> Iterator[torch.Tensor]:
141	    """Wrap an iterator to clean up *model* memory once it is exhausted or abandoned."""
142	    with gpu_model(model):
143	        yield from it
144
145
146	# ---------------------------------------------------------------------------
147	# DiffusionStage
148	# ---------------------------------------------------------------------------
149
150
151	class DiffusionStage:
152	    """Owns transformer lifecycle. Builds on each call, frees on exit.
153	    Replaces the manual ``model_ledger.transformer()`` / ``del transformer``
154	    pattern in every pipeline.
155	    """
156
157	    def __init__(
158	        self,
159	        checkpoint_path: str,
160	        dtype: torch.dtype,
161	        device: torch.device,
162	        loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
163	        quantization: QuantizationPolicy | None = None,
164	        registry: Registry | None = None,
165	        torch_compile: bool = False,
166	    ) -> None:
167	        self._dtype = dtype
168	        self._device = device
169	        self._quantization = quantization
170	        self._torch_compile = torch_compile
171	        self._transformer_builder = Builder(
172	            model_path=checkpoint_path,
173	            model_class_configurator=LTXModelConfigurator,
174	            model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
175	            loras=tuple(loras),
176	            registry=registry or DummyRegistry(),
177	        )
178
179	    def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
180	        target = device or self._device
181	        sd_ops = self._transformer_builder.model_sd_ops
182	        module_ops = self._transformer_builder.module_ops
183	        loras = self._transformer_builder.loras
184	        if self._torch_compile:
185	            module_ops = (*module_ops, COMPILE_TRANSFORMER)
186	            number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
187	            sd_ops = modify_sd_ops_for_compilation(sd_ops, number_of_layers)
188	            loras = tuple(
189	                LoraPathStrengthAndSDOps(
190	                    lora.path,
191	                    lora.strength,
192	                    modify_sd_ops_for_compilation(
193	                        lora.sd_ops if lora.sd_ops is not None else SDOps(name="identity"), number_of_layers
194	                    ),
195	                )
196	                for lora in loras
197	            )
198	        if self._quantization is not None:
199	            module_ops = (*module_ops, *self._quantization.module_ops)
200	            sd_ops = SDOps(
201	                name=f"sd_ops_chain_{sd_ops.name}+{self._quantization.sd_ops.name}",
202	                mapping=(*sd_ops.mapping, *self._quantization.sd_ops.mapping),
203	            )
204
205	        builder = self._transformer_builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
206	        return X0Model(builder.build(device=target, **kwargs)).to(target).eval()
207
208	    def _transformer_ctx(
209	        self,
210	        streaming_prefetch_count: int | None,
211	        **kwargs: object,
212	    ) -> AbstractContextManager:
213	        if streaming_prefetch_count is not None:
214	            return _streaming_model(
215	                self._build_transformer(device=torch.device("cpu"), **kwargs),
216	                layers_attr="velocity_model.transformer_blocks",
217	                target_device=self._device,
218	                prefetch_count=streaming_prefetch_count,
219	            )
220	        return gpu_model(self._build_transformer(**kwargs))
221
222	    def __call__(  # noqa: PLR0913
223	        self,
224	        denoiser: Denoiser,
225	        sigmas: torch.Tensor,
226	        noiser: Noiser,
227	        width: int,
228	        height: int,
229	        frames: int,
230	        fps: float,
231	        video: ModalitySpec | None = None,
232	        audio: ModalitySpec | None = None,
233	        stepper: DiffusionStepProtocol | None = None,
234	        loop: Callable[..., tuple[LatentState | None, LatentState | None]] | None = None,
235	        streaming_prefetch_count: int | None = None,
236	        max_batch_size: int = 1,
237	    ) -> tuple[LatentState | None, LatentState | None]:
238	        """Build transformer → run denoising loop → free transformer.
239	        Args:
240	            width: Output width in pixels.
241	            height: Output height in pixels.
242	            frames: Number of output frames.
243	            fps: Frame rate.
244	            loop: Denoising loop function. Must accept
245	                ``(sigmas, video_state, audio_state, stepper, transformer, denoiser)``
246	                as the first six positional arguments. When ``None``, resolves to
247	                :func:`euler_denoising_loop` at call time.
248	            streaming_prefetch_count: When set, build the transformer on CPU and
249	                wrap with :class:`LayerStreamingWrapper` for memory-efficient
250	                inference, prefetching this many layers ahead.
251	            max_batch_size: Maximum batch size per transformer forward pass.
252	                Guided denoisers make up to 4 transformer calls per step.
253	                When set to a value > 1, the transformer batches multiple
254	                calls together, reducing layer-streaming PCIe transfers.
255	                Default ``1`` preserves sequential behavior.
256	        Returns ``(video_state | None, audio_state | None)`` with cleared
257	        conditionings and unpatchified latents for present modalities.
258	        """
259	        if video is None and audio is None:
260	            raise ValueError("At least one of `video` or `audio` must be provided")
261
262	        if loop is None:
263	            loop = euler_denoising_loop
264
265	        if stepper is None:
266	            stepper = EulerDiffusionStep()
267
268	        pixel_shape = VideoPixelShape(batch=1, frames=frames, height=height, width=width, fps=fps)
269
270	        video_state: LatentState | None = None
271	        video_tools: LatentTools | None = None
272	        if video is not None:
273	            v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
274	            video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, fps)
275	            video_state = _build_state(video, video_tools, noiser, self._dtype, self._device)
276
277	        audio_state: LatentState | None = None
278	        audio_tools: LatentTools | None = None
279	        if audio is not None:
280	            a_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape)
281	            audio_tools = AudioLatentTools(AudioPatchifier(patch_size=1), a_shape)
282	            audio_state = _build_state(audio, audio_tools, noiser, self._dtype, self._device)
283
284	        with self._transformer_ctx(streaming_prefetch_count, video_tools=video_tools) as base_transformer:
285	            transformer = BatchSplitAdapter(base_transformer, max_batch_size=max_batch_size)
286	            video_state, audio_state = loop(
287	                sigmas=sigmas,
288	                video_state=video_state,
289	                audio_state=audio_state,
290	                stepper=stepper,
291	                transformer=transformer,
292	                denoiser=denoiser,
293	            )
294
295	        # Post-process: clear conditionings and unpatchify
296	        if video_state is not None and video_tools is not None:
297	            video_state = video_tools.clear_conditioning(video_state)
298	            video_state = video_tools.unpatchify(video_state)
299
300	        if audio_state is not None and audio_tools is not None:
301	            audio_state = audio_tools.clear_conditioning(audio_state)
302	            audio_state = audio_tools.unpatchify(audio_state)
303
304	        return video_state, audio_state
305
306
307	# ---------------------------------------------------------------------------
308	# PromptEncoder
309	# ---------------------------------------------------------------------------
310
311
312	class PromptEncoder:
313	    """Owns text encoder + embeddings processor lifecycle.
314	    Loads Gemma, encodes prompts, frees Gemma, then loads the embeddings
315	    processor to produce final outputs.
316	    """
317
318	    def __init__(
319	        self,
320	        checkpoint_path: str,
321	        gemma_root: str,
322	        dtype: torch.dtype,
323	        device: torch.device,
324	        registry: Registry | None = None,
325	    ) -> None:
326	        self._dtype = dtype
327	        self._device = device
328
329	        module_ops = module_ops_from_gemma_root(gemma_root)
330	        model_folder = find_matching_file(gemma_root, "model*.safetensors").parent
331	        weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
332
333	        self._text_encoder_builder = Builder(
334	            model_path=tuple(weight_paths),
335	            model_class_configurator=GemmaTextEncoderConfigurator,
336	            model_sd_ops=GEMMA_LLM_KEY_OPS,
337	            module_ops=(GEMMA_MODEL_OPS, *module_ops),
338	            registry=registry or DummyRegistry(),
339	        )
340	        self._embeddings_processor_builder = Builder(
341	            model_path=checkpoint_path,
342	            model_class_configurator=EmbeddingsProcessorConfigurator,
343	            model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS,
344	            registry=registry or DummyRegistry(),
345	        )
346
347	    def _text_encoder_ctx(
348	        self,
349	        streaming_prefetch_count: int | None,
350	    ) -> AbstractContextManager:
351	        if streaming_prefetch_count is not None:
352	            return _streaming_model(
353	                self._text_encoder_builder.build(device=torch.device("cpu"), dtype=self._dtype).eval(),
354	                layers_attr="model.model.language_model.layers",
355	                target_device=self._device,
356	                prefetch_count=streaming_prefetch_count,
357	            )
358	        return gpu_model(self._text_encoder_builder.build(device=self._device, dtype=self._dtype).eval())
359
360	    def __call__(
361	        self,
362	        prompts: list[str],
363	        *,
364	        enhance_first_prompt: bool = False,
365	        enhance_prompt_image: str | None = None,
366	        enhance_prompt_seed: int = 42,
367	        streaming_prefetch_count: int | None = None,
368	    ) -> list[EmbeddingsProcessorOutput]:
369	        """Encode *prompts* through Gemma → embeddings processor, freeing each model after use."""
370	        with self._text_encoder_ctx(streaming_prefetch_count) as text_encoder:
371	            if enhance_first_prompt:
372	                prompts = list(prompts)
373	                prompts[0] = generate_enhanced_prompt(
374	                    text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed
375	                )
376	            raw_outputs = [text_encoder.encode(p) for p in prompts]
377
378	        with gpu_model(
379	            self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
380	        ) as embeddings_processor:
381	            return [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
382
383
384	# ---------------------------------------------------------------------------
385	# ImageConditioner
386	# ---------------------------------------------------------------------------
387
388
389	class ImageConditioner:
390	    """Owns video encoder lifecycle.
391	    Builds the encoder, passes it to the user-supplied callable, then frees it.
392	    """
393
394	    def __init__(
395	        self,
396	        checkpoint_path: str,
397	        dtype: torch.dtype,
398	        device: torch.device,
399	        registry: Registry | None = None,
400	    ) -> None:
401	        self._dtype = dtype
402	        self._device = device
403	        self._encoder_builder = Builder(
404	            model_path=checkpoint_path,
405	            model_class_configurator=VideoEncoderConfigurator,
406	            model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
407	            registry=registry or DummyRegistry(),
408	        )
409
410	    def _build_encoder(self) -> VideoEncoder:
411	        return self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
412
413	    def __call__(self, fn: Callable[[VideoEncoder], T]) -> T:
414	        """Build video encoder → call *fn(encoder)* → free encoder."""
415	        with gpu_model(self._build_encoder()) as encoder:
416	            return fn(encoder)
417
418
419	# ---------------------------------------------------------------------------
420	# VideoUpsampler
421	# ---------------------------------------------------------------------------
422
423
424	class VideoUpsampler:
425	    """Owns video encoder + spatial upsampler lifecycle."""
426
427	    def __init__(
428	        self,
429	        checkpoint_path: str,
430	        upsampler_path: str,
431	        dtype: torch.dtype,
432	        device: torch.device,
433	        registry: Registry | None = None,
434	    ) -> None:
435	        self._dtype = dtype
436	        self._device = device
437	        self._encoder_builder = Builder(
438	            model_path=checkpoint_path,
439	            model_class_configurator=VideoEncoderConfigurator,
440	            model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
441	            registry=registry or DummyRegistry(),
442	        )
443	        self._upsampler_builder = Builder(
444	            model_path=upsampler_path,
445	            model_class_configurator=LatentUpsamplerConfigurator,
446	            registry=registry or DummyRegistry(),
447	        )
448
449	    def __call__(self, latent: torch.Tensor) -> torch.Tensor:
450	        """Upsample *latent* using video encoder + spatial upsampler, then free both."""
451	        with (
452	            gpu_model(
453	                self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
454	            ) as encoder,
455	            gpu_model(
456	                self._upsampler_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
457	            ) as upsampler,
458	        ):
459	            return upsample_video(latent=latent, video_encoder=encoder, upsampler=upsampler)
460
461
462	# ---------------------------------------------------------------------------
463	# VideoDecoder
464	# ---------------------------------------------------------------------------
465
466
467	class VideoDecoder:
468	    """Owns video decoder lifecycle.
469	    Returns an iterator that cleans up the decoder after all chunks are consumed.
470	    """
471
472	    def __init__(
473	        self,
474	        checkpoint_path: str,
475	        dtype: torch.dtype,
476	        device: torch.device,
477	        registry: Registry | None = None,
478	    ) -> None:
479	        self._dtype = dtype
480	        self._device = device
481	        self._decoder_builder = Builder(
482	            model_path=checkpoint_path,
483	            model_class_configurator=VideoDecoderConfigurator,
484	            model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
485	            registry=registry or DummyRegistry(),
486	        )
487
488	    def __call__(
489	        self,
490	        latent: torch.Tensor,
491	        tiling_config: TilingConfig | None = None,
492	        generator: torch.Generator | None = None,
493	    ) -> Iterator[torch.Tensor]:
494	        """Decode *latent* to pixel-space video chunks. Decoder freed after exhaustion."""
495	        decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
496	        return _cleanup_iter(decoder.decode_video(latent, tiling_config, generator), decoder)
497
498
499	# ---------------------------------------------------------------------------
500	# AudioDecoder
501	# ---------------------------------------------------------------------------
502
503
504	class AudioDecoder:
505	    """Owns audio decoder + vocoder lifecycle."""
506
507	    def __init__(
508	        self,
509	        checkpoint_path: str,
510	        dtype: torch.dtype,
511	        device: torch.device,
512	        registry: Registry | None = None,
513	    ) -> None:
514	        self._dtype = dtype
515	        self._device = device
516	        self._decoder_builder = Builder(
517	            model_path=checkpoint_path,
518	            model_class_configurator=AudioDecoderConfigurator,
519	            model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
520	            registry=registry or DummyRegistry(),
521	        )
522	        self._vocoder_builder = Builder(
523	            model_path=checkpoint_path,
524	            model_class_configurator=VocoderConfigurator,
525	            model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
526	            registry=registry or DummyRegistry(),
527	        )
528
529	    def __call__(self, latent: torch.Tensor) -> Audio:
530	        """Decode audio *latent* through VAE decoder + vocoder, then free both."""
531	        with (
532	            gpu_model(
533	                self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
534	            ) as decoder,
535	            gpu_model(
536	                self._vocoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
537	            ) as vocoder,
538	        ):
539	            return vae_decode_audio(latent, decoder, vocoder)
540
541
542	# ---------------------------------------------------------------------------
543	# AudioEncoder
544	# ---------------------------------------------------------------------------
545
546
547	class AudioConditioner:
548	    """Owns audio encoder lifecycle.
549	    Builds the encoder, passes it to the user-supplied callable, then frees it.
550	    Mirrors :class:`ImageConditioner` for the audio modality.
551	    """
552
553	    def __init__(
554	        self,
555	        checkpoint_path: str,
556	        dtype: torch.dtype,
557	        device: torch.device,
558	        registry: Registry | None = None,
559	    ) -> None:
560	        self._dtype = dtype
561	        self._device = device
562	        self._encoder_builder = Builder(
563	            model_path=checkpoint_path,
564	            model_class_configurator=AudioEncoderConfigurator,
565	            model_sd_ops=AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
566	            registry=registry or DummyRegistry(),
567	        )
568
569	    def __call__(self, fn: Callable[[torch.nn.Module], T]) -> T:
570	        """Build audio encoder → call *fn(encoder)* → free encoder."""
571	        with gpu_model(
572	            self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
573	        ) as encoder:
574	            return fn(encoder)
575

"""Pipeline blocks — each block owns its model lifecycle.
Blocks build a model on each ``__call__``, use it, then free GPU memory.
This eliminates manual ``del model; cleanup_memory()`` in pipelines and
removes the need for :class:`ModelLedger`.
"""

from __future__ import annotations

import logging
from collections.abc import Iterator
from contextlib import AbstractContextManager, contextmanager
from dataclasses import replace
from typing import Callable, TypeVar

import torch

from ltx_core.batch_split import BatchSplitAdapter
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.noisers import Noiser
from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.layer_streaming import LayerStreamingWrapper
from ltx_core.loader import SDOps
from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import DummyRegistry, Registry
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
from ltx_core.model.audio_vae import (
    AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
    AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
    VOCODER_COMFY_KEYS_FILTER,
    AudioDecoderConfigurator,
    AudioEncoderConfigurator,
    VocoderConfigurator,
)
from ltx_core.model.audio_vae import (
    decode_audio as vae_decode_audio,
)
from ltx_core.model.transformer import (
    LTXV_MODEL_COMFY_RENAMING_MAP,
    LTXModelConfigurator,
    X0Model,
)
from ltx_core.model.transformer.compiling import COMPILE_TRANSFORMER, modify_sd_ops_for_compilation
from ltx_core.model.upsampler import LatentUpsamplerConfigurator, upsample_video
from ltx_core.model.video_vae import (
    VAE_DECODER_COMFY_KEYS_FILTER,
    VAE_ENCODER_COMFY_KEYS_FILTER,
    TilingConfig,
    VideoDecoderConfigurator,
    VideoEncoder,
    VideoEncoderConfigurator,
)
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import (
    EMBEDDINGS_PROCESSOR_KEY_OPS,
    GEMMA_LLM_KEY_OPS,
    GEMMA_MODEL_OPS,
    EmbeddingsProcessorConfigurator,
    GemmaTextEncoderConfigurator,
    module_ops_from_gemma_root,
)
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
from ltx_core.utils import find_matching_file
from ltx_pipelines.utils.gpu_model import gpu_model
from ltx_pipelines.utils.helpers import (
    cleanup_memory,
    create_noised_state,
    generate_enhanced_prompt,
)
from ltx_pipelines.utils.samplers import euler_denoising_loop
from ltx_pipelines.utils.types import Denoiser, ModalitySpec

logger = logging.getLogger(__name__)

T = TypeVar("T")
_M = TypeVar("_M", bound=torch.nn.Module)

# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

@contextmanager
def _streaming_model(
    model: _M,
    layers_attr: str,
    target_device: torch.device,
    prefetch_count: int,
) -> Iterator[_M]:
    """Wrap *model* with :class:`LayerStreamingWrapper`, yield it, then tear down."""
    wrapped = LayerStreamingWrapper(
        model,
        layers_attr=layers_attr,
        target_device=target_device,
        prefetch_count=prefetch_count,
    )
    try:
        yield wrapped  # type: ignore[misc]
    finally:
        wrapped.teardown()
        wrapped.to("meta")
        cleanup_memory()
        # Flush the host (pinned) memory cache so that freed pinned pages are
        # returned to the OS.  Without this, sequential streaming models
        # (e.g. text encoder then transformer) exhaust host memory because the
        # CachingHostAllocator keeps freed blocks cached indefinitely.
        torch.cuda.synchronize(device=target_device)
        try:
            if hasattr(torch._C, "_host_emptyCache"):
                torch._C._host_emptyCache()
        except Exception:
            logger.warning("Host empty cache cleanup failed; ignoring.", exc_info=True)

def _build_state(
    spec: ModalitySpec,
    tools: LatentTools,
    noiser: Noiser,
    dtype: torch.dtype,
    device: torch.device,
) -> LatentState:
    """Create a noised latent state from a modality spec and tools."""
    state = create_noised_state(
        tools=tools,
        conditionings=spec.conditionings,
        noiser=noiser,
        dtype=dtype,
        device=device,
        noise_scale=spec.noise_scale,
        initial_latent=spec.initial_latent,
    )
    if spec.frozen:
        state = replace(state, denoise_mask=torch.zeros_like(state.denoise_mask))
    return state

def _cleanup_iter(it: Iterator[torch.Tensor], model: torch.nn.Module) -> Iterator[torch.Tensor]:
    """Wrap an iterator to clean up *model* memory once it is exhausted or abandoned."""
    with gpu_model(model):
        yield from it

# ---------------------------------------------------------------------------
# DiffusionStage
# ---------------------------------------------------------------------------

class DiffusionStage:
    """Owns transformer lifecycle. Builds on each call, frees on exit.
    Replaces the manual ``model_ledger.transformer()`` / ``del transformer``
    pattern in every pipeline.
    """

    def __init__(
        self,
        checkpoint_path: str,
        dtype: torch.dtype,
        device: torch.device,
        loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
        quantization: QuantizationPolicy | None = None,
        registry: Registry | None = None,
        torch_compile: bool = False,
    ) -> None:
        self._dtype = dtype
        self._device = device
        self._quantization = quantization
        self._torch_compile = torch_compile
        self._transformer_builder = Builder(
            model_path=checkpoint_path,
            model_class_configurator=LTXModelConfigurator,
            model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
            loras=tuple(loras),
            registry=registry or DummyRegistry(),
        )

    def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
        target = device or self._device
        sd_ops = self._transformer_builder.model_sd_ops
        module_ops = self._transformer_builder.module_ops
        loras = self._transformer_builder.loras
        if self._torch_compile:
            module_ops = (*module_ops, COMPILE_TRANSFORMER)
            number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
            sd_ops = modify_sd_ops_for_compilation(sd_ops, number_of_layers)
            loras = tuple(
                LoraPathStrengthAndSDOps(
                    lora.path,
                    lora.strength,
                    modify_sd_ops_for_compilation(
                        lora.sd_ops if lora.sd_ops is not None else SDOps(name="identity"), number_of_layers
                    ),
                )
                for lora in loras
            )
        if self._quantization is not None:
            module_ops = (*module_ops, *self._quantization.module_ops)
            sd_ops = SDOps(
                name=f"sd_ops_chain_{sd_ops.name}+{self._quantization.sd_ops.name}",
                mapping=(*sd_ops.mapping, *self._quantization.sd_ops.mapping),
            )

        builder = self._transformer_builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
        return X0Model(builder.build(device=target, **kwargs)).to(target).eval()

    def _transformer_ctx(
        self,
        streaming_prefetch_count: int | None,
        **kwargs: object,
    ) -> AbstractContextManager:
        if streaming_prefetch_count is not None:
            return _streaming_model(
                self._build_transformer(device=torch.device("cpu"), **kwargs),
                layers_attr="velocity_model.transformer_blocks",
                target_device=self._device,
                prefetch_count=streaming_prefetch_count,
            )
        return gpu_model(self._build_transformer(**kwargs))

    def __call__(  # noqa: PLR0913
        self,
        denoiser: Denoiser,
        sigmas: torch.Tensor,
        noiser: Noiser,
        width: int,
        height: int,
        frames: int,
        fps: float,
        video: ModalitySpec | None = None,
        audio: ModalitySpec | None = None,
        stepper: DiffusionStepProtocol | None = None,
        loop: Callable[..., tuple[LatentState | None, LatentState | None]] | None = None,
        streaming_prefetch_count: int | None = None,
        max_batch_size: int = 1,
    ) -> tuple[LatentState | None, LatentState | None]:
        """Build transformer → run denoising loop → free transformer.
        Args:
            width: Output width in pixels.
            height: Output height in pixels.
            frames: Number of output frames.
            fps: Frame rate.
            loop: Denoising loop function. Must accept
                ``(sigmas, video_state, audio_state, stepper, transformer, denoiser)``
                as the first six positional arguments. When ``None``, resolves to
                :func:`euler_denoising_loop` at call time.
            streaming_prefetch_count: When set, build the transformer on CPU and
                wrap with :class:`LayerStreamingWrapper` for memory-efficient
                inference, prefetching this many layers ahead.
            max_batch_size: Maximum batch size per transformer forward pass.
                Guided denoisers make up to 4 transformer calls per step.
                When set to a value > 1, the transformer batches multiple
                calls together, reducing layer-streaming PCIe transfers.
                Default ``1`` preserves sequential behavior.
        Returns ``(video_state | None, audio_state | None)`` with cleared
        conditionings and unpatchified latents for present modalities.
        """
        if video is None and audio is None:
            raise ValueError("At least one of `video` or `audio` must be provided")

        if loop is None:
            loop = euler_denoising_loop

        if stepper is None:
            stepper = EulerDiffusionStep()

        pixel_shape = VideoPixelShape(batch=1, frames=frames, height=height, width=width, fps=fps)

        video_state: LatentState | None = None
        video_tools: LatentTools | None = None
        if video is not None:
            v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
            video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, fps)
            video_state = _build_state(video, video_tools, noiser, self._dtype, self._device)

        audio_state: LatentState | None = None
        audio_tools: LatentTools | None = None
        if audio is not None:
            a_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape)
            audio_tools = AudioLatentTools(AudioPatchifier(patch_size=1), a_shape)
            audio_state = _build_state(audio, audio_tools, noiser, self._dtype, self._device)

        with self._transformer_ctx(streaming_prefetch_count, video_tools=video_tools) as base_transformer:
            transformer = BatchSplitAdapter(base_transformer, max_batch_size=max_batch_size)
            video_state, audio_state = loop(
                sigmas=sigmas,
                video_state=video_state,
                audio_state=audio_state,
                stepper=stepper,
                transformer=transformer,
                denoiser=denoiser,
            )

        # Post-process: clear conditionings and unpatchify
        if video_state is not None and video_tools is not None:
            video_state = video_tools.clear_conditioning(video_state)
            video_state = video_tools.unpatchify(video_state)

        if audio_state is not None and audio_tools is not None:
            audio_state = audio_tools.clear_conditioning(audio_state)
            audio_state = audio_tools.unpatchify(audio_state)

        return video_state, audio_state

# ---------------------------------------------------------------------------
# PromptEncoder
# ---------------------------------------------------------------------------

class PromptEncoder:
    """Owns text encoder + embeddings processor lifecycle.
    Loads Gemma, encodes prompts, frees Gemma, then loads the embeddings
    processor to produce final outputs.
    """

    def __init__(
        self,
        checkpoint_path: str,
        gemma_root: str,
        dtype: torch.dtype,
        device: torch.device,
        registry: Registry | None = None,
    ) -> None:
        self._dtype = dtype
        self._device = device

        module_ops = module_ops_from_gemma_root(gemma_root)
        model_folder = find_matching_file(gemma_root, "model*.safetensors").parent
        weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]

        self._text_encoder_builder = Builder(
            model_path=tuple(weight_paths),
            model_class_configurator=GemmaTextEncoderConfigurator,
            model_sd_ops=GEMMA_LLM_KEY_OPS,
            module_ops=(GEMMA_MODEL_OPS, *module_ops),
            registry=registry or DummyRegistry(),
        )
        self._embeddings_processor_builder = Builder(
            model_path=checkpoint_path,
            model_class_configurator=EmbeddingsProcessorConfigurator,
            model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS,
            registry=registry or DummyRegistry(),
        )

    def _text_encoder_ctx(
        self,
        streaming_prefetch_count: int | None,
    ) -> AbstractContextManager:
        if streaming_prefetch_count is not None:
            return _streaming_model(
                self._text_encoder_builder.build(device=torch.device("cpu"), dtype=self._dtype).eval(),
                layers_attr="model.model.language_model.layers",
                target_device=self._device,
                prefetch_count=streaming_prefetch_count,
            )
        return gpu_model(self._text_encoder_builder.build(device=self._device, dtype=self._dtype).eval())

    def __call__(
        self,
        prompts: list[str],
        *,
        enhance_first_prompt: bool = False,
        enhance_prompt_image: str | None = None,
        enhance_prompt_seed: int = 42,
        streaming_prefetch_count: int | None = None,
    ) -> list[EmbeddingsProcessorOutput]:
        """Encode *prompts* through Gemma → embeddings processor, freeing each model after use."""
        with self._text_encoder_ctx(streaming_prefetch_count) as text_encoder:
            if enhance_first_prompt:
                prompts = list(prompts)
                prompts[0] = generate_enhanced_prompt(
                    text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed
                )
            raw_outputs = [text_encoder.encode(p) for p in prompts]

        with gpu_model(
            self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
        ) as embeddings_processor:
            return [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]

# ---------------------------------------------------------------------------
# ImageConditioner
# ---------------------------------------------------------------------------

class ImageConditioner:
    """Owns video encoder lifecycle.
    Builds the encoder, passes it to the user-supplied callable, then frees it.
    """

    def __init__(
        self,
        checkpoint_path: str,
        dtype: torch.dtype,
        device: torch.device,
        registry: Registry | None = None,
    ) -> None:
        self._dtype = dtype
        self._device = device
        self._encoder_builder = Builder(
            model_path=checkpoint_path,
            model_class_configurator=VideoEncoderConfigurator,
            model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
            registry=registry or DummyRegistry(),
        )

    def _build_encoder(self) -> VideoEncoder:
        return self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()

    def __call__(self, fn: Callable[[VideoEncoder], T]) -> T:
        """Build video encoder → call *fn(encoder)* → free encoder."""
        with gpu_model(self._build_encoder()) as encoder:
            return fn(encoder)

# ---------------------------------------------------------------------------
# VideoUpsampler
# ---------------------------------------------------------------------------

class VideoUpsampler:
    """Owns video encoder + spatial upsampler lifecycle."""

    def __init__(
        self,
        checkpoint_path: str,
        upsampler_path: str,
        dtype: torch.dtype,
        device: torch.device,
        registry: Registry | None = None,
    ) -> None:
        self._dtype = dtype
        self._device = device
        self._encoder_builder = Builder(
            model_path=checkpoint_path,
            model_class_configurator=VideoEncoderConfigurator,
            model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
            registry=registry or DummyRegistry(),
        )
        self._upsampler_builder = Builder(
            model_path=upsampler_path,
            model_class_configurator=LatentUpsamplerConfigurator,
            registry=registry or DummyRegistry(),
        )

    def __call__(self, latent: torch.Tensor) -> torch.Tensor:
        """Upsample *latent* using video encoder + spatial upsampler, then free both."""
        with (
            gpu_model(
                self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
            ) as encoder,
            gpu_model(
                self._upsampler_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
            ) as upsampler,
        ):
            return upsample_video(latent=latent, video_encoder=encoder, upsampler=upsampler)

# ---------------------------------------------------------------------------
# VideoDecoder
# ---------------------------------------------------------------------------

class VideoDecoder:
    """Owns video decoder lifecycle.
    Returns an iterator that cleans up the decoder after all chunks are consumed.
    """

    def __init__(
        self,
        checkpoint_path: str,
        dtype: torch.dtype,
        device: torch.device,
        registry: Registry | None = None,
    ) -> None:
        self._dtype = dtype
        self._device = device
        self._decoder_builder = Builder(
            model_path=checkpoint_path,
            model_class_configurator=VideoDecoderConfigurator,
            model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
            registry=registry or DummyRegistry(),
        )

    def __call__(
        self,
        latent: torch.Tensor,
        tiling_config: TilingConfig | None = None,
        generator: torch.Generator | None = None,
    ) -> Iterator[torch.Tensor]:
        """Decode *latent* to pixel-space video chunks. Decoder freed after exhaustion."""
        decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
        return _cleanup_iter(decoder.decode_video(latent, tiling_config, generator), decoder)

# ---------------------------------------------------------------------------
# AudioDecoder
# ---------------------------------------------------------------------------

class AudioDecoder:
    """Owns audio decoder + vocoder lifecycle."""

    def __init__(
        self,
        checkpoint_path: str,
        dtype: torch.dtype,
        device: torch.device,
        registry: Registry | None = None,
    ) -> None:
        self._dtype = dtype
        self._device = device
        self._decoder_builder = Builder(
            model_path=checkpoint_path,
            model_class_configurator=AudioDecoderConfigurator,
            model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
            registry=registry or DummyRegistry(),
        )
        self._vocoder_builder = Builder(
            model_path=checkpoint_path,
            model_class_configurator=VocoderConfigurator,
            model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
            registry=registry or DummyRegistry(),
        )

    def __call__(self, latent: torch.Tensor) -> Audio:
        """Decode audio *latent* through VAE decoder + vocoder, then free both."""
        with (
            gpu_model(
                self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
            ) as decoder,
            gpu_model(
                self._vocoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
            ) as vocoder,
        ):
            return vae_decode_audio(latent, decoder, vocoder)

# ---------------------------------------------------------------------------
# AudioEncoder
# ---------------------------------------------------------------------------

class AudioConditioner:
    """Owns audio encoder lifecycle.
    Builds the encoder, passes it to the user-supplied callable, then frees it.
    Mirrors :class:`ImageConditioner` for the audio modality.
    """

    def __init__(
        self,
        checkpoint_path: str,
        dtype: torch.dtype,
        device: torch.device,
        registry: Registry | None = None,
    ) -> None:
        self._dtype = dtype
        self._device = device
        self._encoder_builder = Builder(
            model_path=checkpoint_path,
            model_class_configurator=AudioEncoderConfigurator,
            model_sd_ops=AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
            registry=registry or DummyRegistry(),
        )

    def __call__(self, fn: Callable[[torch.nn.Module], T]) -> T:
        """Build audio encoder → call *fn(encoder)* → free encoder."""
        with gpu_model(
            self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
        ) as encoder:
            return fn(encoder)


There it is. Line 1-2 says it all:

> **"Blocks build a model on each `__call__`, use it, then free GPU memory."**

Every single block (transformer, text encoder, VAE decoder, upsampler, etc.) rebuilds from disk on every generation call. That's by design for memory management but terrible for latency.

Let me check the `gpu_model` context manager:


1	from collections.abc import Iterator
2	from contextlib import contextmanager
3	from typing import TypeVar
4
5	import torch
6
7	from ltx_pipelines.utils.helpers import cleanup_memory
8
9	_M = TypeVar("_M", bound=torch.nn.Module)
10
11
12	@contextmanager
13	def gpu_model(model: _M) -> Iterator[_M]:
14	    """Context manager that yields a model and releases its memory on exit.
15	    Moves all parameters and buffers to ``meta`` device on exit, which
16	    immediately releases the underlying storage on **both** GPU and CPU,
17	    then runs ``cleanup_memory()`` to reclaim fragmented CUDA memory.
18	    Usage::
19	        with gpu_model(build_encoder()) as encoder:
20	            ...  # use encoder — typed as the concrete class
21	        # GPU + CPU memory freed automatically
22	    """
23	    try:
24	        yield model
25	    finally:
26	        torch.cuda.synchronize()
27	        # .to("meta") releases storage for all parameters/buffers regardless
28	        # of their original device (CUDA or CPU).
29	        model.to("meta")
30	        cleanup_memory()
31

from collections.abc import Iterator
from contextlib import contextmanager
from typing import TypeVar

import torch

from ltx_pipelines.utils.helpers import cleanup_memory

_M = TypeVar("_M", bound=torch.nn.Module)

@contextmanager
def gpu_model(model: _M) -> Iterator[_M]:
    """Context manager that yields a model and releases its memory on exit.
    Moves all parameters and buffers to ``meta`` device on exit, which
    immediately releases the underlying storage on **both** GPU and CPU,
    then runs ``cleanup_memory()`` to reclaim fragmented CUDA memory.
    Usage::
        with gpu_model(build_encoder()) as encoder:
            ...  # use encoder — typed as the concrete class
        # GPU + CPU memory freed automatically
    """
    try:
        yield model
    finally:
        torch.cuda.synchronize()
        # .to("meta") releases storage for all parameters/buffers regardless
        # of their original device (CUDA or CPU).
        model.to("meta")
        cleanup_memory()


Yep — every component is loaded from disk, used, moved to `meta` device, and garbage collected. For a single call pipeline flow:

1. Load Gemma text encoder (~12GB) → encode prompt → free
2. Load video encoder → encode image conditioning → free
3. Load transformer (~28GB FP8) → run stage 1 denoising → free
4. Load video encoder + upsampler → upsample → free
5. Load transformer again → run stage 2 → free
6. Load video decoder → decode → free
7. Load audio decoder + vocoder → decode → free

That's **7 model loads from disk per generation**. The transformer alone loads **twice**.

Let me write a proper persistent inference server that loads everything once.


Write

"""
LTX-2 Persistent Inference Server
Loads all models ONCE into GPU memory, reuses across generations.
Avoids the default pipeline behavior of rebuilding models on every call.
"""

import gc
import os
import time
import logging

os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")

import torch
import ltx_core.loader  # break circular import

from ltx_core.batch_split import BatchSplitAdapter
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
from ltx_core.loader.registry import DummyRegistry
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
from ltx_core.model.audio_vae import (
    AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
    VOCODER_COMFY_KEYS_FILTER,
    AudioDecoderConfigurator,
    VocoderConfigurator,
    decode_audio as vae_decode_audio,
)
from ltx_core.model.transformer import (
    LTXV_MODEL_COMFY_RENAMING_MAP,
    LTXModelConfigurator,
    X0Model,
)
from ltx_core.model.upsampler import LatentUpsamplerConfigurator, upsample_video
from ltx_core.model.video_vae import (
    VAE_DECODER_COMFY_KEYS_FILTER,
    VAE_ENCODER_COMFY_KEYS_FILTER,
    TilingConfig,
    VideoDecoderConfigurator,
    VideoEncoderConfigurator,
    get_video_chunks_number,
)
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import (
    EMBEDDINGS_PROCESSOR_KEY_OPS,
    GEMMA_LLM_KEY_OPS,
    GEMMA_MODEL_OPS,
    EmbeddingsProcessorConfigurator,
    GemmaTextEncoderConfigurator,
    module_ops_from_gemma_root,
)
from ltx_core.tools import AudioLatentTools, VideoLatentTools
from ltx_core.types import AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
from ltx_core.utils import find_matching_file
from ltx_core.loader import SDOps

from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
from ltx_pipelines.utils.helpers import (
    assert_resolution,
    combined_image_conditionings,
    create_noised_state,
)
from ltx_pipelines.utils.samplers import euler_denoising_loop
from ltx_pipelines.utils.denoisers import SimpleDenoiser
from ltx_pipelines.utils.types import ModalitySpec
from ltx_pipelines.utils.media_io import encode_video

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
logger = logging.getLogger(__name__)

class LTX2PersistentEngine:
    """Load all models once, keep in GPU memory, run inference fast."""

    def __init__(
        self,
        distilled_checkpoint_path: str,
        spatial_upsampler_path: str,
        gemma_root: str,
        quantization: QuantizationPolicy | None = None,
        device: torch.device | None = None,
    ):
        self.device = device or torch.device("cuda")
        self.dtype = torch.bfloat16
        registry = DummyRegistry()

        t0 = time.time()

        # ── 1. Text encoder (Gemma 3) ──
        logger.info("Loading Gemma text encoder...")
        t1 = time.time()
        module_ops = module_ops_from_gemma_root(gemma_root)
        model_folder = find_matching_file(gemma_root, "model*.safetensors").parent
        weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
        text_enc_builder = Builder(
            model_path=tuple(weight_paths),
            model_class_configurator=GemmaTextEncoderConfigurator,
            model_sd_ops=GEMMA_LLM_KEY_OPS,
            module_ops=(GEMMA_MODEL_OPS, *module_ops),
            registry=registry,
        )
        self.text_encoder = text_enc_builder.build(device=self.device, dtype=self.dtype).eval()
        logger.info(f"  Text encoder loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        # ── 2. Embeddings processor ──
        logger.info("Loading embeddings processor...")
        t1 = time.time()
        emb_builder = Builder(
            model_path=distilled_checkpoint_path,
            model_class_configurator=EmbeddingsProcessorConfigurator,
            model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS,
            registry=registry,
        )
        self.embeddings_processor = emb_builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()
        logger.info(f"  Embeddings processor loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        # ── 3. Video encoder (for image conditioning + upsampling) ──
        logger.info("Loading video encoder...")
        t1 = time.time()
        enc_builder = Builder(
            model_path=distilled_checkpoint_path,
            model_class_configurator=VideoEncoderConfigurator,
            model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
            registry=registry,
        )
        self.video_encoder = enc_builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()
        logger.info(f"  Video encoder loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        # ── 4. Spatial upsampler ──
        logger.info("Loading spatial upsampler...")
        t1 = time.time()
        up_builder = Builder(
            model_path=spatial_upsampler_path,
            model_class_configurator=LatentUpsamplerConfigurator,
            registry=registry,
        )
        self.upsampler = up_builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()
        logger.info(f"  Upsampler loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        # ── 5. Transformer (the big one) ──
        logger.info("Loading transformer...")
        t1 = time.time()
        trans_builder = Builder(
            model_path=distilled_checkpoint_path,
            model_class_configurator=LTXModelConfigurator,
            model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
            loras=(),
            registry=registry,
        )
        # Apply quantization sd_ops if needed
        sd_ops = trans_builder.model_sd_ops
        module_ops_trans = trans_builder.module_ops
        if quantization is not None:
            module_ops_trans = (*module_ops_trans, *quantization.module_ops)
            sd_ops = SDOps(
                name=f"chain_{sd_ops.name}+{quantization.sd_ops.name}",
                mapping=(*sd_ops.mapping, *quantization.sd_ops.mapping),
            )
        trans_builder = trans_builder.with_module_ops(module_ops_trans).with_sd_ops(sd_ops)
        self.transformer = X0Model(trans_builder.build(device=self.device)).to(self.device).eval()
        logger.info(f"  Transformer loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        # ── 6. Video decoder ──
        logger.info("Loading video decoder...")
        t1 = time.time()
        dec_builder = Builder(
            model_path=distilled_checkpoint_path,
            model_class_configurator=VideoDecoderConfigurator,
            model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
            registry=registry,
        )
        self.video_decoder = dec_builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()
        logger.info(f"  Video decoder loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        # ── 7. Audio decoder + vocoder ──
        logger.info("Loading audio decoder + vocoder...")
        t1 = time.time()
        audio_dec_builder = Builder(
            model_path=distilled_checkpoint_path,
            model_class_configurator=AudioDecoderConfigurator,
            model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
            registry=registry,
        )
        self.audio_decoder = audio_dec_builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()

        vocoder_builder = Builder(
            model_path=distilled_checkpoint_path,
            model_class_configurator=VocoderConfigurator,
            model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
            registry=registry,
        )
        self.vocoder = vocoder_builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()
        logger.info(f"  Audio decoder+vocoder loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        total_load = time.time() - t0
        logger.info(f"ALL MODELS LOADED in {total_load:.1f}s")
        logger.info(f"Total VRAM: {torch.cuda.memory_allocated()/1024**3:.2f} GB / {torch.cuda.get_device_properties(0).total_memory/1024**3:.1f} GB")

    @torch.inference_mode()
    def generate(
        self,
        prompt: str,
        seed: int = 42,
        height: int = 512,
        width: int = 768,
        num_frames: int = 41,
        frame_rate: float = 24.0,
        output_path: str | None = None,
    ) -> dict:
        """Generate a video. Returns timing dict."""
        assert_resolution(height=height, width=width, is_two_stage=True)
        timings = {}
        t_total = time.time()

        generator = torch.Generator(device=self.device).manual_seed(seed)
        noiser = GaussianNoiser(generator=generator)

        # ── Encode prompt ──
        t0 = time.time()
        raw_hs, mask = self.text_encoder.encode(prompt)
        ctx = self.embeddings_processor.process_hidden_states(raw_hs, mask)
        video_context, audio_context = ctx.video_encoding, ctx.audio_encoding
        timings["prompt_encode"] = time.time() - t0

        # ── Stage 1: Low-res generation ──
        t0 = time.time()
        stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
        s1_w, s1_h = width // 2, height // 2

        # Image conditioning (none for text-to-video)
        stage_1_conditionings = combined_image_conditionings(
            images=[], height=s1_h, width=s1_w,
            video_encoder=self.video_encoder, dtype=self.dtype, device=self.device,
        )

        pixel_shape = VideoPixelShape(batch=1, frames=num_frames, height=s1_h, width=s1_w, fps=frame_rate)
        v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
        video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, frame_rate)
        a_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape)
        audio_tools = AudioLatentTools(AudioPatchifier(patch_size=1), a_shape)

        video_spec = ModalitySpec(context=video_context, conditionings=stage_1_conditionings)
        audio_spec = ModalitySpec(context=audio_context)

        video_state = create_noised_state(
            tools=video_tools, conditionings=video_spec.conditionings,
            noiser=noiser, dtype=self.dtype, device=self.device,
        )
        audio_state = create_noised_state(
            tools=audio_tools, conditionings=audio_spec.conditionings,
            noiser=noiser, dtype=self.dtype, device=self.device,
        )

        denoiser = SimpleDenoiser(video_context, audio_context)
        stepper = EulerDiffusionStep()
        transformer = BatchSplitAdapter(self.transformer, max_batch_size=1)

        video_state, audio_state = euler_denoising_loop(
            sigmas=stage_1_sigmas,
            video_state=video_state,
            audio_state=audio_state,
            stepper=stepper,
            transformer=transformer,
            denoiser=denoiser,
        )

        video_state = video_tools.clear_conditioning(video_state)
        video_state = video_tools.unpatchify(video_state)
        audio_state = audio_tools.clear_conditioning(audio_state)
        audio_state = audio_tools.unpatchify(audio_state)
        timings["stage1_denoise"] = time.time() - t0

        # ── Upsample ──
        t0 = time.time()
        upscaled_video_latent = upsample_video(
            latent=video_state.latent[:1],
            video_encoder=self.video_encoder,
            upsampler=self.upsampler,
        )
        timings["upsample"] = time.time() - t0

        # ── Stage 2: Refine at full resolution ──
        t0 = time.time()
        stage_2_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
        stage_2_conditionings = combined_image_conditionings(
            images=[], height=height, width=width,
            video_encoder=self.video_encoder, dtype=self.dtype, device=self.device,
        )

        pixel_shape_2 = VideoPixelShape(batch=1, frames=num_frames, height=height, width=width, fps=frame_rate)
        v_shape_2 = VideoLatentShape.from_pixel_shape(pixel_shape_2)
        video_tools_2 = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape_2, frame_rate)
        a_shape_2 = AudioLatentShape.from_video_pixel_shape(pixel_shape_2)
        audio_tools_2 = AudioLatentTools(AudioPatchifier(patch_size=1), a_shape_2)

        video_spec_2 = ModalitySpec(
            context=video_context, conditionings=stage_2_conditionings,
            noise_scale=stage_2_sigmas[0].item(), initial_latent=upscaled_video_latent,
        )
        audio_spec_2 = ModalitySpec(
            context=audio_context,
            noise_scale=stage_2_sigmas[0].item(), initial_latent=audio_state.latent,
        )

        video_state_2 = create_noised_state(
            tools=video_tools_2, conditionings=video_spec_2.conditionings,
            noiser=noiser, dtype=self.dtype, device=self.device,
            noise_scale=video_spec_2.noise_scale, initial_latent=video_spec_2.initial_latent,
        )
        audio_state_2 = create_noised_state(
            tools=audio_tools_2, conditionings=audio_spec_2.conditionings,
            noiser=noiser, dtype=self.dtype, device=self.device,
            noise_scale=audio_spec_2.noise_scale, initial_latent=audio_spec_2.initial_latent,
        )

        denoiser_2 = SimpleDenoiser(video_context, audio_context)
        video_state_2, audio_state_2 = euler_denoising_loop(
            sigmas=stage_2_sigmas,
            video_state=video_state_2,
            audio_state=audio_state_2,
            stepper=stepper,
            transformer=transformer,
            denoiser=denoiser_2,
        )

        video_state_2 = video_tools_2.clear_conditioning(video_state_2)
        video_state_2 = video_tools_2.unpatchify(video_state_2)
        audio_state_2 = audio_tools_2.clear_conditioning(audio_state_2)
        audio_state_2 = audio_tools_2.unpatchify(audio_state_2)
        timings["stage2_denoise"] = time.time() - t0

        # ── Decode ──
        t0 = time.time()
        tiling_config = TilingConfig.default()
        video_chunks = self.video_decoder.decode_video(video_state_2.latent, tiling_config, generator)
        audio = vae_decode_audio(audio_state_2.latent, self.audio_decoder, self.vocoder)
        timings["decode"] = time.time() - t0

        # ── Encode to file ──
        if output_path:
            t0 = time.time()
            video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
            encode_video(video=video_chunks, fps=frame_rate, audio=audio,
                         output_path=output_path, video_chunks_number=video_chunks_number)
            timings["file_encode"] = time.time() - t0

        timings["total"] = time.time() - t_total
        return timings

def main():
    MODEL_DIR = "/home/ubuntu/ltx2-models"
    DISTILLED_CKPT = f"{MODEL_DIR}/ltx-2.3-22b-distilled-fp8.safetensors"
    SPATIAL_UP = f"{MODEL_DIR}/ltx-2.3-spatial-upscaler-x2-1.0.safetensors"
    GEMMA_ROOT = f"{MODEL_DIR}/gemma3"

    engine = LTX2PersistentEngine(
        distilled_checkpoint_path=DISTILLED_CKPT,
        spatial_upsampler_path=SPATIAL_UP,
        gemma_root=GEMMA_ROOT,
        quantization=QuantizationPolicy.fp8_cast(),
    )

    prompts = [
        "A golden retriever running through a field of wildflowers at sunset, cinematic lighting",
        "A futuristic cityscape at night with neon lights reflecting on wet streets, aerial drone shot",
        "Ocean waves crashing against rocky cliffs, slow motion, dramatic sky with storm clouds",
    ]

    os.makedirs("/home/ubuntu/ltx2_bench_output", exist_ok=True)

    # Warmup
    logger.info("=" * 60)
    logger.info("WARMUP RUN")
    logger.info("=" * 60)
    t = engine.generate(
        prompt=prompts[0], seed=42, num_frames=41,
        output_path="/home/ubuntu/ltx2_bench_output/persistent_warmup.mp4",
    )
    logger.info(f"Warmup: {t}")

    # Timed runs
    for i, prompt in enumerate(prompts):
        logger.info("=" * 60)
        logger.info(f"RUN {i+1}")
        logger.info("=" * 60)
        t = engine.generate(
            prompt=prompt, seed=42 + i, num_frames=41,
            output_path=f"/home/ubuntu/ltx2_bench_output/persistent_{i}.mp4",
        )
        logger.info(f"Run {i+1} timings: {t}")
        for k, v in t.items():
            logger.info(f"  {k}: {v:.2f}s")

    logger.info(f"Peak VRAM: {torch.cuda.max_memory_allocated()/1024**3:.2f} GB")

if __name__ == "__main__":
    main()


File created successfully at: /home/ubuntu/ltx2_server.py

"""
LTX-2 Persistent Inference Server
Loads all models ONCE into GPU memory, reuses across generations.
Avoids the default pipeline behavior of rebuilding models on every call.
"""

import gc
import os
import time
import logging

os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")

import torch
import ltx_core.loader  # break circular import

from ltx_core.batch_split import BatchSplitAdapter
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
from ltx_core.loader.registry import DummyRegistry
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
from ltx_core.model.audio_vae import (
    AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
    VOCODER_COMFY_KEYS_FILTER,
    AudioDecoderConfigurator,
    VocoderConfigurator,
    decode_audio as vae_decode_audio,
)
from ltx_core.model.transformer import (
    LTXV_MODEL_COMFY_RENAMING_MAP,
    LTXModelConfigurator,
    X0Model,
)
from ltx_core.model.upsampler import LatentUpsamplerConfigurator, upsample_video
from ltx_core.model.video_vae import (
    VAE_DECODER_COMFY_KEYS_FILTER,
    VAE_ENCODER_COMFY_KEYS_FILTER,
    TilingConfig,
    VideoDecoderConfigurator,
    VideoEncoderConfigurator,
    get_video_chunks_number,
)
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import (
    EMBEDDINGS_PROCESSOR_KEY_OPS,
    GEMMA_LLM_KEY_OPS,
    GEMMA_MODEL_OPS,
    EmbeddingsProcessorConfigurator,
    GemmaTextEncoderConfigurator,
    module_ops_from_gemma_root,
)
from ltx_core.tools import AudioLatentTools, VideoLatentTools
from ltx_core.types import AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
from ltx_core.utils import find_matching_file
from ltx_core.loader import SDOps

from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
from ltx_pipelines.utils.helpers import (
    assert_resolution,
    combined_image_conditionings,
    create_noised_state,
)
from ltx_pipelines.utils.samplers import euler_denoising_loop
from ltx_pipelines.utils.denoisers import SimpleDenoiser
from ltx_pipelines.utils.types import ModalitySpec
from ltx_pipelines.utils.media_io import encode_video

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
logger = logging.getLogger(__name__)

class LTX2PersistentEngine:
    """Load all models once, keep in GPU memory, run inference fast."""

    def __init__(
        self,
        distilled_checkpoint_path: str,
        spatial_upsampler_path: str,
        gemma_root: str,
        quantization: QuantizationPolicy | None = None,
        device: torch.device | None = None,
    ):
        self.device = device or torch.device("cuda")
        self.dtype = torch.bfloat16
        registry = DummyRegistry()

        t0 = time.time()

        # ── 1. Text encoder (Gemma 3) ──
        logger.info("Loading Gemma text encoder...")
        t1 = time.time()
        module_ops = module_ops_from_gemma_root(gemma_root)
        model_folder = find_matching_file(gemma_root, "model*.safetensors").parent
        weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
        text_enc_builder = Builder(
            model_path=tuple(weight_paths),
            model_class_configurator=GemmaTextEncoderConfigurator,
            model_sd_ops=GEMMA_LLM_KEY_OPS,
            module_ops=(GEMMA_MODEL_OPS, *module_ops),
            registry=registry,
        )
        self.text_encoder = text_enc_builder.build(device=self.device, dtype=self.dtype).eval()
        logger.info(f"  Text encoder loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        # ── 2. Embeddings processor ──
        logger.info("Loading embeddings processor...")
        t1 = time.time()
        emb_builder = Builder(
            model_path=distilled_checkpoint_path,
            model_class_configurator=EmbeddingsProcessorConfigurator,
            model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS,
            registry=registry,
        )
        self.embeddings_processor = emb_builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()
        logger.info(f"  Embeddings processor loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        # ── 3. Video encoder (for image conditioning + upsampling) ──
        logger.info("Loading video encoder...")
        t1 = time.time()
        enc_builder = Builder(
            model_path=distilled_checkpoint_path,
            model_class_configurator=VideoEncoderConfigurator,
            model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
            registry=registry,
        )
        self.video_encoder = enc_builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()
        logger.info(f"  Video encoder loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        # ── 4. Spatial upsampler ──
        logger.info("Loading spatial upsampler...")
        t1 = time.time()
        up_builder = Builder(
            model_path=spatial_upsampler_path,
            model_class_configurator=LatentUpsamplerConfigurator,
            registry=registry,
        )
        self.upsampler = up_builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()
        logger.info(f"  Upsampler loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        # ── 5. Transformer (the big one) ──
        logger.info("Loading transformer...")
        t1 = time.time()
        trans_builder = Builder(
            model_path=distilled_checkpoint_path,
            model_class_configurator=LTXModelConfigurator,
            model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
            loras=(),
            registry=registry,
        )
        # Apply quantization sd_ops if needed
        sd_ops = trans_builder.model_sd_ops
        module_ops_trans = trans_builder.module_ops
        if quantization is not None:
            module_ops_trans = (*module_ops_trans, *quantization.module_ops)
            sd_ops = SDOps(
                name=f"chain_{sd_ops.name}+{quantization.sd_ops.name}",
                mapping=(*sd_ops.mapping, *quantization.sd_ops.mapping),
            )
        trans_builder = trans_builder.with_module_ops(module_ops_trans).with_sd_ops(sd_ops)
        self.transformer = X0Model(trans_builder.build(device=self.device)).to(self.device).eval()
        logger.info(f"  Transformer loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        # ── 6. Video decoder ──
        logger.info("Loading video decoder...")
        t1 = time.time()
        dec_builder = Builder(
            model_path=distilled_checkpoint_path,
            model_class_configurator=VideoDecoderConfigurator,
            model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
            registry=registry,
        )
        self.video_decoder = dec_builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()
        logger.info(f"  Video decoder loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        # ── 7. Audio decoder + vocoder ──
        logger.info("Loading audio decoder + vocoder...")
        t1 = time.time()
        audio_dec_builder = Builder(
            model_path=distilled_checkpoint_path,
            model_class_configurator=AudioDecoderConfigurator,
            model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
            registry=registry,
        )
        self.audio_decoder = audio_dec_builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()

        vocoder_builder = Builder(
            model_path=distilled_checkpoint_path,
            model_class_configurator=VocoderConfigurator,
            model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
            registry=registry,
        )
        self.vocoder = vocoder_builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()
        logger.info(f"  Audio decoder+vocoder loaded in {time.time()-t1:.1f}s | VRAM: {torch.cuda.memory_allocated()/1024**3:.1f}GB")

        total_load = time.time() - t0
        logger.info(f"ALL MODELS LOADED in {total_load:.1f}s")
        logger.info(f"Total VRAM: {torch.cuda.memory_allocated()/1024**3:.2f} GB / {torch.cuda.get_device_properties(0).total_memory/1024**3:.1f} GB")

    @torch.inference_mode()
    def generate(
        self,
        prompt: str,
        seed: int = 42,
        height: int = 512,
        width: int = 768,
        num_frames: int = 41,
        frame_rate: float = 24.0,
        output_path: str | None = None,
    ) -> dict:
        """Generate a video. Returns timing dict."""
        assert_resolution(height=height, width=width, is_two_stage=True)
        timings = {}
        t_total = time.time()

        generator = torch.Generator(device=self.device).manual_seed(seed)
        noiser = GaussianNoiser(generator=generator)

        # ── Encode prompt ──
        t0 = time.time()
        raw_hs, mask = self.text_encoder.encode(prompt)
        ctx = self.embeddings_processor.process_hidden_states(raw_hs, mask)
        video_context, audio_context = ctx.video_encoding, ctx.audio_encoding
        timings["prompt_encode"] = time.time() - t0

        # ── Stage 1: Low-res generation ──
        t0 = time.time()
        stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
        s1_w, s1_h = width // 2, height // 2

        # Image conditioning (none for text-to-video)
        stage_1_conditionings = combined_image_conditionings(
            images=[], height=s1_h, width=s1_w,
            video_encoder=self.video_encoder, dtype=self.dtype, device=self.device,
        )

        pixel_shape = VideoPixelShape(batch=1, frames=num_frames, height=s1_h, width=s1_w, fps=frame_rate)
        v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
        video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, frame_rate)
        a_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape)
        audio_tools = AudioLatentTools(AudioPatchifier(patch_size=1), a_shape)

        video_spec = ModalitySpec(context=video_context, conditionings=stage_1_conditionings)
        audio_spec = ModalitySpec(context=audio_context)

        video_state = create_noised_state(
            tools=video_tools, conditionings=video_spec.conditionings,
            noiser=noiser, dtype=self.dtype, device=self.device,
        )
        audio_state = create_noised_state(
            tools=audio_tools, conditionings=audio_spec.conditionings,
            noiser=noiser, dtype=self.dtype, device=self.device,
        )

        denoiser = SimpleDenoiser(video_context, audio_context)
        stepper = EulerDiffusionStep()
        transformer = BatchSplitAdapter(self.transformer, max_batch_size=1)

        video_state, audio_state = euler_denoising_loop(
            sigmas=stage_1_sigmas,
            video_state=video_state,
            audio_state=audio_state,
            stepper=stepper,
            transformer=transformer,
            denoiser=denoiser,
        )

        video_state = video_tools.clear_conditioning(video_state)
        video_state = video_tools.unpatchify(video_state)
        audio_state = audio_tools.clear_conditioning(audio_state)
        audio_state = audio_tools.unpatchify(audio_state)
        timings["stage1_denoise"] = time.time() - t0

        # ── Upsample ──
        t0 = time.time()
        upscaled_video_latent = upsample_video(
            latent=video_state.latent[:1],
            video_encoder=self.video_encoder,
            upsampler=self.upsampler,
        )
        timings["upsample"] = time.time() - t0

        # ── Stage 2: Refine at full resolution ──
        t0 = time.time()
        stage_2_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
        stage_2_conditionings = combined_image_conditionings(
            images=[], height=height, width=width,
            video_encoder=self.video_encoder, dtype=self.dtype, device=self.device,
        )

        pixel_shape_2 = VideoPixelShape(batch=1, frames=num_frames, height=height, width=width, fps=frame_rate)
        v_shape_2 = VideoLatentShape.from_pixel_shape(pixel_shape_2)
        video_tools_2 = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape_2, frame_rate)
        a_shape_2 = AudioLatentShape.from_video_pixel_shape(pixel_shape_2)
        audio_tools_2 = AudioLatentTools(AudioPatchifier(patch_size=1), a_shape_2)

        video_spec_2 = ModalitySpec(
            context=video_context, conditionings=stage_2_conditionings,
            noise_scale=stage_2_sigmas[0].item(), initial_latent=upscaled_video_latent,
        )
        audio_spec_2 = ModalitySpec(
            context=audio_context,
            noise_scale=stage_2_sigmas[0].item(), initial_latent=audio_state.latent,
        )

        video_state_2 = create_noised_state(
            tools=video_tools_2, conditionings=video_spec_2.conditionings,
            noiser=noiser, dtype=self.dtype, device=self.device,
            noise_scale=video_spec_2.noise_scale, initial_latent=video_spec_2.initial_latent,
        )
        audio_state_2 = create_noised_state(
            tools=audio_tools_2, conditionings=audio_spec_2.conditionings,
            noiser=noiser, dtype=self.dtype, device=self.device,
            noise_scale=audio_spec_2.noise_scale, initial_latent=audio_spec_2.initial_latent,
        )

        denoiser_2 = SimpleDenoiser(video_context, audio_context)
        video_state_2, audio_state_2 = euler_denoising_loop(
            sigmas=stage_2_sigmas,
            video_state=video_state_2,
            audio_state=audio_state_2,
            stepper=stepper,
            transformer=transformer,
            denoiser=denoiser_2,
        )

        video_state_2 = video_tools_2.clear_conditioning(video_state_2)
        video_state_2 = video_tools_2.unpatchify(video_state_2)
        audio_state_2 = audio_tools_2.clear_conditioning(audio_state_2)
        audio_state_2 = audio_tools_2.unpatchify(audio_state_2)
        timings["stage2_denoise"] = time.time() - t0

        # ── Decode ──
        t0 = time.time()
        tiling_config = TilingConfig.default()
        video_chunks = self.video_decoder.decode_video(video_state_2.latent, tiling_config, generator)
        audio = vae_decode_audio(audio_state_2.latent, self.audio_decoder, self.vocoder)
        timings["decode"] = time.time() - t0

        # ── Encode to file ──
        if output_path:
            t0 = time.time()
            video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
            encode_video(video=video_chunks, fps=frame_rate, audio=audio,
                         output_path=output_path, video_chunks_number=video_chunks_number)
            timings["file_encode"] = time.time() - t0

        timings["total"] = time.time() - t_total
        return timings

def main():
    MODEL_DIR = "/home/ubuntu/ltx2-models"
    DISTILLED_CKPT = f"{MODEL_DIR}/ltx-2.3-22b-distilled-fp8.safetensors"
    SPATIAL_UP = f"{MODEL_DIR}/ltx-2.3-spatial-upscaler-x2-1.0.safetensors"
    GEMMA_ROOT = f"{MODEL_DIR}/gemma3"

    engine = LTX2PersistentEngine(
        distilled_checkpoint_path=DISTILLED_CKPT,
        spatial_upsampler_path=SPATIAL_UP,
        gemma_root=GEMMA_ROOT,
        quantization=QuantizationPolicy.fp8_cast(),
    )

    prompts = [
        "A golden retriever running through a field of wildflowers at sunset, cinematic lighting",
        "A futuristic cityscape at night with neon lights reflecting on wet streets, aerial drone shot",
        "Ocean waves crashing against rocky cliffs, slow motion, dramatic sky with storm clouds",
    ]

    os.makedirs("/home/ubuntu/ltx2_bench_output", exist_ok=True)

    # Warmup
    logger.info("=" * 60)
    logger.info("WARMUP RUN")
    logger.info("=" * 60)
    t = engine.generate(
        prompt=prompts[0], seed=42, num_frames=41,
        output_path="/home/ubuntu/ltx2_bench_output/persistent_warmup.mp4",
    )
    logger.info(f"Warmup: {t}")

    # Timed runs
    for i, prompt in enumerate(prompts):
        logger.info("=" * 60)
        logger.info(f"RUN {i+1}")
        logger.info("=" * 60)
        t = engine.generate(
            prompt=prompt, seed=42 + i, num_frames=41,
            output_path=f"/home/ubuntu/ltx2_bench_output/persistent_{i}.mp4",
        )
        logger.info(f"Run {i+1} timings: {t}")
        for k, v in t.items():
            logger.info(f"  {k}: {v:.2f}s")

    logger.info(f"Peak VRAM: {torch.cuda.max_memory_allocated()/1024**3:.2f} GB")

if __name__ == "__main__":
    main()


Bash

/home/ubuntu/LTX-2/.venv/bin/python -u /home/ubuntu/ltx2_server.py 2>&1

Run persistent inference engine


2026-04-04 14:42:00,324 Loading Gemma text encoder...
Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. `use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`.
2026-04-04 14:42:07,095   Text encoder loaded in 6.8s | VRAM: 22.8GB
2026-04-04 14:42:07,095 Loading embeddings processor...
2026-04-04 14:42:08,284   Embeddings processor loaded in 1.2s | VRAM: 28.7GB
2026-04-04 14:42:08,284 Loading video encoder...
2026-04-04 14:42:08,470   Video encoder loaded in 0.2s | VRAM: 29.3GB
2026-04-04 14:42:08,470 Loading spatial upsampler...
2026-04-04 14:42:08,678   Upsampler loaded in 0.2s | VRAM: 30.2GB
2026-04-04 14:42:08,678 Loading transformer...
2026-04-04 14:42:14,634   Transformer loaded in 6.0s | VRAM: 48.3GB
2026-04-04 14:42:14,634 Loading video decoder...
2026-04-04 14:42:14,835   Video decoder loaded in 0.2s | VRAM: 49.1GB
2026-04-04 14:42:14,835 Loading audio decoder + vocoder...
2026-04-04 14:42:15,735   Audio decoder+vocoder loaded in 0.9s | VRAM: 49.4GB
2026-04-04 14:42:15,736 ALL MODELS LOADED in 15.4s
2026-04-04 14:42:15,736 Total VRAM: 49.41 GB / 79.3 GB
2026-04-04 14:42:15,737 ============================================================
2026-04-04 14:42:15,737 WARMUP RUN
2026-04-04 14:42:15,737 ============================================================

  0%|          | 0/8 [00:00<?, ?it/s]
 12%|█▎        | 1/8 [00:00<00:02,  3.05it/s]
 25%|██▌       | 2/8 [00:00<00:01,  3.21it/s]
 38%|███▊      | 3/8 [00:00<00:01,  3.27it/s]
 50%|█████     | 4/8 [00:01<00:01,  3.30it/s]
 62%|██████▎   | 5/8 [00:01<00:00,  3.31it/s]
 75%|███████▌  | 6/8 [00:01<00:00,  3.32it/s]
 88%|████████▊ | 7/8 [00:02<00:00,  3.33it/s]
100%|██████████| 8/8 [00:02<00:00,  3.34it/s]
100%|██████████| 8/8 [00:02<00:00,  3.30it/s]

  0%|          | 0/3 [00:00<?, ?it/s]
 33%|███▎      | 1/3 [00:00<00:01,  1.78it/s]
 67%|██████▋   | 2/3 [00:01<00:00,  1.78it/s]
100%|██████████| 3/3 [00:01<00:00,  1.78it/s]
100%|██████████| 3/3 [00:01<00:00,  1.78it/s]

  0%|          | 0/1 [00:00<?, ?it/s]
100%|██████████| 1/1 [00:00<00:00,  2.33it/s]
100%|██████████| 1/1 [00:00<00:00,  2.33it/s]
2026-04-04 14:42:23,091 Video saved to /home/ubuntu/ltx2_bench_output/persistent_warmup.mp4
2026-04-04 14:42:23,102 Warmup: {'prompt_encode': 0.6089284420013428, 'stage1_denoise': 2.429271936416626, 'upsample': 0.26369595527648926, 'stage2_denoise': 1.6863317489624023, 'decode': 0.35330843925476074, 'file_encode': 2.023315191268921, 'total': 7.364931106567383}
2026-04-04 14:42:23,103 ============================================================
2026-04-04 14:42:23,103 RUN 1
2026-04-04 14:42:23,103 ============================================================

  0%|          | 0/8 [00:00<?, ?it/s]
 12%|█▎        | 1/8 [00:00<00:02,  3.24it/s]
 25%|██▌       | 2/8 [00:00<00:01,  3.28it/s]
 38%|███▊      | 3/8 [00:00<00:01,  3.24it/s]
 50%|█████     | 4/8 [00:01<00:01,  3.27it/s]
 62%|██████▎   | 5/8 [00:01<00:00,  3.28it/s]
 75%|███████▌  | 6/8 [00:01<00:00,  3.30it/s]
 88%|████████▊ | 7/8 [00:02<00:00,  3.30it/s]
100%|██████████| 8/8 [00:02<00:00,  3.31it/s]
100%|██████████| 8/8 [00:02<00:00,  3.29it/s]

  0%|          | 0/3 [00:00<?, ?it/s]
 33%|███▎      | 1/3 [00:00<00:01,  1.78it/s]
 67%|██████▋   | 2/3 [00:01<00:00,  1.78it/s]
100%|██████████| 3/3 [00:01<00:00,  1.78it/s]
100%|██████████| 3/3 [00:01<00:00,  1.78it/s]

  0%|          | 0/1 [00:00<?, ?it/s]
100%|██████████| 1/1 [00:00<00:00,  2.40it/s]
100%|██████████| 1/1 [00:00<00:00,  2.39it/s]
2026-04-04 14:42:29,720 Video saved to /home/ubuntu/ltx2_bench_output/persistent_0.mp4
2026-04-04 14:42:29,724 Run 1 timings: {'prompt_encode': 0.23945307731628418, 'stage1_denoise': 2.436739206314087, 'upsample': 0.15655183792114258, 'stage2_denoise': 1.6862413883209229, 'decode': 0.10430765151977539, 'file_encode': 1.9978010654449463, 'total': 6.6211841106414795}
2026-04-04 14:42:29,725   prompt_encode: 0.24s
2026-04-04 14:42:29,725   stage1_denoise: 2.44s
2026-04-04 14:42:29,725   upsample: 0.16s
2026-04-04 14:42:29,725   stage2_denoise: 1.69s
2026-04-04 14:42:29,725   decode: 0.10s
2026-04-04 14:42:29,725   file_encode: 2.00s
2026-04-04 14:42:29,725   total: 6.62s
2026-04-04 14:42:29,725 ============================================================
2026-04-04 14:42:29,725 RUN 2
2026-04-04 14:42:29,725 ============================================================

  0%|          | 0/8 [00:00<?, ?it/s]
 12%|█▎        | 1/8 [00:00<00:02,  3.28it/s]
 25%|██▌       | 2/8 [00:00<00:01,  3.30it/s]
 38%|███▊      | 3/8 [00:00<00:01,  3.30it/s]
 50%|█████     | 4/8 [00:01<00:01,  3.31it/s]
 62%|██████▎   | 5/8 [00:01<00:00,  3.30it/s]
 75%|███████▌  | 6/8 [00:01<00:00,  3.31it/s]
 88%|████████▊ | 7/8 [00:02<00:00,  3.27it/s]
100%|██████████| 8/8 [00:02<00:00,  3.28it/s]
100%|██████████| 8/8 [00:02<00:00,  3.29it/s]

  0%|          | 0/3 [00:00<?, ?it/s]
 33%|███▎      | 1/3 [00:00<00:01,  1.78it/s]
 67%|██████▋   | 2/3 [00:01<00:00,  1.78it/s]
100%|██████████| 3/3 [00:01<00:00,  1.78it/s]
100%|██████████| 3/3 [00:01<00:00,  1.78it/s]

  0%|          | 0/1 [00:00<?, ?it/s]
100%|██████████| 1/1 [00:00<00:00,  2.54it/s]
100%|██████████| 1/1 [00:00<00:00,  2.54it/s]
2026-04-04 14:42:36,327 Video saved to /home/ubuntu/ltx2_bench_output/persistent_1.mp4
2026-04-04 14:42:36,330 Run 2 timings: {'prompt_encode': 0.23836755752563477, 'stage1_denoise': 2.4360907077789307, 'upsample': 0.15742015838623047, 'stage2_denoise': 1.6884264945983887, 'decode': 0.10558104515075684, 'file_encode': 1.9780371189117432, 'total': 6.60401463508606}
2026-04-04 14:42:36,330   prompt_encode: 0.24s
2026-04-04 14:42:36,330   stage1_denoise: 2.44s
2026-04-04 14:42:36,330   upsample: 0.16s
2026-04-04 14:42:36,330   stage2_denoise: 1.69s
2026-04-04 14:42:36,330   decode: 0.11s
2026-04-04 14:42:36,330   file_encode: 1.98s
2026-04-04 14:42:36,330   total: 6.60s
2026-04-04 14:42:36,330 ============================================================
2026-04-04 14:42:36,330 RUN 3
2026-04-04 14:42:36,330 ============================================================

  0%|          | 0/8 [00:00<?, ?it/s]
 12%|█▎        | 1/8 [00:00<00:02,  3.23it/s]
 25%|██▌       | 2/8 [00:00<00:01,  3.25it/s]
 38%|███▊      | 3/8 [00:00<00:01,  3.24it/s]
 50%|█████     | 4/8 [00:01<00:01,  3.23it/s]
 62%|██████▎   | 5/8 [00:01<00:00,  3.23it/s]
 75%|███████▌  | 6/8 [00:01<00:00,  3.24it/s]
 88%|████████▊ | 7/8 [00:02<00:00,  3.25it/s]
100%|██████████| 8/8 [00:02<00:00,  3.28it/s]
100%|██████████| 8/8 [00:02<00:00,  3.25it/s]

  0%|          | 0/3 [00:00<?, ?it/s]
 33%|███▎      | 1/3 [00:00<00:01,  1.78it/s]
 67%|██████▋   | 2/3 [00:01<00:00,  1.78it/s]
100%|██████████| 3/3 [00:01<00:00,  1.78it/s]
100%|██████████| 3/3 [00:01<00:00,  1.78it/s]

  0%|          | 0/1 [00:00<?, ?it/s]
100%|██████████| 1/1 [00:00<00:00,  2.53it/s]
100%|██████████| 1/1 [00:00<00:00,  2.53it/s]
2026-04-04 14:42:42,967 Video saved to /home/ubuntu/ltx2_bench_output/persistent_2.mp4
2026-04-04 14:42:42,969 Run 3 timings: {'prompt_encode': 0.23949337005615234, 'stage1_denoise': 2.4636967182159424, 'upsample': 0.1576066017150879, 'stage2_denoise': 1.6888017654418945, 'decode': 0.10845232009887695, 'file_encode': 1.9801409244537354, 'total': 6.638295412063599}
2026-04-04 14:42:42,969   prompt_encode: 0.24s
2026-04-04 14:42:42,969   stage1_denoise: 2.46s
2026-04-04 14:42:42,969   upsample: 0.16s
2026-04-04 14:42:42,969   stage2_denoise: 1.69s
2026-04-04 14:42:42,969   decode: 0.11s
2026-04-04 14:42:42,969   file_encode: 1.98s
2026-04-04 14:42:42,969   total: 6.64s
2026-04-04 14:42:42,969 Peak VRAM: 53.52 GB


Massive difference. Here's the comparison:

## Before vs After: Persistent Model Loading

| Metric | Before (reload each call) | After (persistent) | Speedup |
|---|---|---|---|
| **Total per generation** | **~35s** | **~6.6s** | **5.3x** |
| Prompt encode | ~26s (model load dominated) | 0.24s | 108x |
| Stage 1 denoise (8 steps) | ~2.4s | 2.44s | same |
| Upsample | included in overhead | 0.16s | - |
| Stage 2 denoise (3 steps) | ~1.7s | 1.69s | same |
| Decode | included in overhead | 0.11s | - |
| File encode (MP4) | ~2.1s | 1.98s | same |

**VRAM usage: 49.4 GB resident / 53.5 GB peak** out of 80GB — meaning there's ~27GB headroom. Still tight for concurrent generation but the actual compute is only ~4.6s of diffusion per video, with ~2s of file encoding overhead.

The script is at `/home/ubuntu/ltx2_server.py`. Want me to turn this into an actual HTTP server with queuing, or test concurrent generation now that we know the real VRAM footprint?


per min  how manyw e can do


Quick math from the results: **6.6s per generation = ~9 videos/min**.

But ~2s of that is MP4 file encoding on CPU. If we skip file I/O or do it async, pure GPU time is **~4.6s = ~13 videos/min**.

Let me verify with a quick burst test.


Write

"""LTX-2 throughput test — how many videos per minute with persistent engine."""

import os
import time

os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")

import torch
import ltx_core.loader

from ltx_core.quantization import QuantizationPolicy
from ltx2_server import LTX2PersistentEngine

MODEL_DIR = "/home/ubuntu/ltx2-models"

engine = LTX2PersistentEngine(
    distilled_checkpoint_path=f"{MODEL_DIR}/ltx-2.3-22b-distilled-fp8.safetensors",
    spatial_upsampler_path=f"{MODEL_DIR}/ltx-2.3-spatial-upscaler-x2-1.0.safetensors",
    gemma_root=f"{MODEL_DIR}/gemma3",
    quantization=QuantizationPolicy.fp8_cast(),
)

prompts = [
    "A golden retriever running through wildflowers at sunset",
    "Futuristic cityscape at night with neon lights on wet streets",
    "Ocean waves crashing against rocky cliffs, slow motion",
    "Chef preparing sushi in a Japanese restaurant, close-up",
    "Astronaut floating in space with Earth in background",
    "Rain falling on a quiet Japanese garden with koi pond",
    "Sports car drifting on a mountain road, drone shot",
    "Northern lights over a frozen lake, timelapse",
    "Street musician playing guitar in Paris at dusk",
    "Underwater coral reef with tropical fish swimming",
]

os.makedirs("/home/ubuntu/ltx2_bench_output", exist_ok=True)

# Warmup
print("Warmup...")
engine.generate(prompt=prompts[0], seed=0, num_frames=41)

# Burst: 10 generations, with file save
print("\n" + "=" * 60)
print("BURST TEST: 10 videos WITH file save")
print("=" * 60)
t0 = time.time()
for i, p in enumerate(prompts):
    t = engine.generate(
        prompt=p, seed=i, num_frames=41,
        output_path=f"/home/ubuntu/ltx2_bench_output/burst_{i}.mp4",
    )
    print(f"  [{i+1}/10] {t['total']:.2f}s  (denoise: {t['stage1_denoise']+t['stage2_denoise']:.2f}s, encode: {t.get('file_encode',0):.2f}s)")
burst_with_save = time.time() - t0
print(f"Total: {burst_with_save:.1f}s | Rate: {10/burst_with_save*60:.1f} videos/min")

# Burst: 10 generations, NO file save (pure GPU throughput)
print("\n" + "=" * 60)
print("BURST TEST: 10 videos WITHOUT file save (pure GPU)")
print("=" * 60)
t0 = time.time()
for i, p in enumerate(prompts):
    t = engine.generate(prompt=p, seed=i, num_frames=41)
    print(f"  [{i+1}/10] {t['total']:.2f}s")
burst_no_save = time.time() - t0
print(f"Total: {burst_no_save:.1f}s | Rate: {10/burst_no_save*60:.1f} videos/min")

print("\n" + "=" * 60)
print("SUMMARY")
print("=" * 60)
print(f"With file save:    {10/burst_with_save*60:.1f} videos/min  ({burst_with_save/10:.2f}s each)")
print(f"Without file save: {10/burst_no_save*60:.1f} videos/min  ({burst_no_save/10:.2f}s each)")
print(f"Peak VRAM: {torch.cuda.max_memory_allocated()/1024**3:.2f} GB")


File created successfully at: /home/ubuntu/ltx2_throughput.py

"""LTX-2 throughput test — how many videos per minute with persistent engine."""

import os
import time

os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")

import torch
import ltx_core.loader

from ltx_core.quantization import QuantizationPolicy
from ltx2_server import LTX2PersistentEngine

MODEL_DIR = "/home/ubuntu/ltx2-models"

engine = LTX2PersistentEngine(
    distilled_checkpoint_path=f"{MODEL_DIR}/ltx-2.3-22b-distilled-fp8.safetensors",
    spatial_upsampler_path=f"{MODEL_DIR}/ltx-2.3-spatial-upscaler-x2-1.0.safetensors",
    gemma_root=f"{MODEL_DIR}/gemma3",
    quantization=QuantizationPolicy.fp8_cast(),
)

prompts = [
    "A golden retriever running through wildflowers at sunset",
    "Futuristic cityscape at night with neon lights on wet streets",
    "Ocean waves crashing against rocky cliffs, slow motion",
    "Chef preparing sushi in a Japanese restaurant, close-up",
    "Astronaut floating in space with Earth in background",
    "Rain falling on a quiet Japanese garden with koi pond",
    "Sports car drifting on a mountain road, drone shot",
    "Northern lights over a frozen lake, timelapse",
    "Street musician playing guitar in Paris at dusk",
    "Underwater coral reef with tropical fish swimming",
]

os.makedirs("/home/ubuntu/ltx2_bench_output", exist_ok=True)

# Warmup
print("Warmup...")
engine.generate(prompt=prompts[0], seed=0, num_frames=41)

# Burst: 10 generations, with file save
print("\n" + "=" * 60)
print("BURST TEST: 10 videos WITH file save")
print("=" * 60)
t0 = time.time()
for i, p in enumerate(prompts):
    t = engine.generate(
        prompt=p, seed=i, num_frames=41,
        output_path=f"/home/ubuntu/ltx2_bench_output/burst_{i}.mp4",
    )
    print(f"  [{i+1}/10] {t['total']:.2f}s  (denoise: {t['stage1_denoise']+t['stage2_denoise']:.2f}s, encode: {t.get('file_encode',0):.2f}s)")
burst_with_save = time.time() - t0
print(f"Total: {burst_with_save:.1f}s | Rate: {10/burst_with_save*60:.1f} videos/min")

# Burst: 10 generations, NO file save (pure GPU throughput)
print("\n" + "=" * 60)
print("BURST TEST: 10 videos WITHOUT file save (pure GPU)")
print("=" * 60)
t0 = time.time()
for i, p in enumerate(prompts):
    t = engine.generate(prompt=p, seed=i, num_frames=41)
    print(f"  [{i+1}/10] {t['total']:.2f}s")
burst_no_save = time.time() - t0
print(f"Total: {burst_no_save:.1f}s | Rate: {10/burst_no_save*60:.1f} videos/min")

print("\n" + "=" * 60)
print("SUMMARY")
print("=" * 60)
print(f"With file save:    {10/burst_with_save*60:.1f} videos/min  ({burst_with_save/10:.2f}s each)")
print(f"Without file save: {10/burst_no_save*60:.1f} videos/min  ({burst_no_save/10:.2f}s each)")
print(f"Peak VRAM: {torch.cuda.max_memory_allocated()/1024**3:.2f} GB")


Bash

/home/ubuntu/LTX-2/.venv/bin/python -u /home/ubuntu/ltx2_throughput.py 2>&1

Run throughput burst test