github-actions[bot]
874 lines
36 KiB
Python
874 lines
36 KiB
Python
"""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: each
|
|
block is self-contained, so no central model-coordinator object is needed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
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.block_streaming import DISK_CPU_SLOTS, StreamingModelBuilder
|
|
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.loader import SDOps
|
|
from ltx_core.loader.attention_ops import set_attention_module_op
|
|
from ltx_core.loader.fuse_loras import bf16_fuse_rule
|
|
from ltx_core.loader.module_ops import ModuleOps
|
|
from ltx_core.loader.primitives import BuilderProtocol, LoraPathStrengthAndSDOps, ModelBuilderProtocol
|
|
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.model_protocol import LTXModelProtocol, ModelConfigurator
|
|
from ltx_core.model.transformer import (
|
|
LTXV_MODEL_COMFY_RENAMING_MAP,
|
|
LTXModelConfigurator,
|
|
X0Model,
|
|
)
|
|
from ltx_core.model.transformer.attention import (
|
|
AttentionCallable,
|
|
AttentionFunction,
|
|
)
|
|
from ltx_core.model.transformer.compiling import (
|
|
CompilationConfig,
|
|
build_compile_transformer_op,
|
|
modify_sd_ops_for_compilation,
|
|
)
|
|
from ltx_core.model.upsampler import LatentUpsamplerConfigurator, upsample_video
|
|
from ltx_core.model.video_vae import (
|
|
MEMORY_EFFICIENT_DECODE,
|
|
VAE_DECODER_COMFY_KEYS_FILTER,
|
|
VAE_ENCODER_COMFY_KEYS_FILTER,
|
|
TilingConfig,
|
|
VideoDecoderConfigurator,
|
|
VideoEncoder,
|
|
VideoEncoderConfigurator,
|
|
)
|
|
from ltx_core.quantization import QuantizationPolicy, fp8_cast_fuse_rule
|
|
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 EmbeddingsProcessor, 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.allocator_trim_strategy import AllocatorTrimStrategy
|
|
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, OffloadMode
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
T = TypeVar("T")
|
|
_M = TypeVar("_M", bound=torch.nn.Module)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Internal helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _chain_quantization(
|
|
sd_ops: SDOps,
|
|
module_ops: tuple[ModuleOps, ...],
|
|
quantization: QuantizationPolicy,
|
|
) -> tuple[SDOps, tuple[ModuleOps, ...]]:
|
|
chained_sd_ops = sd_ops
|
|
if quantization.sd_ops is not None:
|
|
chained_sd_ops = SDOps(
|
|
name=f"sd_ops_chain_{sd_ops.name}+{quantization.sd_ops.name}",
|
|
mapping=(*sd_ops.mapping, *quantization.sd_ops.mapping),
|
|
)
|
|
return chained_sd_ops, (*module_ops, *quantization.module_ops)
|
|
|
|
|
|
def _apply_compile_ops(
|
|
sd_ops: SDOps,
|
|
module_ops: tuple[ModuleOps, ...],
|
|
loras: tuple[LoraPathStrengthAndSDOps, ...],
|
|
number_of_layers: int,
|
|
compilation_config: CompilationConfig,
|
|
) -> tuple[SDOps, tuple[ModuleOps, ...], tuple[LoraPathStrengthAndSDOps, ...]]:
|
|
"""Rewrite sd_ops/module_ops/LoRAs for compiled blocks (params land under ``_orig_mod``)."""
|
|
sd_ops = modify_sd_ops_for_compilation(sd_ops, number_of_layers)
|
|
compile_op = build_compile_transformer_op(compilation_config)
|
|
module_ops = (*module_ops, compile_op)
|
|
loras = tuple(
|
|
LoraPathStrengthAndSDOps(
|
|
lora.path,
|
|
lora.strength,
|
|
modify_sd_ops_for_compilation(lora.sd_ops, number_of_layers),
|
|
)
|
|
for lora in loras
|
|
)
|
|
return sd_ops, module_ops, loras
|
|
|
|
|
|
@contextmanager
|
|
def _streaming_model(
|
|
builder: StreamingModelBuilder,
|
|
target_device: torch.device,
|
|
dtype: torch.dtype,
|
|
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
|
) -> Iterator:
|
|
"""Build a streaming wrapper, yield it, then tear down and free memory.
|
|
The builder's own ``cpu_slots_count`` selects RAM vs disk streaming.
|
|
``teardown()`` always runs -- it releases non-memory resources (forward
|
|
hooks, the disk I/O worker thread, open file handles) that GC would not
|
|
reclaim promptly. ``alloc_trim_strategy=DEFER`` only skips the eager allocator
|
|
reclaim (``to("meta")`` + ``cleanup_memory()``), leaving param storage for GC.
|
|
"""
|
|
wrapped = builder.build(device=target_device, dtype=dtype)
|
|
try:
|
|
yield wrapped
|
|
finally:
|
|
wrapped.teardown()
|
|
if alloc_trim_strategy == AllocatorTrimStrategy.TRIM:
|
|
wrapped.to("meta")
|
|
cleanup_memory()
|
|
|
|
|
|
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,
|
|
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
|
) -> Iterator[torch.Tensor]:
|
|
"""Wrap an iterator to release *model* memory (per ``alloc_trim_strategy``) once exhausted or abandoned."""
|
|
with gpu_model(model, alloc_trim_strategy=alloc_trim_strategy):
|
|
yield from it
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DiffusionStage
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class DiffusionStage:
|
|
"""Owns transformer lifecycle. Builds on each call, frees on exit.
|
|
Replaces the manual build-transformer / ``del transformer`` pattern that
|
|
every pipeline previously repeated.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
transformer_builder: ModelBuilderProtocol[LTXModelProtocol],
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
*,
|
|
quantization: QuantizationPolicy | None = None,
|
|
compilation_config: CompilationConfig | None = None,
|
|
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
|
) -> None:
|
|
"""Construct a stage from a single pre-built transformer ``builder``.
|
|
Holds only that builder plus build-time configuration (dtype, device,
|
|
quantization, compilation). Turning a checkpoint path + LoRA set into a
|
|
builder -- and choosing a :class:`StreamingModelBuilder` when offloading --
|
|
lives in :meth:`from_checkpoint`, which is how pipelines normally create a
|
|
stage. A :class:`StreamingModelBuilder` selects the block-streaming build
|
|
path; any other builder uses the standard (all-on-GPU) path.
|
|
``quantization`` and ``compilation_config`` are applied lazily on the
|
|
standard path; on the streaming path they are already baked into the
|
|
streaming builder by :meth:`from_checkpoint` and these fields are unused.
|
|
"""
|
|
self._transformer_builder = transformer_builder
|
|
self._dtype = dtype
|
|
self._device = device
|
|
self._quantization = quantization
|
|
self._compilation_config = compilation_config
|
|
self._alloc_trim_strategy = alloc_trim_strategy
|
|
|
|
@classmethod
|
|
def from_checkpoint( # noqa: PLR0913
|
|
cls,
|
|
checkpoint_path: str,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
|
|
quantization: QuantizationPolicy | None = None,
|
|
registry: Registry | None = None,
|
|
compilation_config: CompilationConfig | None = None,
|
|
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
|
offload_mode: OffloadMode = OffloadMode.NONE,
|
|
model_configurator: type[ModelConfigurator] = LTXModelConfigurator,
|
|
model_sd_ops: SDOps = LTXV_MODEL_COMFY_RENAMING_MAP,
|
|
) -> "DiffusionStage":
|
|
"""Build a stage from a checkpoint path and LoRA set.
|
|
Constructs a single transformer builder from ``checkpoint_path`` +
|
|
``loras`` + ``quantization`` and delegates to ``__init__``. When
|
|
``offload_mode != OffloadMode.NONE`` that builder is a
|
|
:class:`StreamingModelBuilder` (with quantization/compilation baked in and
|
|
its ``cpu_slots_count`` set for the requested mode); otherwise it is the
|
|
standard single-GPU builder. This is the high-level entry point used by
|
|
pipelines; ``__init__`` itself takes an already-built builder.
|
|
``model_configurator`` / ``model_sd_ops`` let callers (e.g. the audio-only
|
|
T2A pipeline) override the model class configurator and the state-dict key
|
|
mapping. A quantization policy that pins its own configurator takes
|
|
precedence over ``model_configurator``.
|
|
"""
|
|
# A quantization policy may pin its own configurator; otherwise use the one
|
|
# provided by the caller (defaults to the audio-video LTXModelConfigurator).
|
|
configurator = (
|
|
quantization.model_configurator
|
|
if quantization is not None and quantization.model_configurator is not None
|
|
else model_configurator
|
|
)
|
|
|
|
transformer_builder: ModelBuilderProtocol[LTXModelProtocol]
|
|
if offload_mode == OffloadMode.NONE:
|
|
transformer_builder = Builder(
|
|
model_path=checkpoint_path,
|
|
model_class_configurator=configurator,
|
|
model_sd_ops=model_sd_ops,
|
|
loras=tuple(loras),
|
|
registry=registry or DummyRegistry(),
|
|
)
|
|
else:
|
|
transformer_builder = cls._build_streaming_builder(
|
|
checkpoint_path=checkpoint_path,
|
|
configurator=configurator,
|
|
model_sd_ops=model_sd_ops,
|
|
loras=tuple(loras),
|
|
quantization=quantization,
|
|
registry=registry or DummyRegistry(),
|
|
offload_mode=offload_mode,
|
|
)
|
|
|
|
return cls(
|
|
transformer_builder,
|
|
dtype,
|
|
device,
|
|
quantization=quantization,
|
|
compilation_config=compilation_config,
|
|
alloc_trim_strategy=alloc_trim_strategy,
|
|
)
|
|
|
|
@staticmethod
|
|
def _build_streaming_builder(
|
|
*,
|
|
checkpoint_path: str,
|
|
configurator: type[ModelConfigurator],
|
|
model_sd_ops: SDOps,
|
|
loras: tuple[LoraPathStrengthAndSDOps, ...],
|
|
quantization: QuantizationPolicy | None,
|
|
registry: Registry,
|
|
offload_mode: OffloadMode,
|
|
) -> StreamingModelBuilder:
|
|
"""Construct the streaming transformer builder for an offloading stage.
|
|
Holds only raw config (``model_sd_ops`` / ``loras``); compilation and
|
|
quantization are applied at build time by :meth:`_prepared_builder`, exactly
|
|
as on the standard path -- so the builder's LoRA set stays raw and
|
|
:meth:`with_loras` swaps it consistently. ``cpu_slots_count`` is pinned for
|
|
the requested ``offload_mode`` (disk streaming uses a small slot count;
|
|
CPU/RAM streaming pins every block).
|
|
"""
|
|
# WeightsProvider currently only supports plain bf16 + fp8_cast LoRA fusion
|
|
# (no companion-key emission). Quantization policies that emit
|
|
# companion keys (e.g. ``.weight_scale``) cannot be streamed yet.
|
|
if quantization is not None and quantization.fuse_rule is not fp8_cast_fuse_rule:
|
|
raise ValueError(
|
|
"Block streaming is not supported with this quantization policy "
|
|
"(only bf16 and fp8_cast are currently supported)."
|
|
)
|
|
return StreamingModelBuilder(
|
|
model_class_configurator=configurator,
|
|
model_path=checkpoint_path,
|
|
model_sd_ops=model_sd_ops,
|
|
loras=loras,
|
|
registry=registry,
|
|
fuse_rule=quantization.fuse_rule if quantization is not None else bf16_fuse_rule,
|
|
blocks_attr="transformer_blocks",
|
|
blocks_prefix="transformer_blocks",
|
|
cpu_slots_count=DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None,
|
|
)
|
|
|
|
def with_attention(self, attention: AttentionFunction | AttentionCallable | None) -> "DiffusionStage":
|
|
"""Return a new ``DiffusionStage`` that pins the transformer build to ``attention``.
|
|
Functional: never mutates ``self``. The returned stage shares all other
|
|
configuration with the original; only the underlying builders' ``module_ops``
|
|
gain a ``set_attention_module_op(attention)`` entry so subsequent transformer
|
|
builds use that kernel. ``attention=None`` is a no-op (returns ``self``).
|
|
"""
|
|
if attention is None:
|
|
return self
|
|
op = set_attention_module_op(attention)
|
|
new = copy.copy(self)
|
|
new._transformer_builder = self._transformer_builder.with_module_ops(
|
|
(*self._transformer_builder.module_ops, op),
|
|
)
|
|
return new
|
|
|
|
def with_builder(self, builder: ModelBuilderProtocol[LTXModelProtocol]) -> "DiffusionStage":
|
|
"""Return a new ``DiffusionStage`` that builds its transformer from ``builder``.
|
|
Functional: never mutates ``self``; shares all other configuration (dtype, device,
|
|
quantization, compilation). Affects the standard (non-offload) build path.
|
|
"""
|
|
new = copy.copy(self)
|
|
new._transformer_builder = builder
|
|
return new
|
|
|
|
def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "DiffusionStage":
|
|
"""Return a new ``DiffusionStage`` built with exactly ``loras`` (replacing the current set)."""
|
|
return self.with_builder(self._transformer_builder.with_loras(loras))
|
|
|
|
def _prepared_builder(self) -> ModelBuilderProtocol[LTXModelProtocol]:
|
|
"""Return the configured builder with the stage's build-time ops applied.
|
|
Compilation and quantization live on the stage (not on the builder) and are
|
|
applied here, lazily, for both the standard and streaming paths. This keeps
|
|
the builder holding only raw sd_ops/module_ops/LoRAs, so ``with_loras`` /
|
|
``with_builder`` swap them consistently regardless of the build path. The
|
|
returned copy preserves the builder's concrete type (e.g. a
|
|
``StreamingModelBuilder`` stays one).
|
|
"""
|
|
builder = self._transformer_builder
|
|
sd_ops = builder.model_sd_ops
|
|
module_ops = builder.module_ops
|
|
loras = builder.loras
|
|
if self._compilation_config is not None:
|
|
number_of_layers = builder.model_config()["transformer"]["num_layers"]
|
|
sd_ops, module_ops, loras = _apply_compile_ops(
|
|
sd_ops, module_ops, loras, number_of_layers, self._compilation_config
|
|
)
|
|
if self._quantization is not None:
|
|
sd_ops, module_ops = _chain_quantization(sd_ops, module_ops, self._quantization)
|
|
builder = builder.with_fuse_rule(self._quantization.fuse_rule)
|
|
return builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
|
|
|
|
def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
|
|
target = device or self._device
|
|
return X0Model(self._prepared_builder().build(device=target, **kwargs)).to(target).eval()
|
|
|
|
@property
|
|
def _is_streaming(self) -> bool:
|
|
"""Whether the configured builder uses the block-streaming build path."""
|
|
return isinstance(self._transformer_builder, StreamingModelBuilder)
|
|
|
|
@contextmanager
|
|
def _streaming_transformer_ctx(self) -> Iterator[X0Model]:
|
|
builder = self._prepared_builder()
|
|
assert isinstance(builder, StreamingModelBuilder)
|
|
with _streaming_model(builder, self._device, self._dtype, self._alloc_trim_strategy) as streaming_wrapper:
|
|
yield X0Model(streaming_wrapper).eval()
|
|
|
|
def _transformer_ctx(self, **kwargs: object) -> AbstractContextManager:
|
|
if self._is_streaming:
|
|
return self._streaming_transformer_ctx()
|
|
return gpu_model(self._build_transformer(**kwargs), alloc_trim_strategy=self._alloc_trim_strategy)
|
|
|
|
def model_context(self, **kwargs: object) -> AbstractContextManager:
|
|
"""Build the transformer, yield it, then free its memory on exit.
|
|
Keyword arguments are forwarded to the underlying builder (e.g.
|
|
``video_tools`` required by ``TiledDataParallelBuilder``).
|
|
"""
|
|
return self._transformer_ctx(**kwargs)
|
|
|
|
def run( # noqa: PLR0913
|
|
self,
|
|
transformer: object,
|
|
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,
|
|
max_batch_size: int = 1,
|
|
) -> tuple[LatentState | None, LatentState | None]:
|
|
"""Run denoising with a pre-built transformer.
|
|
Same semantics as ``__call__`` but accepts a pre-built transformer so
|
|
the model can be shared across multiple calls (e.g. tiled inference
|
|
inside a single ``model_context()`` block). Audio supports
|
|
``ModalitySpec(frozen=True)`` to keep the latent unchanged throughout
|
|
denoising while still providing cross-modal context to the transformer.
|
|
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)
|
|
|
|
wrapped = BatchSplitAdapter(transformer, max_batch_size=max_batch_size) # type: ignore[arg-type]
|
|
video_state, audio_state = loop(
|
|
sigmas=sigmas,
|
|
video_state=video_state,
|
|
audio_state=audio_state,
|
|
stepper=stepper,
|
|
transformer=wrapped,
|
|
denoiser=denoiser,
|
|
)
|
|
|
|
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
|
|
|
|
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,
|
|
max_batch_size: int = 1,
|
|
) -> tuple[LatentState | None, LatentState | None]:
|
|
"""Build transformer -> run denoising loop -> free transformer.
|
|
Returns ``(video_state | None, audio_state | None)`` with cleared
|
|
conditionings and unpatchified latents for present modalities.
|
|
"""
|
|
# Build video_tools up front so it can be forwarded to the transformer
|
|
# context (required by TiledDataParallelBuilder in multi-GPU mode).
|
|
# `run()` rebuilds its own tools internally; the duplication is cheap.
|
|
video_tools: LatentTools | None = None
|
|
if video is not None:
|
|
pixel_shape = VideoPixelShape(batch=1, frames=frames, height=height, width=width, fps=fps)
|
|
v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
|
|
video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, fps)
|
|
|
|
mode = "streaming" if self._is_streaming else "standard"
|
|
logger.info("Building transformer (%s) from %s", mode, self._transformer_builder.checkpoint)
|
|
with self._transformer_ctx(video_tools=video_tools) as transformer:
|
|
logger.info(
|
|
"Running denoising loop (%d steps, %dx%d %d frames @ %.1f fps)",
|
|
len(sigmas) - 1,
|
|
width,
|
|
height,
|
|
frames,
|
|
fps,
|
|
)
|
|
return self.run(
|
|
transformer,
|
|
denoiser,
|
|
sigmas,
|
|
noiser,
|
|
width,
|
|
height,
|
|
frames,
|
|
fps,
|
|
video,
|
|
audio,
|
|
stepper,
|
|
loop,
|
|
max_batch_size,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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,
|
|
offload_mode: OffloadMode = OffloadMode.NONE,
|
|
text_encoder_builder: BuilderProtocol | None = None,
|
|
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
|
) -> None:
|
|
self._gemma_root = gemma_root
|
|
self._checkpoint_path = checkpoint_path
|
|
self._dtype = dtype
|
|
self._device = device
|
|
self._offload_mode = offload_mode
|
|
self._alloc_trim_strategy = alloc_trim_strategy
|
|
|
|
if text_encoder_builder is not None:
|
|
if offload_mode != OffloadMode.NONE:
|
|
raise ValueError(
|
|
"text_encoder_builder cannot be used with offload_mode != OffloadMode.NONE "
|
|
"because no streaming text encoder builder is available."
|
|
)
|
|
self._text_encoder_builder = text_encoder_builder
|
|
self._streaming_text_encoder_builder = None
|
|
else:
|
|
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._streaming_text_encoder_builder = StreamingModelBuilder(
|
|
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(),
|
|
blocks_attr="model.model.language_model.layers",
|
|
blocks_prefix="model.model.language_model.layers",
|
|
cpu_slots_count=DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None,
|
|
)
|
|
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 _build_text_encoder(self) -> torch.nn.Module:
|
|
"""Build the Gemma text encoder (non-streaming path)."""
|
|
return self._text_encoder_builder.build(device=self._device, dtype=self._dtype).eval()
|
|
|
|
def _build_embeddings_processor(self) -> EmbeddingsProcessor:
|
|
"""Build the embeddings processor on the target device."""
|
|
return self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).eval()
|
|
|
|
def _text_encoder_ctx(self) -> AbstractContextManager:
|
|
if self._offload_mode != OffloadMode.NONE:
|
|
return _streaming_model(
|
|
self._streaming_text_encoder_builder, self._device, self._dtype, self._alloc_trim_strategy
|
|
)
|
|
return gpu_model(self._build_text_encoder(), alloc_trim_strategy=self._alloc_trim_strategy)
|
|
|
|
def __call__(
|
|
self,
|
|
prompts: list[str],
|
|
*,
|
|
enhance_first_prompt: bool = False,
|
|
enhance_prompt_image: str | None = None,
|
|
enhance_prompt_seed: int = 42,
|
|
) -> list[EmbeddingsProcessorOutput]:
|
|
"""Encode *prompts* through Gemma -> embeddings processor, freeing each model after use."""
|
|
logger.info("Building text encoder from %s", self._gemma_root)
|
|
with self._text_encoder_ctx() 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(prompts)
|
|
logger.info("Text encoder done, building embeddings processor from %s", self._checkpoint_path)
|
|
|
|
with gpu_model(
|
|
self._build_embeddings_processor(), alloc_trim_strategy=self._alloc_trim_strategy
|
|
) as embeddings_processor:
|
|
result = [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
|
|
logger.info("Prompt encoding complete")
|
|
return result
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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,
|
|
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
|
) -> 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._alloc_trim_strategy = alloc_trim_strategy
|
|
|
|
def _build_encoder(self) -> VideoEncoder:
|
|
return self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()
|
|
|
|
def __call__(self, fn: Callable[[VideoEncoder], T]) -> T:
|
|
"""Build video encoder → call *fn(encoder)* → free encoder."""
|
|
with gpu_model(self._build_encoder(), alloc_trim_strategy=self._alloc_trim_strategy) 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,
|
|
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
|
) -> None:
|
|
self._upsampler_path = upsampler_path
|
|
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(),
|
|
)
|
|
self._alloc_trim_strategy = alloc_trim_strategy
|
|
|
|
def __call__(self, latent: torch.Tensor) -> torch.Tensor:
|
|
"""Upsample *latent* using video encoder + spatial upsampler, then free both."""
|
|
logger.info("Building video encoder + spatial upsampler from %s", self._upsampler_path)
|
|
with (
|
|
gpu_model(
|
|
self._encoder_builder.build(device=self._device, dtype=self._dtype).eval(),
|
|
alloc_trim_strategy=self._alloc_trim_strategy,
|
|
) as encoder,
|
|
gpu_model(
|
|
self._upsampler_builder.build(device=self._device, dtype=self._dtype).eval(),
|
|
alloc_trim_strategy=self._alloc_trim_strategy,
|
|
) 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,
|
|
memory_efficient: bool = True,
|
|
decoder_builder: BuilderProtocol | None = None,
|
|
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
|
) -> None:
|
|
self._checkpoint_path = checkpoint_path
|
|
self._dtype = dtype
|
|
self._device = device
|
|
if decoder_builder is not None:
|
|
self._decoder_builder = decoder_builder
|
|
else:
|
|
self._decoder_builder = Builder(
|
|
model_path=checkpoint_path,
|
|
model_class_configurator=VideoDecoderConfigurator,
|
|
model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
|
|
registry=registry or DummyRegistry(),
|
|
module_ops=(MEMORY_EFFICIENT_DECODE,) if memory_efficient else (),
|
|
)
|
|
self._alloc_trim_strategy = alloc_trim_strategy
|
|
|
|
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."""
|
|
logger.info("Building video decoder from %s", self._checkpoint_path)
|
|
decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()
|
|
return _cleanup_iter(
|
|
decoder.decode_video(latent, tiling_config, generator),
|
|
decoder,
|
|
alloc_trim_strategy=self._alloc_trim_strategy,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# AudioDecoder
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class AudioDecoder:
|
|
"""Owns audio decoder + vocoder lifecycle."""
|
|
|
|
def __init__(
|
|
self,
|
|
checkpoint_path: str,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
registry: Registry | None = None,
|
|
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
|
) -> None:
|
|
self._checkpoint_path = checkpoint_path
|
|
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(),
|
|
)
|
|
self._alloc_trim_strategy = alloc_trim_strategy
|
|
|
|
def __call__(self, latent: torch.Tensor) -> Audio:
|
|
"""Decode audio *latent* through VAE decoder + vocoder, then free both."""
|
|
logger.info("Building audio decoder + vocoder from %s", self._checkpoint_path)
|
|
# The vocoder always runs in fp32 (bf16 accumulation degrades spectral
|
|
# metrics). On CUDA/CPU it is stored in bf16 and autocast upcasts per-op to
|
|
# save memory; MPS has no fp32 autocast, so store it in fp32 directly and
|
|
# avoid the per-call cast. Negligible footprint for this small model.
|
|
vocoder_dtype = torch.float32 if self._device.type == "mps" else self._dtype
|
|
with (
|
|
gpu_model(
|
|
self._decoder_builder.build(device=self._device, dtype=self._dtype).eval(),
|
|
alloc_trim_strategy=self._alloc_trim_strategy,
|
|
) as decoder,
|
|
gpu_model(
|
|
self._vocoder_builder.build(device=self._device, dtype=vocoder_dtype).eval(),
|
|
alloc_trim_strategy=self._alloc_trim_strategy,
|
|
) 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,
|
|
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
|
) -> None:
|
|
self._dtype = dtype
|
|
self._device = device
|
|
self._alloc_trim_strategy = alloc_trim_strategy
|
|
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).eval(),
|
|
alloc_trim_strategy=self._alloc_trim_strategy,
|
|
) as encoder:
|
|
return fn(encoder)
|