Automated PR - 2026-03-30
This commit is contained in:
github-actions[bot]
@@ -1,35 +1,46 @@
|
||||
from ltx_pipelines.utils.blocks import (
|
||||
AudioConditioner,
|
||||
AudioDecoder,
|
||||
DiffusionStage,
|
||||
ImageConditioner,
|
||||
PromptEncoder,
|
||||
VideoDecoder,
|
||||
VideoUpsampler,
|
||||
)
|
||||
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser, GuidedDenoiser, SimpleDenoiser
|
||||
from ltx_pipelines.utils.helpers import (
|
||||
assert_resolution,
|
||||
cleanup_memory,
|
||||
combined_image_conditionings,
|
||||
denoise_audio_video,
|
||||
encode_prompts,
|
||||
generate_enhanced_prompt,
|
||||
get_device,
|
||||
multi_modal_guider_denoising_func,
|
||||
multi_modal_guider_factory_denoising_func,
|
||||
simple_denoising_func,
|
||||
image_conditionings_by_adding_guiding_latent,
|
||||
)
|
||||
from ltx_pipelines.utils.model_ledger import ModelLedger
|
||||
from ltx_pipelines.utils.samplers import (
|
||||
euler_denoising_loop,
|
||||
gradient_estimating_euler_denoising_loop,
|
||||
res2s_audio_video_denoising_loop,
|
||||
)
|
||||
from ltx_pipelines.utils.types import Denoiser, ModalitySpec
|
||||
|
||||
__all__ = [
|
||||
"ModelLedger",
|
||||
"AudioConditioner",
|
||||
"AudioDecoder",
|
||||
"Denoiser",
|
||||
"DiffusionStage",
|
||||
"FactoryGuidedDenoiser",
|
||||
"GuidedDenoiser",
|
||||
"ImageConditioner",
|
||||
"ModalitySpec",
|
||||
"PromptEncoder",
|
||||
"SimpleDenoiser",
|
||||
"VideoDecoder",
|
||||
"VideoUpsampler",
|
||||
"assert_resolution",
|
||||
"cleanup_memory",
|
||||
"combined_image_conditionings",
|
||||
"denoise_audio_video",
|
||||
"encode_prompts",
|
||||
"euler_denoising_loop",
|
||||
"generate_enhanced_prompt",
|
||||
"get_device",
|
||||
"gradient_estimating_euler_denoising_loop",
|
||||
"multi_modal_guider_denoising_func",
|
||||
"multi_modal_guider_factory_denoising_func",
|
||||
"image_conditionings_by_adding_guiding_latent",
|
||||
"res2s_audio_video_denoising_loop",
|
||||
"simple_denoising_func",
|
||||
]
|
||||
|
||||
@@ -173,6 +173,15 @@ def basic_arg_parser(
|
||||
required=True,
|
||||
help="Path to LTX-2 model checkpoint (.safetensors file).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-inference-steps",
|
||||
type=int,
|
||||
default=params.num_inference_steps,
|
||||
help=(
|
||||
f"Number of denoising steps in the diffusion sampling process. "
|
||||
f"Higher values improve quality but increase generation time (default: {params.num_inference_steps})."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gemma-root",
|
||||
type=resolve_path,
|
||||
@@ -197,6 +206,85 @@ def basic_arg_parser(
|
||||
default=params.seed,
|
||||
help=f"Random seed for reproducible generation (default: {params.seed}).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lora",
|
||||
dest="lora",
|
||||
action=LoraAction,
|
||||
nargs="+", # Accept 1-2 arguments per use (path and optional strength); validation is handled in LoraAction
|
||||
metavar=("PATH", "STRENGTH"),
|
||||
default=[],
|
||||
help=(
|
||||
"LoRA (Low-Rank Adaptation) model: path to model file and optional strength "
|
||||
f"(default strength: {DEFAULT_LORA_STRENGTH}). Can be specified multiple times. "
|
||||
"Example: --lora path/to/lora1.safetensors 0.8 --lora path/to/lora2.safetensors"
|
||||
),
|
||||
)
|
||||
|
||||
parser.add_argument("--enhance-prompt", action="store_true")
|
||||
|
||||
def _positive_int(value: str) -> int:
|
||||
try:
|
||||
int_value = int(value)
|
||||
if int_value < 1:
|
||||
raise argparse.ArgumentTypeError("must be >= 1")
|
||||
return int_value
|
||||
except ValueError as e:
|
||||
raise argparse.ArgumentTypeError(f"must be an integer, got {value}") from e
|
||||
|
||||
# Layer streaming
|
||||
parser.add_argument(
|
||||
"--streaming-prefetch-count",
|
||||
type=_positive_int,
|
||||
default=None,
|
||||
metavar="N",
|
||||
help=(
|
||||
"Enable layer streaming prefetching N layers ahead. "
|
||||
"At most 1 + N layers reside on GPU at once. "
|
||||
"Must be >= 1. Example: --streaming-prefetch-count 2"
|
||||
),
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--max-batch-size",
|
||||
type=_positive_int,
|
||||
default=1,
|
||||
metavar="N",
|
||||
help=(
|
||||
"Maximum batch size per transformer forward pass. "
|
||||
"Guided denoisers batch up to 4 guidance passes into a single call. "
|
||||
"Default 1 runs passes sequentially. Set to 4 to batch all passes "
|
||||
"together, which reduces layer-streaming PCIe transfers. "
|
||||
"Example: --max-batch-size 4"
|
||||
),
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
dest="quantization",
|
||||
action=QuantizationAction,
|
||||
nargs="+",
|
||||
metavar=("POLICY", "AMAX_PATH"),
|
||||
default=None,
|
||||
help=(
|
||||
f"Quantization policy: {', '.join(QUANTIZATION_POLICIES)}. "
|
||||
"fp8-cast uses FP8 casting with upcasting during inference. "
|
||||
"fp8-scaled-mm uses FP8 scaled matrix multiplication (optionally provide amax calibration file path). "
|
||||
"Example: --quantization fp8-cast or --quantization fp8-scaled-mm /path/to/amax.json"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--compile",
|
||||
action="store_true",
|
||||
help="Enable torch.compile for transformer blocks to optimize performance.",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def new_video_gen_arg_parser(
|
||||
params: PipelineParams = LTX_2_3_PARAMS,
|
||||
distilled: bool = False,
|
||||
) -> argparse.ArgumentParser:
|
||||
parser = basic_arg_parser(params=params, distilled=distilled)
|
||||
parser.add_argument(
|
||||
"--height",
|
||||
type=int,
|
||||
@@ -222,15 +310,6 @@ def basic_arg_parser(
|
||||
default=params.frame_rate,
|
||||
help=f"Frame rate of the generated video (fps) (default: {params.frame_rate}).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-inference-steps",
|
||||
type=int,
|
||||
default=params.num_inference_steps,
|
||||
help=(
|
||||
f"Number of denoising steps in the diffusion sampling process. "
|
||||
f"Higher values improve quality but increase generation time (default: {params.num_inference_steps})."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image",
|
||||
dest="images",
|
||||
@@ -247,42 +326,28 @@ def basic_arg_parser(
|
||||
"--image path/to/image2.jpg 160 0.9 0"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lora",
|
||||
dest="lora",
|
||||
action=LoraAction,
|
||||
nargs="+", # Accept 1-2 arguments per use (path and optional strength); validation is handled in LoraAction
|
||||
metavar=("PATH", "STRENGTH"),
|
||||
default=[],
|
||||
help=(
|
||||
"LoRA (Low-Rank Adaptation) model: path to model file and optional strength "
|
||||
f"(default strength: {DEFAULT_LORA_STRENGTH}). Can be specified multiple times. "
|
||||
"Example: --lora path/to/lora1.safetensors 0.8 --lora path/to/lora2.safetensors"
|
||||
),
|
||||
)
|
||||
|
||||
parser.add_argument("--enhance-prompt", action="store_true")
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
dest="quantization",
|
||||
action=QuantizationAction,
|
||||
nargs="+",
|
||||
metavar=("POLICY", "AMAX_PATH"),
|
||||
default=None,
|
||||
help=(
|
||||
f"Quantization policy: {', '.join(QUANTIZATION_POLICIES)}. "
|
||||
"fp8-cast uses FP8 casting with upcasting during inference. "
|
||||
"fp8-scaled-mm uses FP8 scaled matrix multiplication (optionally provide amax calibration file path). "
|
||||
"Example: --quantization fp8-cast or --quantization fp8-scaled-mm /path/to/amax.json"
|
||||
),
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def video_editing_arg_parser(
|
||||
distilled: bool = True,
|
||||
) -> argparse.ArgumentParser:
|
||||
"""Base argument parser for video-editing pipelines (retake, extension, inpainting, sticker movement).
|
||||
Uses the same actions and conventions as basic_arg_parser but only the args needed for editing
|
||||
(no height/width/num-frames; resolution comes from input video). Default is distilled checkpoint only.
|
||||
"""
|
||||
parser = basic_arg_parser(distilled=distilled)
|
||||
parser.add_argument("--video-path", type=resolve_path, required=True, help="Path to the source video.")
|
||||
parser.add_argument("--start-time", type=float, required=True, help="Start time of the region to regenerate (s).")
|
||||
parser.add_argument("--end-time", type=float, required=True, help="End time of the region to regenerate (s).")
|
||||
return parser
|
||||
|
||||
|
||||
def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
|
||||
video_guider = params.video_guider_params
|
||||
audio_guider = params.audio_guider_params
|
||||
parser = basic_arg_parser(params=params)
|
||||
parser = new_video_gen_arg_parser(params=params)
|
||||
parser.add_argument(
|
||||
"--negative-prompt",
|
||||
type=str,
|
||||
@@ -476,7 +541,7 @@ def hq_2_stage_arg_parser(params: PipelineParams = LTX_2_3_HQ_PARAMS) -> argpars
|
||||
|
||||
|
||||
def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
|
||||
parser = basic_arg_parser(params=params, distilled=True)
|
||||
parser = new_video_gen_arg_parser(params=params, distilled=True)
|
||||
parser.set_defaults(height=params.stage_2_height, width=params.stage_2_width)
|
||||
# Update help text to reflect 2-stage defaults
|
||||
for action in parser._actions:
|
||||
|
||||
@@ -0,0 +1,574 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,305 @@
|
||||
"""Flat denoiser classes — transformer received at call time, not stored.
|
||||
Three implementations of the :class:`~ltx_pipelines.utils.types.Denoiser` protocol:
|
||||
* :class:`SimpleDenoiser` — single transformer call, no guidance.
|
||||
* :class:`GuidedDenoiser` — static guiders, handles CFG + STG + isolated modality.
|
||||
* :class:`FactoryGuidedDenoiser` — resolves guiders per-step from sigma.
|
||||
``GuidedDenoiser`` and ``FactoryGuidedDenoiser`` share the core multi-pass
|
||||
logic via the module-level :func:`_guided_denoise` function, which batches
|
||||
all guidance passes into a single transformer call.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderFactory, MultiModalGuiderParams
|
||||
from ltx_core.guidance.perturbations import (
|
||||
BatchedPerturbationConfig,
|
||||
Perturbation,
|
||||
PerturbationConfig,
|
||||
PerturbationType,
|
||||
)
|
||||
from ltx_core.model.transformer import X0Model
|
||||
from ltx_core.types import LatentState
|
||||
from ltx_pipelines.utils.helpers import modality_from_latent_state
|
||||
|
||||
_POSITIVE_ONLY_GUIDER = MultiModalGuider(
|
||||
params=MultiModalGuiderParams(cfg_scale=1.0, stg_scale=0.0, modality_scale=1.0),
|
||||
)
|
||||
"""Guider that only runs the conditioned pass and returns cond unchanged."""
|
||||
|
||||
|
||||
def _ensure_guider(guider: MultiModalGuider | None) -> MultiModalGuider:
|
||||
"""Return the guider as-is, or a positive-only guider for absent modalities."""
|
||||
return guider if guider is not None else _POSITIVE_ONLY_GUIDER
|
||||
|
||||
|
||||
def _repeat_state(state: LatentState, n: int) -> LatentState:
|
||||
"""Repeat a ``LatentState`` *n* times along the batch dimension.
|
||||
``(B, ...) → (n*B, ...)`` by tiling the whole tensor n times, so the
|
||||
ordering is ``[item0, item1, ..., item0, item1, ...]`` — matching
|
||||
``torch.cat`` of n per-pass contexts.
|
||||
"""
|
||||
|
||||
def _repeat(t: torch.Tensor) -> torch.Tensor:
|
||||
repeats = [1] * t.dim()
|
||||
repeats[0] = n
|
||||
return t.repeat(repeats)
|
||||
|
||||
return LatentState(
|
||||
latent=_repeat(state.latent),
|
||||
denoise_mask=_repeat(state.denoise_mask),
|
||||
positions=_repeat(state.positions),
|
||||
clean_latent=_repeat(state.clean_latent),
|
||||
attention_mask=_repeat(state.attention_mask) if state.attention_mask is not None else None,
|
||||
)
|
||||
|
||||
|
||||
def _guided_denoise( # noqa: PLR0913
|
||||
transformer: X0Model,
|
||||
video_state: LatentState | None,
|
||||
audio_state: LatentState | None,
|
||||
sigma: torch.Tensor,
|
||||
video_guider: MultiModalGuider,
|
||||
audio_guider: MultiModalGuider,
|
||||
v_context: torch.Tensor | None,
|
||||
a_context: torch.Tensor | None,
|
||||
*,
|
||||
last_denoised_video: torch.Tensor | None,
|
||||
last_denoised_audio: torch.Tensor | None,
|
||||
step_index: int,
|
||||
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
|
||||
"""Core guided denoising — batches all guidance passes into one transformer call.
|
||||
Collects per-pass contexts first, then builds a single batched Modality
|
||||
per present modality via :func:`modality_from_latent_state`. When wrapped
|
||||
with :class:`~ltx_core.batch_split.BatchSplitAdapter`, the transformer may
|
||||
split this batch into sequential chunks internally.
|
||||
Guiders must not be ``None``. For absent modalities, callers should pass
|
||||
:data:`_POSITIVE_ONLY_GUIDER` (via :func:`_ensure_guider`) so that only
|
||||
the conditioned pass runs and ``calculate()`` returns cond unchanged.
|
||||
"""
|
||||
v_skip = video_guider.should_skip_step(step_index)
|
||||
a_skip = audio_guider.should_skip_step(step_index)
|
||||
|
||||
if v_skip and a_skip:
|
||||
return last_denoised_video, last_denoised_audio
|
||||
|
||||
if video_state is not None and v_context is None:
|
||||
raise ValueError("v_context is required when video_state is provided")
|
||||
if audio_state is not None and a_context is None:
|
||||
raise ValueError("a_context is required when audio_state is provided")
|
||||
# Define passes: (name, video_context, audio_context, perturbation_config).
|
||||
# Context is None for absent modalities — filtered out during collection.
|
||||
_pass = tuple[str, torch.Tensor | None, torch.Tensor | None, PerturbationConfig]
|
||||
passes: list[_pass] = [("cond", v_context, a_context, PerturbationConfig.empty())]
|
||||
|
||||
if video_guider.do_unconditional_generation() or audio_guider.do_unconditional_generation():
|
||||
if video_guider.do_unconditional_generation() and video_guider.negative_context is None:
|
||||
raise ValueError("Negative context is required for unconditioned denoising")
|
||||
if audio_guider.do_unconditional_generation() and audio_guider.negative_context is None:
|
||||
raise ValueError("Negative context is required for unconditioned denoising")
|
||||
v_neg = video_guider.negative_context if video_guider.negative_context is not None else v_context
|
||||
a_neg = audio_guider.negative_context if audio_guider.negative_context is not None else a_context
|
||||
passes.append(("uncond", v_neg, a_neg, PerturbationConfig.empty()))
|
||||
|
||||
stg_perturbations: list[Perturbation] = []
|
||||
if video_guider.do_perturbed_generation():
|
||||
stg_perturbations.append(
|
||||
Perturbation(type=PerturbationType.SKIP_VIDEO_SELF_ATTN, blocks=video_guider.params.stg_blocks)
|
||||
)
|
||||
if audio_guider.do_perturbed_generation():
|
||||
stg_perturbations.append(
|
||||
Perturbation(type=PerturbationType.SKIP_AUDIO_SELF_ATTN, blocks=audio_guider.params.stg_blocks)
|
||||
)
|
||||
if stg_perturbations:
|
||||
passes.append(("ptb", v_context, a_context, PerturbationConfig(stg_perturbations)))
|
||||
|
||||
if video_guider.do_isolated_modality_generation() or audio_guider.do_isolated_modality_generation():
|
||||
passes.append(
|
||||
(
|
||||
"mod",
|
||||
v_context,
|
||||
a_context,
|
||||
PerturbationConfig(
|
||||
[
|
||||
Perturbation(type=PerturbationType.SKIP_A2V_CROSS_ATTN, blocks=None),
|
||||
Perturbation(type=PerturbationType.SKIP_V2A_CROSS_ATTN, blocks=None),
|
||||
]
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
# Collect contexts, repeat states, and build batched modalities.
|
||||
pass_names = [name for name, _, _, _ in passes]
|
||||
ptb_configs = [ptb for _, _, _, ptb in passes]
|
||||
n = len(passes)
|
||||
|
||||
def _batched_sigma(state: LatentState) -> torch.Tensor:
|
||||
"""Expand scalar sigma to (n * B,) matching the repeated state."""
|
||||
return sigma.expand(state.latent.shape[0] * n)
|
||||
|
||||
batched_video = None
|
||||
if video_state is not None:
|
||||
v_context = torch.cat([vc for _, vc, _, _ in passes], dim=0)
|
||||
batched_video = modality_from_latent_state(
|
||||
_repeat_state(video_state, n),
|
||||
v_context,
|
||||
_batched_sigma(video_state),
|
||||
enabled=not v_skip,
|
||||
)
|
||||
|
||||
batched_audio = None
|
||||
if audio_state is not None:
|
||||
a_context = torch.cat([ac for _, _, ac, _ in passes], dim=0)
|
||||
batched_audio = modality_from_latent_state(
|
||||
_repeat_state(audio_state, n),
|
||||
a_context,
|
||||
_batched_sigma(audio_state),
|
||||
enabled=not a_skip,
|
||||
)
|
||||
|
||||
all_v, all_a = transformer(
|
||||
video=batched_video, audio=batched_audio, perturbations=BatchedPerturbationConfig(ptb_configs)
|
||||
)
|
||||
|
||||
# Split results back and combine via guiders.
|
||||
splits_v = list(all_v.chunk(n)) if all_v is not None else [0.0] * n
|
||||
splits_a = list(all_a.chunk(n)) if all_a is not None else [0.0] * n
|
||||
r = dict(zip(pass_names, zip(splits_v, splits_a, strict=True), strict=True))
|
||||
|
||||
cond_v, cond_a = r["cond"]
|
||||
uncond_v, uncond_a = r.get("uncond", (0.0, 0.0))
|
||||
ptb_v, ptb_a = r.get("ptb", (0.0, 0.0))
|
||||
mod_v, mod_a = r.get("mod", (0.0, 0.0))
|
||||
|
||||
denoised_video = last_denoised_video if v_skip else video_guider.calculate(cond_v, uncond_v, ptb_v, mod_v)
|
||||
denoised_audio = last_denoised_audio if a_skip else audio_guider.calculate(cond_a, uncond_a, ptb_a, mod_a)
|
||||
return denoised_video, denoised_audio
|
||||
|
||||
|
||||
class SimpleDenoiser:
|
||||
"""Single transformer call, no guidance.
|
||||
Passes ``None`` Modality for absent modalities.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
v_context: torch.Tensor | None,
|
||||
a_context: torch.Tensor | None,
|
||||
) -> None:
|
||||
self.v_context = v_context
|
||||
self.a_context = a_context
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
transformer: X0Model,
|
||||
video_state: LatentState | None,
|
||||
audio_state: LatentState | None,
|
||||
sigmas: torch.Tensor,
|
||||
step_index: int,
|
||||
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
|
||||
sigma = sigmas[step_index]
|
||||
pos_video = modality_from_latent_state(video_state, self.v_context, sigma) if video_state is not None else None
|
||||
pos_audio = modality_from_latent_state(audio_state, self.a_context, sigma) if audio_state is not None else None
|
||||
return transformer(video=pos_video, audio=pos_audio, perturbations=None)
|
||||
|
||||
|
||||
class GuidedDenoiser:
|
||||
"""Static guiders — handles CFG + STG + isolated modality.
|
||||
Context/guider can be ``None`` for absent modalities (a positive-only
|
||||
guider is substituted at call time).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
v_context: torch.Tensor | None,
|
||||
a_context: torch.Tensor | None,
|
||||
video_guider: MultiModalGuider | None = None,
|
||||
audio_guider: MultiModalGuider | None = None,
|
||||
) -> None:
|
||||
self.v_context = v_context
|
||||
self.a_context = a_context
|
||||
self.video_guider = video_guider
|
||||
self.audio_guider = audio_guider
|
||||
self._last_denoised_video: torch.Tensor | None = None
|
||||
self._last_denoised_audio: torch.Tensor | None = None
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
transformer: X0Model,
|
||||
video_state: LatentState | None,
|
||||
audio_state: LatentState | None,
|
||||
sigmas: torch.Tensor,
|
||||
step_index: int,
|
||||
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
|
||||
denoised_video, denoised_audio = _guided_denoise(
|
||||
transformer=transformer,
|
||||
video_state=video_state,
|
||||
audio_state=audio_state,
|
||||
sigma=sigmas[step_index],
|
||||
video_guider=_ensure_guider(self.video_guider),
|
||||
audio_guider=_ensure_guider(self.audio_guider),
|
||||
v_context=self.v_context,
|
||||
a_context=self.a_context,
|
||||
last_denoised_video=self._last_denoised_video,
|
||||
last_denoised_audio=self._last_denoised_audio,
|
||||
step_index=step_index,
|
||||
)
|
||||
self._last_denoised_video = denoised_video
|
||||
self._last_denoised_audio = denoised_audio
|
||||
return denoised_video, denoised_audio
|
||||
|
||||
|
||||
class FactoryGuidedDenoiser:
|
||||
"""Resolves guiders per-step from sigma, then delegates to shared guided logic."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
v_context: torch.Tensor | None,
|
||||
a_context: torch.Tensor | None,
|
||||
video_guider_factory: MultiModalGuiderFactory | None = None,
|
||||
audio_guider_factory: MultiModalGuiderFactory | None = None,
|
||||
) -> None:
|
||||
self.v_context = v_context
|
||||
self.a_context = a_context
|
||||
self.video_guider_factory = video_guider_factory
|
||||
self.audio_guider_factory = audio_guider_factory
|
||||
self._last_denoised_video: torch.Tensor | None = None
|
||||
self._last_denoised_audio: torch.Tensor | None = None
|
||||
self._sigma_vals_cached: list[float] | None = None
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
transformer: X0Model,
|
||||
video_state: LatentState | None,
|
||||
audio_state: LatentState | None,
|
||||
sigmas: torch.Tensor,
|
||||
step_index: int,
|
||||
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
|
||||
if self._sigma_vals_cached is None:
|
||||
self._sigma_vals_cached = sigmas.detach().cpu().tolist()
|
||||
sigma_val = self._sigma_vals_cached[step_index]
|
||||
|
||||
video_guider = _ensure_guider(
|
||||
self.video_guider_factory.build_from_sigma(sigma_val) if self.video_guider_factory else None
|
||||
)
|
||||
audio_guider = _ensure_guider(
|
||||
(self.audio_guider_factory or self.video_guider_factory).build_from_sigma(sigma_val)
|
||||
if self.video_guider_factory or self.audio_guider_factory
|
||||
else None
|
||||
)
|
||||
|
||||
denoised_video, denoised_audio = _guided_denoise(
|
||||
transformer=transformer,
|
||||
video_state=video_state,
|
||||
audio_state=audio_state,
|
||||
sigma=sigmas[step_index],
|
||||
video_guider=video_guider,
|
||||
audio_guider=audio_guider,
|
||||
v_context=self.v_context,
|
||||
a_context=self.a_context,
|
||||
last_denoised_video=self._last_denoised_video,
|
||||
last_denoised_audio=self._last_denoised_audio,
|
||||
step_index=step_index,
|
||||
)
|
||||
self._last_denoised_video = denoised_video
|
||||
self._last_denoised_audio = denoised_audio
|
||||
return denoised_video, denoised_audio
|
||||
@@ -0,0 +1,30 @@
|
||||
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()
|
||||
@@ -1,41 +1,35 @@
|
||||
import gc
|
||||
import logging
|
||||
from dataclasses import replace
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderFactory
|
||||
from ltx_core.components.noisers import Noiser
|
||||
from ltx_core.components.protocols import DiffusionStepProtocol, GuiderProtocol
|
||||
from ltx_core.conditioning import (
|
||||
ConditioningItem,
|
||||
VideoConditionByKeyframeIndex,
|
||||
VideoConditionByLatentIndex,
|
||||
)
|
||||
from ltx_core.guidance.perturbations import (
|
||||
BatchedPerturbationConfig,
|
||||
Perturbation,
|
||||
PerturbationConfig,
|
||||
PerturbationType,
|
||||
)
|
||||
from ltx_core.model.transformer import Modality, X0Model
|
||||
from ltx_core.model.video_vae import VideoEncoder
|
||||
from ltx_core.model.audio_vae import encode_audio
|
||||
from ltx_core.model.transformer import Modality
|
||||
from ltx_core.model.video_vae import TilingConfig, VideoEncoder
|
||||
from ltx_core.text_encoders.gemma import GemmaTextEncoder
|
||||
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
|
||||
from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
|
||||
from ltx_core.tools import LatentTools
|
||||
from ltx_core.types import AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
|
||||
from ltx_pipelines.utils.args import ImageConditioningInput
|
||||
from ltx_pipelines.utils.media_io import decode_image, load_image_conditioning, resize_aspect_ratio_preserving
|
||||
from ltx_pipelines.utils.types import (
|
||||
DenoisingFunc,
|
||||
DenoisingLoopFunc,
|
||||
PipelineComponents,
|
||||
from ltx_pipelines.utils.media_io import (
|
||||
decode_audio_from_file,
|
||||
decode_image,
|
||||
decode_video_from_file,
|
||||
get_videostream_fps,
|
||||
load_image_and_preprocess,
|
||||
resize_aspect_ratio_preserving,
|
||||
video_preprocess,
|
||||
)
|
||||
|
||||
|
||||
def get_device() -> torch.device:
|
||||
if torch.cuda.is_available():
|
||||
return torch.device("cuda")
|
||||
return torch.device("cuda", torch.cuda.current_device())
|
||||
return torch.device("cpu")
|
||||
|
||||
|
||||
@@ -45,45 +39,89 @@ def cleanup_memory() -> None:
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
def encode_prompts(
|
||||
prompts: list[str],
|
||||
model_ledger: object,
|
||||
*,
|
||||
enhance_prompt_image: str | None = None,
|
||||
enhance_prompt_seed: int = 42,
|
||||
enhance_first_prompt: bool = False,
|
||||
) -> list[EmbeddingsProcessorOutput]:
|
||||
"""Encode prompts through Gemma → embeddings processor, freeing each after use.
|
||||
Loads the text encoder from *model_ledger*, optionally enhances the first
|
||||
prompt, encodes all *prompts*, frees the text encoder, then loads the
|
||||
embeddings processor to produce the final outputs. Because the text encoder
|
||||
is loaded and freed entirely within this function, there are no lingering
|
||||
references that could prevent GPU memory reclamation.
|
||||
Args:
|
||||
prompts: Text prompts to encode.
|
||||
model_ledger: ModelLedger instance (used to load text encoder and embeddings processor).
|
||||
enhance_prompt_image: Optional image path for prompt enhancement.
|
||||
enhance_prompt_seed: Seed for prompt enhancement (default 42).
|
||||
enhance_first_prompt: If True, enhance ``prompts[0]`` before encoding.
|
||||
Returns:
|
||||
List of EmbeddingsProcessorOutput, one per prompt.
|
||||
"""
|
||||
text_encoder = model_ledger.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]
|
||||
torch.cuda.synchronize()
|
||||
del text_encoder
|
||||
cleanup_memory()
|
||||
def _conform_latent_length(latent: torch.Tensor, expected_frames_count: int) -> torch.Tensor:
|
||||
actual_frames = latent.shape[2]
|
||||
if actual_frames > expected_frames_count:
|
||||
latent = latent[:, :, :expected_frames_count]
|
||||
elif actual_frames < expected_frames_count:
|
||||
shape_as_list = list(latent.shape)
|
||||
shape_as_list[2] = expected_frames_count - actual_frames
|
||||
pad = torch.zeros(
|
||||
shape_as_list,
|
||||
device=latent.device,
|
||||
dtype=latent.dtype,
|
||||
)
|
||||
latent = torch.cat([latent, pad], dim=2)
|
||||
return latent
|
||||
|
||||
embeddings_processor = model_ledger.gemma_embeddings_processor()
|
||||
results: list[EmbeddingsProcessorOutput] = [
|
||||
embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs
|
||||
]
|
||||
del embeddings_processor
|
||||
cleanup_memory()
|
||||
return results
|
||||
|
||||
def video_latent_from_file(
|
||||
video_encoder: VideoEncoder,
|
||||
file_path: str,
|
||||
output_shape: VideoPixelShape,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
start_time: float = 0.0,
|
||||
max_duration: float | None = None,
|
||||
tiling_config: TilingConfig | None = None,
|
||||
) -> torch.Tensor | None:
|
||||
"""Load video from a file, and construct the video latent conforming to video output shape.
|
||||
Args:
|
||||
video_encoder: Model used to encode pixel frames to latent space.
|
||||
file_path: Path to the video file.
|
||||
output_shape: Target pixel shape (height, width, frames, fps) for the conditioning.
|
||||
device: Device to run the encoder and hold tensors on.
|
||||
dtype: Dtype for the output latents.
|
||||
start_time: Start time in seconds to begin reading the video (default 0.0).
|
||||
max_duration: Maximum duration in seconds. If None, uses output_shape.frames at
|
||||
output_shape.fps (default None).
|
||||
tiling_config: Tiling configuration for the encoder. Defaults to TilingConfig.default().
|
||||
Returns:
|
||||
Encoded video latents of shape (1, C, T, H, W) with T = required_latent_frames, or
|
||||
None (currently this function always returns a tensor).
|
||||
"""
|
||||
fps = get_videostream_fps(file_path)
|
||||
if fps != output_shape.fps:
|
||||
raise ValueError(f"Input video FPS {fps} does not match output FPS {output_shape.fps}, not supported")
|
||||
max_duration = max_duration or output_shape.frames / fps
|
||||
frame_gen = decode_video_from_file(path=file_path, device=device, start_time=start_time, max_duration=max_duration)
|
||||
frames = video_preprocess(frame_gen, output_shape.height, output_shape.width, dtype, device)
|
||||
latents = video_encoder.tiled_encode(frames, tiling_config or TilingConfig.default())
|
||||
required_latent_frames = VideoLatentShape.from_pixel_shape(output_shape).frames
|
||||
return _conform_latent_length(latents, required_latent_frames)
|
||||
|
||||
|
||||
def audio_latent_from_file(
|
||||
audio_encoder: torch.nn.Module,
|
||||
file_path: str,
|
||||
output_shape: VideoPixelShape,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
start_time: float = 0.0,
|
||||
max_duration: float | None = None,
|
||||
) -> torch.Tensor | None:
|
||||
"""Load audio from a file, and construct the audio latent conforming to video output shape.
|
||||
Args:
|
||||
audio_encoder: Model used to encode audio to latent space.
|
||||
file_path: Path to the audio or video file containing an audio stream.
|
||||
output_shape: Target video pixel shape; used to derive required latent frames
|
||||
and, when max_duration is None, the audio duration (output_shape.frames / fps).
|
||||
device: Device to run the encoder and hold tensors on.
|
||||
dtype: Dtype for the output latents.
|
||||
start_time: Start time in seconds to begin reading the audio (default 0.0).
|
||||
max_duration: Maximum duration in seconds. If None, uses the full span implied
|
||||
by output_shape (default None).
|
||||
Returns:
|
||||
Encoded audio latents of shape (1, C, T, ...) with T = required_latent_frames, or
|
||||
None if the file has no audio stream.
|
||||
"""
|
||||
max_duration = max_duration or output_shape.frames / output_shape.fps
|
||||
audio_in = decode_audio_from_file(file_path, device, start_time, max_duration)
|
||||
if audio_in is None:
|
||||
return None
|
||||
latents = encode_audio(audio_in, audio_encoder, None).to(device, dtype)
|
||||
required_latent_frames = AudioLatentShape.from_video_pixel_shape(output_shape).frames
|
||||
return _conform_latent_length(latents, required_latent_frames)
|
||||
|
||||
|
||||
def combined_image_conditionings(
|
||||
@@ -98,7 +136,7 @@ def combined_image_conditionings(
|
||||
and using other encoded images as the keyframe conditionings."""
|
||||
conditionings = []
|
||||
for img in images:
|
||||
image = load_image_conditioning(
|
||||
image = load_image_and_preprocess(
|
||||
image_path=img.path,
|
||||
height=height,
|
||||
width=width,
|
||||
@@ -133,7 +171,7 @@ def image_conditionings_by_replacing_latent(
|
||||
) -> list[ConditioningItem]:
|
||||
conditionings = []
|
||||
for img in images:
|
||||
image = load_image_conditioning(
|
||||
image = load_image_and_preprocess(
|
||||
image_path=img.path,
|
||||
height=height,
|
||||
width=width,
|
||||
@@ -163,7 +201,7 @@ def image_conditionings_by_adding_guiding_latent(
|
||||
) -> list[ConditioningItem]:
|
||||
conditionings = []
|
||||
for img in images:
|
||||
image = load_image_conditioning(
|
||||
image = load_image_and_preprocess(
|
||||
image_path=img.path,
|
||||
height=height,
|
||||
width=width,
|
||||
@@ -178,72 +216,6 @@ def image_conditionings_by_adding_guiding_latent(
|
||||
return conditionings
|
||||
|
||||
|
||||
def noise_video_state(
|
||||
output_shape: VideoPixelShape,
|
||||
noiser: Noiser,
|
||||
conditionings: list[ConditioningItem],
|
||||
components: PipelineComponents,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
noise_scale: float = 1.0,
|
||||
initial_latent: torch.Tensor | None = None,
|
||||
) -> tuple[LatentState, VideoLatentTools]:
|
||||
"""Initialize and noise a video latent state for the diffusion pipeline.
|
||||
Creates a video latent state from the output shape, applies conditionings,
|
||||
and adds noise using the provided noiser. Returns the noised state and
|
||||
video latent tools for further processing. If initial_latent is provided, it will be used to create the initial
|
||||
state, otherwise an empty initial state will be created.
|
||||
"""
|
||||
video_latent_shape = VideoLatentShape.from_pixel_shape(
|
||||
shape=output_shape,
|
||||
latent_channels=components.video_latent_channels,
|
||||
scale_factors=components.video_scale_factors,
|
||||
)
|
||||
video_tools = VideoLatentTools(components.video_patchifier, video_latent_shape, output_shape.fps)
|
||||
video_state = create_noised_state(
|
||||
tools=video_tools,
|
||||
conditionings=conditionings,
|
||||
noiser=noiser,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
noise_scale=noise_scale,
|
||||
initial_latent=initial_latent,
|
||||
)
|
||||
|
||||
return video_state, video_tools
|
||||
|
||||
|
||||
def noise_audio_state(
|
||||
output_shape: VideoPixelShape,
|
||||
noiser: Noiser,
|
||||
conditionings: list[ConditioningItem],
|
||||
components: PipelineComponents,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
noise_scale: float = 1.0,
|
||||
initial_latent: torch.Tensor | None = None,
|
||||
) -> tuple[LatentState, AudioLatentTools]:
|
||||
"""Initialize and noise an audio latent state for the diffusion pipeline.
|
||||
Creates an audio latent state from the output shape, applies conditionings,
|
||||
and adds noise using the provided noiser. Returns the noised state and
|
||||
audio latent tools for further processing. If initial_latent is provided, it will be used to create the initial
|
||||
state, otherwise an empty initial state will be created.
|
||||
"""
|
||||
audio_latent_shape = AudioLatentShape.from_video_pixel_shape(output_shape)
|
||||
audio_tools = AudioLatentTools(components.audio_patchifier, audio_latent_shape)
|
||||
audio_state = create_noised_state(
|
||||
tools=audio_tools,
|
||||
conditionings=conditionings,
|
||||
noiser=noiser,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
noise_scale=noise_scale,
|
||||
initial_latent=initial_latent,
|
||||
)
|
||||
|
||||
return audio_state, audio_tools
|
||||
|
||||
|
||||
def create_noised_state(
|
||||
tools: LatentTools,
|
||||
conditionings: list[ConditioningItem],
|
||||
@@ -308,301 +280,14 @@ def timesteps_from_mask(denoise_mask: torch.Tensor, sigma: float | torch.Tensor)
|
||||
"""Compute timesteps from a denoise mask and sigma value.
|
||||
Multiplies the denoise mask by sigma to produce timesteps for each position
|
||||
in the latent state. Areas where the mask is 0 will have zero timesteps.
|
||||
When sigma is ``(B,)`` it is reshaped to ``(B, 1, ...)`` so the batch
|
||||
dimension aligns correctly with ``denoise_mask``.
|
||||
"""
|
||||
if isinstance(sigma, torch.Tensor) and sigma.dim() == 1:
|
||||
sigma = sigma.view(-1, *([1] * (denoise_mask.dim() - 1)))
|
||||
return denoise_mask * sigma
|
||||
|
||||
|
||||
def simple_denoising_func(
|
||||
video_context: torch.Tensor, audio_context: torch.Tensor, transformer: X0Model
|
||||
) -> DenoisingFunc:
|
||||
def simple_denoising_step(
|
||||
video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
sigma = sigmas[step_index]
|
||||
pos_video = modality_from_latent_state(video_state, video_context, sigma)
|
||||
pos_audio = modality_from_latent_state(audio_state, audio_context, sigma)
|
||||
|
||||
denoised_video, denoised_audio = transformer(video=pos_video, audio=pos_audio, perturbations=None)
|
||||
return denoised_video, denoised_audio
|
||||
|
||||
return simple_denoising_step
|
||||
|
||||
|
||||
def guider_denoising_func(
|
||||
guider: GuiderProtocol,
|
||||
v_context_p: torch.Tensor,
|
||||
v_context_n: torch.Tensor,
|
||||
a_context_p: torch.Tensor,
|
||||
a_context_n: torch.Tensor,
|
||||
transformer: X0Model,
|
||||
) -> DenoisingFunc:
|
||||
def guider_denoising_step(
|
||||
video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
sigma = sigmas[step_index]
|
||||
pos_video = modality_from_latent_state(video_state, v_context_p, sigma)
|
||||
pos_audio = modality_from_latent_state(audio_state, a_context_p, sigma)
|
||||
|
||||
denoised_video, denoised_audio = transformer(video=pos_video, audio=pos_audio, perturbations=None)
|
||||
if guider.enabled():
|
||||
neg_video = modality_from_latent_state(video_state, v_context_n, sigma)
|
||||
neg_audio = modality_from_latent_state(audio_state, a_context_n, sigma)
|
||||
|
||||
neg_denoised_video, neg_denoised_audio = transformer(video=neg_video, audio=neg_audio, perturbations=None)
|
||||
|
||||
denoised_video = denoised_video + guider.delta(denoised_video, neg_denoised_video)
|
||||
denoised_audio = denoised_audio + guider.delta(denoised_audio, neg_denoised_audio)
|
||||
|
||||
return denoised_video, denoised_audio
|
||||
|
||||
return guider_denoising_step
|
||||
|
||||
|
||||
def multi_modal_guider_denoising_func(
|
||||
video_guider: MultiModalGuider,
|
||||
audio_guider: MultiModalGuider,
|
||||
v_context: torch.Tensor,
|
||||
a_context: torch.Tensor,
|
||||
transformer: X0Model,
|
||||
*,
|
||||
last_denoised_video: torch.Tensor | None = None,
|
||||
last_denoised_audio: torch.Tensor | None = None,
|
||||
) -> DenoisingFunc:
|
||||
def guider_denoising_step(
|
||||
video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
nonlocal last_denoised_video, last_denoised_audio
|
||||
|
||||
if video_guider.should_skip_step(step_index) and audio_guider.should_skip_step(step_index):
|
||||
return last_denoised_video, last_denoised_audio
|
||||
|
||||
sigma = sigmas[step_index]
|
||||
pos_video_modality = modality_from_latent_state(
|
||||
video_state, v_context, sigma, enabled=not video_guider.should_skip_step(step_index)
|
||||
)
|
||||
pos_audio_modality = modality_from_latent_state(
|
||||
audio_state, a_context, sigma, enabled=not audio_guider.should_skip_step(step_index)
|
||||
)
|
||||
|
||||
denoised_video, denoised_audio = transformer(
|
||||
video=pos_video_modality, audio=pos_audio_modality, perturbations=None
|
||||
)
|
||||
neg_denoised_video, neg_denoised_audio = 0.0, 0.0
|
||||
if video_guider.do_unconditional_generation() or audio_guider.do_unconditional_generation():
|
||||
if video_guider.do_unconditional_generation() and video_guider.negative_context is None:
|
||||
raise ValueError("Negative context is required for unconditioned denoising")
|
||||
if audio_guider.do_unconditional_generation() and audio_guider.negative_context is None:
|
||||
raise ValueError("Negative context is required for unconditioned denoising")
|
||||
neg_video_modality = modality_from_latent_state(
|
||||
video_state,
|
||||
video_guider.negative_context
|
||||
if video_guider.negative_context is not None
|
||||
else pos_video_modality.context,
|
||||
sigma,
|
||||
)
|
||||
neg_audio_modality = modality_from_latent_state(
|
||||
audio_state,
|
||||
audio_guider.negative_context
|
||||
if audio_guider.negative_context is not None
|
||||
else pos_audio_modality.context,
|
||||
sigma,
|
||||
)
|
||||
|
||||
neg_denoised_video, neg_denoised_audio = transformer(
|
||||
video=neg_video_modality, audio=neg_audio_modality, perturbations=None
|
||||
)
|
||||
|
||||
ptb_denoised_video, ptb_denoised_audio = 0.0, 0.0
|
||||
if video_guider.do_perturbed_generation() or audio_guider.do_perturbed_generation():
|
||||
perturbations = []
|
||||
if video_guider.do_perturbed_generation():
|
||||
perturbations.append(
|
||||
Perturbation(type=PerturbationType.SKIP_VIDEO_SELF_ATTN, blocks=video_guider.params.stg_blocks)
|
||||
)
|
||||
if audio_guider.do_perturbed_generation():
|
||||
perturbations.append(
|
||||
Perturbation(type=PerturbationType.SKIP_AUDIO_SELF_ATTN, blocks=audio_guider.params.stg_blocks)
|
||||
)
|
||||
perturbation_config = PerturbationConfig(perturbations=perturbations)
|
||||
ptb_denoised_video, ptb_denoised_audio = transformer(
|
||||
video=pos_video_modality,
|
||||
audio=pos_audio_modality,
|
||||
perturbations=BatchedPerturbationConfig(perturbations=[perturbation_config]),
|
||||
)
|
||||
|
||||
mod_denoised_video, mod_denoised_audio = 0.0, 0.0
|
||||
if video_guider.do_isolated_modality_generation() or audio_guider.do_isolated_modality_generation():
|
||||
perturbations = [
|
||||
Perturbation(type=PerturbationType.SKIP_A2V_CROSS_ATTN, blocks=None),
|
||||
Perturbation(type=PerturbationType.SKIP_V2A_CROSS_ATTN, blocks=None),
|
||||
]
|
||||
perturbation_config = PerturbationConfig(perturbations=perturbations)
|
||||
mod_denoised_video, mod_denoised_audio = transformer(
|
||||
video=pos_video_modality,
|
||||
audio=pos_audio_modality,
|
||||
perturbations=BatchedPerturbationConfig(perturbations=[perturbation_config]),
|
||||
)
|
||||
|
||||
if video_guider.should_skip_step(step_index):
|
||||
denoised_video = last_denoised_video
|
||||
else:
|
||||
denoised_video = video_guider.calculate(
|
||||
denoised_video, neg_denoised_video, ptb_denoised_video, mod_denoised_video
|
||||
)
|
||||
|
||||
if audio_guider.should_skip_step(step_index):
|
||||
denoised_audio = last_denoised_audio
|
||||
else:
|
||||
denoised_audio = audio_guider.calculate(
|
||||
denoised_audio, neg_denoised_audio, ptb_denoised_audio, mod_denoised_audio
|
||||
)
|
||||
|
||||
last_denoised_video = denoised_video
|
||||
last_denoised_audio = denoised_audio
|
||||
|
||||
return denoised_video, denoised_audio
|
||||
|
||||
return guider_denoising_step
|
||||
|
||||
|
||||
def multi_modal_guider_factory_denoising_func(
|
||||
video_guider_factory: MultiModalGuiderFactory,
|
||||
audio_guider_factory: MultiModalGuiderFactory | None,
|
||||
v_context: torch.Tensor,
|
||||
a_context: torch.Tensor,
|
||||
transformer: X0Model,
|
||||
) -> DenoisingFunc:
|
||||
"""Resolve guiders per step via factory.build_from_sigma, then multi_modal_guider_denoising_func."""
|
||||
last_denoised_video: torch.Tensor | None = None
|
||||
last_denoised_audio: torch.Tensor | None = None
|
||||
sigma_vals_cached: list[float] | None = None
|
||||
|
||||
def guider_denoising_step(
|
||||
video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
nonlocal last_denoised_video, last_denoised_audio, sigma_vals_cached
|
||||
if sigma_vals_cached is None:
|
||||
sigma_vals_cached = sigmas.detach().cpu().tolist()
|
||||
sigma_val = sigma_vals_cached[step_index]
|
||||
video_guider = video_guider_factory.build_from_sigma(sigma_val)
|
||||
audio_guider = (audio_guider_factory or video_guider_factory).build_from_sigma(sigma_val)
|
||||
denoise_fn = multi_modal_guider_denoising_func(
|
||||
video_guider,
|
||||
audio_guider,
|
||||
v_context,
|
||||
a_context,
|
||||
transformer,
|
||||
last_denoised_video=last_denoised_video,
|
||||
last_denoised_audio=last_denoised_audio,
|
||||
)
|
||||
denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_index)
|
||||
last_denoised_video, last_denoised_audio = denoised_video, denoised_audio
|
||||
return denoised_video, denoised_audio
|
||||
|
||||
return guider_denoising_step
|
||||
|
||||
|
||||
def denoise_audio_video( # noqa: PLR0913
|
||||
output_shape: VideoPixelShape,
|
||||
conditionings: list[ConditioningItem],
|
||||
noiser: Noiser,
|
||||
sigmas: torch.Tensor,
|
||||
stepper: DiffusionStepProtocol,
|
||||
denoising_loop_fn: DenoisingLoopFunc,
|
||||
components: PipelineComponents,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
noise_scale: float = 1.0,
|
||||
initial_video_latent: torch.Tensor | None = None,
|
||||
initial_audio_latent: torch.Tensor | None = None,
|
||||
) -> tuple[LatentState, LatentState]:
|
||||
video_state, video_tools = noise_video_state(
|
||||
output_shape=output_shape,
|
||||
noiser=noiser,
|
||||
conditionings=conditionings,
|
||||
components=components,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
noise_scale=noise_scale,
|
||||
initial_latent=initial_video_latent,
|
||||
)
|
||||
audio_state, audio_tools = noise_audio_state(
|
||||
output_shape=output_shape,
|
||||
noiser=noiser,
|
||||
conditionings=[],
|
||||
components=components,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
noise_scale=noise_scale,
|
||||
initial_latent=initial_audio_latent,
|
||||
)
|
||||
|
||||
video_state, audio_state = denoising_loop_fn(
|
||||
sigmas,
|
||||
video_state,
|
||||
audio_state,
|
||||
stepper,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
return video_state, audio_state
|
||||
|
||||
|
||||
def denoise_video_only( # noqa: PLR0913
|
||||
output_shape: VideoPixelShape,
|
||||
conditionings: list[ConditioningItem],
|
||||
noiser: Noiser,
|
||||
sigmas: torch.Tensor,
|
||||
stepper: DiffusionStepProtocol,
|
||||
denoising_loop_fn: DenoisingLoopFunc,
|
||||
components: PipelineComponents,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
noise_scale: float = 1.0,
|
||||
initial_video_latent: torch.Tensor | None = None,
|
||||
initial_audio_latent: torch.Tensor | None = None,
|
||||
) -> LatentState:
|
||||
video_state, video_tools = noise_video_state(
|
||||
output_shape=output_shape,
|
||||
noiser=noiser,
|
||||
conditionings=conditionings,
|
||||
components=components,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
noise_scale=noise_scale,
|
||||
initial_latent=initial_video_latent,
|
||||
)
|
||||
|
||||
audio_state, _ = noise_audio_state(
|
||||
output_shape=output_shape,
|
||||
noiser=noiser,
|
||||
conditionings=[],
|
||||
components=components,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
noise_scale=0.0,
|
||||
initial_latent=initial_audio_latent,
|
||||
)
|
||||
|
||||
audio_state = replace(audio_state, denoise_mask=torch.zeros_like(audio_state.denoise_mask))
|
||||
|
||||
video_state, audio_state = denoising_loop_fn(
|
||||
sigmas,
|
||||
video_state,
|
||||
audio_state,
|
||||
stepper,
|
||||
)
|
||||
|
||||
video_state = video_tools.clear_conditioning(video_state)
|
||||
video_state = video_tools.unpatchify(video_state)
|
||||
|
||||
return video_state
|
||||
|
||||
|
||||
_UNICODE_REPLACEMENTS = str.maketrans("\u2018\u2019\u201c\u201d\u2014\u2013\u00a0\u2032\u2212", "''\"\"-- '-")
|
||||
|
||||
|
||||
|
||||
@@ -12,7 +12,7 @@ from PIL import Image
|
||||
from torch._prims_common import DeviceLikeType
|
||||
from tqdm import tqdm
|
||||
|
||||
from ltx_core.types import Audio
|
||||
from ltx_core.types import Audio, VideoPixelShape
|
||||
from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -79,7 +79,7 @@ def normalize_latent(latent: torch.Tensor, device: torch.device, dtype: torch.dt
|
||||
return (latent / 127.5 - 1.0).to(device=device, dtype=dtype)
|
||||
|
||||
|
||||
def load_image_conditioning(
|
||||
def load_image_and_preprocess(
|
||||
image_path: str,
|
||||
height: int,
|
||||
width: int,
|
||||
@@ -99,14 +99,23 @@ def load_image_conditioning(
|
||||
return image
|
||||
|
||||
|
||||
def load_video_conditioning(
|
||||
video_path: str, height: int, width: int, frame_cap: int, dtype: torch.dtype, device: torch.device
|
||||
def video_preprocess(
|
||||
frames: Generator[torch.Tensor],
|
||||
height: int,
|
||||
width: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
"""Preprocesses a video frame generator for conditioning.
|
||||
Args:
|
||||
frames: Generator of video frames as tensors of shape (1, H, W, C), dtype uint8.
|
||||
height: Target height in pixels.
|
||||
width: Target width in pixels.
|
||||
dtype: Target dtype for the output tensor.
|
||||
device: Target device for the output tensor.
|
||||
Returns:
|
||||
Tensor of shape (1, C, F, height, width) with values in [-1, 1].
|
||||
"""
|
||||
Loads a video from a path and preprocesses it for conditioning.
|
||||
Note: The video is resized to the nearest multiple of 2 for compatibility with video codecs.
|
||||
"""
|
||||
frames = decode_video_from_file(path=video_path, frame_cap=frame_cap, device=device)
|
||||
result = None
|
||||
for f in frames:
|
||||
frame = resize_and_center_crop(f.to(torch.float32), height, width)
|
||||
@@ -257,9 +266,23 @@ def _audio_frame_to_float(frame: av.AudioFrame) -> np.ndarray:
|
||||
return arr
|
||||
|
||||
|
||||
def get_videostream_metadata(path: str) -> tuple[float, int, int, int]:
|
||||
"""Read video stream metadata: (fps, num_frames, width, height).
|
||||
def get_videostream_fps(path: str) -> float:
|
||||
"""Read video stream FPS."""
|
||||
container = av.open(path)
|
||||
try:
|
||||
video_stream = next(s for s in container.streams if s.type == "video")
|
||||
return float(video_stream.average_rate)
|
||||
finally:
|
||||
container.close()
|
||||
|
||||
|
||||
def get_videostream_metadata(path: str) -> VideoPixelShape:
|
||||
"""Read video stream metadata as a VideoPixelShape with batch=1.
|
||||
If frame count is missing in the container, decodes the stream to count frames.
|
||||
Args:
|
||||
path: Path to the video file.
|
||||
Returns:
|
||||
VideoPixelShape with batch=1, frames, height, width, and fps populated from the stream.
|
||||
"""
|
||||
container = av.open(path)
|
||||
try:
|
||||
@@ -270,7 +293,7 @@ def get_videostream_metadata(path: str) -> tuple[float, int, int, int]:
|
||||
num_frames = sum(1 for _ in container.decode(video_stream))
|
||||
width = video_stream.codec_context.width
|
||||
height = video_stream.codec_context.height
|
||||
return fps, num_frames, width, height
|
||||
return VideoPixelShape(batch=1, frames=num_frames, height=height, width=width, fps=fps)
|
||||
finally:
|
||||
container.close()
|
||||
|
||||
@@ -338,16 +361,85 @@ def decode_audio_from_file(
|
||||
return Audio(waveform=waveform, sampling_rate=sample_rate)
|
||||
|
||||
|
||||
def decode_video_from_file(path: str, frame_cap: int, device: DeviceLikeType) -> Generator[torch.Tensor]:
|
||||
def decode_video_by_frame(
|
||||
path: str,
|
||||
device: DeviceLikeType,
|
||||
starting_frame: int = 0,
|
||||
frame_cap: int | None = None,
|
||||
) -> Generator[torch.Tensor]:
|
||||
"""Decodes video from a file by sequential frame index, without relying on pts.
|
||||
Args:
|
||||
path: Path to the video file.
|
||||
device: Device to place the resulting tensors on.
|
||||
starting_frame: Number of leading frames to skip (default 0).
|
||||
frame_cap: Maximum number of frames to yield. If None, no frame limit (default None).
|
||||
Yields:
|
||||
Frames as tensors of shape (1, H, W, C), dtype uint8.
|
||||
"""
|
||||
container = av.open(path)
|
||||
try:
|
||||
video_stream = next(s for s in container.streams if s.type == "video")
|
||||
for frame in container.decode(video_stream):
|
||||
for index, frame in enumerate(container.decode(video_stream)):
|
||||
if index < starting_frame:
|
||||
continue
|
||||
tensor = torch.tensor(frame.to_rgb().to_ndarray(), dtype=torch.uint8, device=device).unsqueeze(0)
|
||||
yield tensor
|
||||
frame_cap = frame_cap - 1
|
||||
if frame_cap == 0:
|
||||
if frame_cap is not None:
|
||||
frame_cap -= 1
|
||||
if frame_cap == 0:
|
||||
break
|
||||
finally:
|
||||
container.close()
|
||||
|
||||
|
||||
def decode_video_from_file(
|
||||
path: str,
|
||||
device: DeviceLikeType,
|
||||
start_time: float = 0.0,
|
||||
max_duration: float | None = None,
|
||||
) -> Generator[torch.Tensor]:
|
||||
"""Decodes video from a file using presentation timestamps for time-based trimming.
|
||||
If a frame with no pts is encountered, falls back to :func:`decode_video_by_frame`
|
||||
using FPS-derived frame indices.
|
||||
Args:
|
||||
path: Path to the video file.
|
||||
device: Device to place the resulting tensors on.
|
||||
start_time: Start time in seconds (default 0.0).
|
||||
max_duration: Maximum duration in seconds to decode. If None, reads to end of
|
||||
stream (default None).
|
||||
Yields:
|
||||
Frames as tensors of shape (1, H, W, C), dtype uint8.
|
||||
"""
|
||||
container = av.open(path)
|
||||
try:
|
||||
video_stream = next(s for s in container.streams if s.type == "video")
|
||||
time_base = float(video_stream.time_base)
|
||||
|
||||
if start_time > 0:
|
||||
container.seek(int(start_time / time_base), stream=video_stream)
|
||||
|
||||
end_time = start_time + max_duration if max_duration is not None else None
|
||||
|
||||
for frame in container.decode(video_stream):
|
||||
# PyAV may leave pts unset when the demuxer does not expose per-frame
|
||||
# timestamps (e.g. some raw/elementary streams, stripped or missing
|
||||
# metadata, or certain remux paths). Without pts we cannot map frames to
|
||||
# wall-clock time, so we fall back to sequential frame indices using the
|
||||
# stream's average frame rate.
|
||||
if frame.pts is None:
|
||||
fps = float(video_stream.average_rate)
|
||||
starting_frame = round(start_time * fps)
|
||||
frame_cap = round(max_duration * fps) if max_duration is not None else None
|
||||
yield from decode_video_by_frame(
|
||||
path=path, device=device, starting_frame=starting_frame, frame_cap=frame_cap
|
||||
)
|
||||
return
|
||||
frame_time = frame.pts * time_base
|
||||
if frame_time < start_time:
|
||||
continue
|
||||
if end_time is not None and frame_time >= end_time:
|
||||
break
|
||||
yield torch.tensor(frame.to_rgb().to_ndarray(), dtype=torch.uint8, device=device).unsqueeze(0)
|
||||
finally:
|
||||
container.close()
|
||||
|
||||
|
||||
@@ -1,304 +0,0 @@
|
||||
from dataclasses import replace
|
||||
|
||||
import torch
|
||||
|
||||
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,
|
||||
AudioDecoder,
|
||||
AudioDecoderConfigurator,
|
||||
AudioEncoder,
|
||||
AudioEncoderConfigurator,
|
||||
Vocoder,
|
||||
VocoderConfigurator,
|
||||
)
|
||||
from ltx_core.model.transformer import (
|
||||
LTXV_MODEL_COMFY_RENAMING_MAP,
|
||||
LTXModelConfigurator,
|
||||
X0Model,
|
||||
)
|
||||
from ltx_core.model.upsampler import LatentUpsampler, LatentUpsamplerConfigurator
|
||||
from ltx_core.model.video_vae import (
|
||||
VAE_DECODER_COMFY_KEYS_FILTER,
|
||||
VAE_ENCODER_COMFY_KEYS_FILTER,
|
||||
VideoDecoder,
|
||||
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,
|
||||
EmbeddingsProcessor,
|
||||
EmbeddingsProcessorConfigurator,
|
||||
GemmaTextEncoder,
|
||||
GemmaTextEncoderConfigurator,
|
||||
module_ops_from_gemma_root,
|
||||
)
|
||||
from ltx_core.utils import find_matching_file
|
||||
|
||||
|
||||
class ModelLedger:
|
||||
"""
|
||||
Central coordinator for loading and building models used in an LTX pipeline.
|
||||
The ledger wires together multiple model builders (transformer, video VAE encoder/decoder,
|
||||
audio VAE decoder, vocoder, text encoder, and optional latent upsampler) and exposes
|
||||
factory methods for constructing model instances.
|
||||
### Model Building
|
||||
Each model method (e.g. :meth:`transformer`, :meth:`video_decoder`, :meth:`text_encoder`)
|
||||
constructs a new model instance on each call. The builder uses the
|
||||
:class:`~ltx_core.loader.registry.Registry` to load weights from the checkpoint,
|
||||
instantiates the model with the configured ``dtype``, and moves it to ``self.device``.
|
||||
.. note::
|
||||
Models are **not cached**. Each call to a model method creates a new instance.
|
||||
Callers are responsible for storing references to models they wish to reuse
|
||||
and for freeing GPU memory (e.g. by deleting references and calling
|
||||
``torch.cuda.empty_cache()``).
|
||||
### Constructor parameters
|
||||
dtype:
|
||||
Torch dtype used when constructing all models (e.g. ``torch.bfloat16``).
|
||||
device:
|
||||
Target device to which models are moved after construction (e.g. ``torch.device("cuda")``).
|
||||
checkpoint_path:
|
||||
Path to a checkpoint directory or file containing the core model weights
|
||||
(transformer, video VAE, audio VAE, text encoder, vocoder). If ``None``, the
|
||||
corresponding builders are not created and calling those methods will raise
|
||||
a :class:`ValueError`.
|
||||
gemma_root_path:
|
||||
Base path to Gemma-compatible CLIP/text encoder weights. Required to
|
||||
initialize the text encoder builder; if omitted, :meth:`text_encoder` cannot be used.
|
||||
spatial_upsampler_path:
|
||||
Optional path to a latent upsampler checkpoint. If provided, the
|
||||
:meth:`spatial_upsampler` method becomes available; otherwise calling it raises
|
||||
a :class:`ValueError`.
|
||||
loras:
|
||||
Tuple of LoRA configurations (path, strength, sd_ops) applied on top of the base
|
||||
transformer weights. Use ``()`` for none.
|
||||
registry:
|
||||
Optional :class:`Registry` instance for weight caching across builders.
|
||||
Defaults to :class:`DummyRegistry` which performs no cross-builder caching.
|
||||
quantization:
|
||||
Optional :class:`QuantizationPolicy` controlling how transformer weights
|
||||
are stored and how matmul is executed. Defaults to None, which means no quantization.
|
||||
### Creating Variants
|
||||
Use :meth:`with_additional_loras` to create a new ``ModelLedger`` instance that
|
||||
includes additional LoRA configurations or :meth:`with_loras` to replace existing
|
||||
lora configurations while sharing the same registry for weight caching.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
checkpoint_path: str | None = None,
|
||||
gemma_root_path: str | None = None,
|
||||
spatial_upsampler_path: str | None = None,
|
||||
loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
|
||||
registry: Registry | None = None,
|
||||
quantization: QuantizationPolicy | None = None,
|
||||
):
|
||||
self.dtype = dtype
|
||||
self.device = device
|
||||
self.checkpoint_path = checkpoint_path
|
||||
self.gemma_root_path = gemma_root_path
|
||||
self.spatial_upsampler_path = spatial_upsampler_path
|
||||
self.loras = loras
|
||||
self.registry = registry or DummyRegistry()
|
||||
self.quantization = quantization
|
||||
self.build_model_builders()
|
||||
|
||||
def build_model_builders(self) -> None:
|
||||
if self.checkpoint_path is not None:
|
||||
self.transformer_builder = Builder(
|
||||
model_path=self.checkpoint_path,
|
||||
model_class_configurator=LTXModelConfigurator,
|
||||
model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
|
||||
loras=tuple(self.loras),
|
||||
registry=self.registry,
|
||||
)
|
||||
|
||||
self.vae_decoder_builder = Builder(
|
||||
model_path=self.checkpoint_path,
|
||||
model_class_configurator=VideoDecoderConfigurator,
|
||||
model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
|
||||
registry=self.registry,
|
||||
)
|
||||
|
||||
self.vae_encoder_builder = Builder(
|
||||
model_path=self.checkpoint_path,
|
||||
model_class_configurator=VideoEncoderConfigurator,
|
||||
model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
|
||||
registry=self.registry,
|
||||
)
|
||||
|
||||
self.audio_encoder_builder = Builder[AudioEncoder](
|
||||
model_path=self.checkpoint_path,
|
||||
model_class_configurator=AudioEncoderConfigurator,
|
||||
model_sd_ops=AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
|
||||
registry=self.registry,
|
||||
)
|
||||
|
||||
self.audio_decoder_builder = Builder(
|
||||
model_path=self.checkpoint_path,
|
||||
model_class_configurator=AudioDecoderConfigurator,
|
||||
model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
|
||||
registry=self.registry,
|
||||
)
|
||||
|
||||
self.vocoder_builder = Builder(
|
||||
model_path=self.checkpoint_path,
|
||||
model_class_configurator=VocoderConfigurator,
|
||||
model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
|
||||
registry=self.registry,
|
||||
)
|
||||
|
||||
# Embeddings processor only needs the LTX checkpoint (no Gemma weights)
|
||||
self.embeddings_processor_builder = Builder(
|
||||
model_path=self.checkpoint_path,
|
||||
model_class_configurator=EmbeddingsProcessorConfigurator,
|
||||
model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS,
|
||||
registry=self.registry,
|
||||
)
|
||||
|
||||
if self.gemma_root_path is not None:
|
||||
module_ops = module_ops_from_gemma_root(self.gemma_root_path)
|
||||
model_folder = find_matching_file(self.gemma_root_path, "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,
|
||||
registry=self.registry,
|
||||
module_ops=(GEMMA_MODEL_OPS, *module_ops),
|
||||
)
|
||||
|
||||
if self.spatial_upsampler_path is not None:
|
||||
self.upsampler_builder = Builder(
|
||||
model_path=self.spatial_upsampler_path,
|
||||
model_class_configurator=LatentUpsamplerConfigurator,
|
||||
registry=self.registry,
|
||||
)
|
||||
|
||||
def _target_device(self) -> torch.device:
|
||||
if isinstance(self.registry, DummyRegistry) or self.registry is None:
|
||||
return self.device
|
||||
else:
|
||||
return torch.device("cpu")
|
||||
|
||||
def with_additional_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "ModelLedger":
|
||||
"""Add new lora configurations to the existing ones."""
|
||||
return self.with_loras((*self.loras, *loras))
|
||||
|
||||
def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "ModelLedger":
|
||||
"""Replace existing lora configurations with new ones."""
|
||||
return ModelLedger(
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
checkpoint_path=self.checkpoint_path,
|
||||
gemma_root_path=self.gemma_root_path,
|
||||
spatial_upsampler_path=self.spatial_upsampler_path,
|
||||
loras=loras,
|
||||
registry=self.registry,
|
||||
quantization=self.quantization,
|
||||
)
|
||||
|
||||
def transformer(self) -> X0Model:
|
||||
if not hasattr(self, "transformer_builder"):
|
||||
raise ValueError(
|
||||
"Transformer not initialized. Please provide a checkpoint path to the ModelLedger constructor."
|
||||
)
|
||||
|
||||
if self.quantization is None:
|
||||
return (
|
||||
X0Model(self.transformer_builder.build(device=self._target_device(), dtype=self.dtype))
|
||||
.to(self.device)
|
||||
.eval()
|
||||
)
|
||||
else:
|
||||
sd_ops = self.transformer_builder.model_sd_ops
|
||||
if self.quantization.sd_ops is not None:
|
||||
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 = replace(
|
||||
self.transformer_builder,
|
||||
module_ops=(*self.transformer_builder.module_ops, *self.quantization.module_ops),
|
||||
model_sd_ops=sd_ops,
|
||||
)
|
||||
return X0Model(builder.build(device=self._target_device())).to(self.device).eval()
|
||||
|
||||
def video_decoder(self) -> VideoDecoder:
|
||||
if not hasattr(self, "vae_decoder_builder"):
|
||||
raise ValueError(
|
||||
"Video decoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
|
||||
)
|
||||
|
||||
return self.vae_decoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
|
||||
|
||||
def video_encoder(self) -> VideoEncoder:
|
||||
if not hasattr(self, "vae_encoder_builder"):
|
||||
raise ValueError(
|
||||
"Video encoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
|
||||
)
|
||||
|
||||
return self.vae_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
|
||||
|
||||
def text_encoder(self) -> GemmaTextEncoder:
|
||||
if not hasattr(self, "text_encoder_builder"):
|
||||
raise ValueError(
|
||||
"Text encoder not initialized. Please provide a checkpoint path and gemma root path to the "
|
||||
"ModelLedger constructor."
|
||||
)
|
||||
|
||||
return self.text_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
|
||||
|
||||
def gemma_embeddings_processor(self) -> EmbeddingsProcessor:
|
||||
if not hasattr(self, "embeddings_processor_builder"):
|
||||
raise ValueError(
|
||||
"Embeddings processor not initialized. Please provide a checkpoint path to the ModelLedger constructor."
|
||||
)
|
||||
|
||||
return (
|
||||
self.embeddings_processor_builder.build(device=self._target_device(), dtype=self.dtype)
|
||||
.to(self.device)
|
||||
.eval()
|
||||
)
|
||||
|
||||
def audio_encoder(self) -> AudioEncoder:
|
||||
if not hasattr(self, "audio_encoder_builder"):
|
||||
raise ValueError(
|
||||
"Audio encoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
|
||||
)
|
||||
|
||||
return self.audio_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
|
||||
|
||||
def audio_decoder(self) -> AudioDecoder:
|
||||
if not hasattr(self, "audio_decoder_builder"):
|
||||
raise ValueError(
|
||||
"Audio decoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
|
||||
)
|
||||
|
||||
return self.audio_decoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
|
||||
|
||||
def vocoder(self) -> Vocoder:
|
||||
if not hasattr(self, "vocoder_builder"):
|
||||
raise ValueError(
|
||||
"Vocoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
|
||||
)
|
||||
|
||||
return self.vocoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
|
||||
|
||||
def spatial_upsampler(self) -> LatentUpsampler:
|
||||
if not hasattr(self, "upsampler_builder"):
|
||||
raise ValueError("Upsampler not initialized. Please provide upsampler path to the ModelLedger constructor.")
|
||||
|
||||
return self.upsampler_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
|
||||
@@ -8,86 +8,91 @@ from tqdm import tqdm
|
||||
|
||||
from ltx_core.components.diffusion_steps import Res2sDiffusionStep
|
||||
from ltx_core.components.protocols import DiffusionStepProtocol
|
||||
from ltx_core.model.transformer import X0Model
|
||||
from ltx_core.utils import to_denoised, to_velocity
|
||||
from ltx_pipelines.utils.helpers import post_process_latent, timesteps_from_mask
|
||||
from ltx_pipelines.utils.res2s import get_res2s_coefficients
|
||||
from ltx_pipelines.utils.types import DenoisingFunc, LatentState
|
||||
from ltx_pipelines.utils.types import Denoiser, LatentState
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _step_state(
|
||||
state: LatentState | None,
|
||||
denoised: torch.Tensor | None,
|
||||
stepper: DiffusionStepProtocol,
|
||||
sigmas: torch.Tensor,
|
||||
step_idx: int,
|
||||
) -> LatentState | None:
|
||||
"""Advance one diffusion step for a single modality, or return ``None`` if absent."""
|
||||
if state is None or denoised is None:
|
||||
return state
|
||||
denoised = post_process_latent(denoised, state.denoise_mask, state.clean_latent)
|
||||
return replace(state, latent=stepper.step(state.latent, denoised, sigmas, step_idx))
|
||||
|
||||
|
||||
def euler_denoising_loop(
|
||||
sigmas: torch.Tensor,
|
||||
video_state: LatentState,
|
||||
audio_state: LatentState,
|
||||
video_state: LatentState | None,
|
||||
audio_state: LatentState | None,
|
||||
stepper: DiffusionStepProtocol,
|
||||
denoise_fn: DenoisingFunc,
|
||||
) -> tuple[LatentState, LatentState]:
|
||||
transformer: X0Model,
|
||||
denoiser: Denoiser,
|
||||
) -> tuple[LatentState | None, LatentState | None]:
|
||||
"""
|
||||
Perform the joint audio-video denoising loop over a diffusion schedule.
|
||||
This function iterates over all but the final value in ``sigmas`` and, at
|
||||
each diffusion step, calls ``denoise_fn`` to obtain denoised video and
|
||||
audio latents. The denoised latents are post-processed with their
|
||||
respective denoise masks and clean latents, then passed to ``stepper`` to
|
||||
advance the noisy latents one step along the diffusion schedule.
|
||||
Either ``video_state`` or ``audio_state`` may be ``None`` for absent
|
||||
modalities; the absent modality is passed through unchanged.
|
||||
### Parameters
|
||||
sigmas:
|
||||
A 1D tensor of noise levels (diffusion sigmas) defining the sampling
|
||||
schedule. All steps except the last element are iterated over.
|
||||
video_state:
|
||||
The current video :class:`LatentState`, containing the noisy latent,
|
||||
its clean reference latent, and the denoising mask.
|
||||
The current video :class:`LatentState`, or ``None`` if video is absent.
|
||||
audio_state:
|
||||
The current audio :class:`LatentState`, analogous to ``video_state``
|
||||
but for the audio modality.
|
||||
The current audio :class:`LatentState`, or ``None`` if audio is absent.
|
||||
stepper:
|
||||
An implementation of :class:`DiffusionStepProtocol` that updates a
|
||||
latent given the current latent, its denoised estimate, the full
|
||||
``sigmas`` schedule, and the current step index.
|
||||
denoise_fn:
|
||||
A callable implementing :class:`DenoisingFunc`. It is invoked as
|
||||
``denoise_fn(video_state, audio_state, sigmas, step_index)`` and must
|
||||
return a tuple ``(denoised_video, denoised_audio)``, where each element
|
||||
is a tensor with the same shape as the corresponding latent.
|
||||
transformer:
|
||||
The diffusion model passed to the denoiser at each step.
|
||||
denoiser:
|
||||
A callable implementing :class:`Denoiser`. It is invoked as
|
||||
``denoiser(transformer, video_state, audio_state, sigmas, step_index)``
|
||||
and must return ``(denoised_video, denoised_audio)``.
|
||||
### Returns
|
||||
tuple[LatentState, LatentState]
|
||||
A pair ``(video_state, audio_state)`` containing the final video and
|
||||
audio latent states after completing the denoising loop.
|
||||
tuple[LatentState | None, LatentState | None]
|
||||
Final ``(video_state, audio_state)`` after the denoising loop.
|
||||
"""
|
||||
for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
|
||||
denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_idx)
|
||||
denoised_video, denoised_audio = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
|
||||
|
||||
denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
|
||||
denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
|
||||
|
||||
video_state = replace(video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx))
|
||||
audio_state = replace(audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx))
|
||||
video_state = _step_state(video_state, denoised_video, stepper, sigmas, step_idx)
|
||||
audio_state = _step_state(audio_state, denoised_audio, stepper, sigmas, step_idx)
|
||||
|
||||
return (video_state, audio_state)
|
||||
|
||||
|
||||
def gradient_estimating_euler_denoising_loop(
|
||||
sigmas: torch.Tensor,
|
||||
video_state: LatentState,
|
||||
audio_state: LatentState,
|
||||
video_state: LatentState | None,
|
||||
audio_state: LatentState | None,
|
||||
stepper: DiffusionStepProtocol,
|
||||
denoise_fn: DenoisingFunc,
|
||||
transformer: X0Model,
|
||||
denoiser: Denoiser,
|
||||
ge_gamma: float = 2.0,
|
||||
) -> tuple[LatentState, LatentState]:
|
||||
) -> tuple[LatentState | None, LatentState | None]:
|
||||
"""
|
||||
Perform the joint audio-video denoising loop using gradient-estimation sampling.
|
||||
This function is similar to :func:`euler_denoising_loop`, but applies
|
||||
gradient estimation to improve the denoised estimates by tracking velocity
|
||||
changes across steps. See the referenced function for detailed parameter
|
||||
documentation.
|
||||
Same interface as :func:`euler_denoising_loop` with an additional
|
||||
``ge_gamma`` parameter for velocity correction.
|
||||
### Parameters
|
||||
ge_gamma:
|
||||
Gradient estimation coefficient controlling the velocity correction term.
|
||||
Default is 2.0. Paper: https://openreview.net/pdf?id=o2ND9v0CeK
|
||||
sigmas, video_state, audio_state, stepper, denoise_fn:
|
||||
See :func:`euler_denoising_loop` for parameter descriptions.
|
||||
### Returns
|
||||
tuple[LatentState, LatentState]
|
||||
tuple[LatentState | None, LatentState | None]
|
||||
See :func:`euler_denoising_loop` for return value description.
|
||||
"""
|
||||
|
||||
@@ -105,23 +110,35 @@ def gradient_estimating_euler_denoising_loop(
|
||||
return current_velocity, denoised_sample
|
||||
|
||||
for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
|
||||
denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_idx)
|
||||
denoised_video, denoised_audio = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
|
||||
|
||||
denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
|
||||
denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
|
||||
if video_state is not None and denoised_video is not None:
|
||||
denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
|
||||
if audio_state is not None and denoised_audio is not None:
|
||||
denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
|
||||
|
||||
if sigmas[step_idx + 1] == 0:
|
||||
return replace(video_state, latent=denoised_video), replace(audio_state, latent=denoised_audio)
|
||||
if video_state is not None and denoised_video is not None:
|
||||
video_state = replace(video_state, latent=denoised_video)
|
||||
if audio_state is not None and denoised_audio is not None:
|
||||
audio_state = replace(audio_state, latent=denoised_audio)
|
||||
return video_state, audio_state
|
||||
|
||||
previous_video_velocity, denoised_video = update_velocity_and_sample(
|
||||
video_state.latent, denoised_video, sigmas[step_idx], previous_video_velocity
|
||||
)
|
||||
previous_audio_velocity, denoised_audio = update_velocity_and_sample(
|
||||
audio_state.latent, denoised_audio, sigmas[step_idx], previous_audio_velocity
|
||||
)
|
||||
if video_state is not None and denoised_video is not None:
|
||||
previous_video_velocity, denoised_video = update_velocity_and_sample(
|
||||
video_state.latent, denoised_video, sigmas[step_idx], previous_video_velocity
|
||||
)
|
||||
video_state = replace(
|
||||
video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx)
|
||||
)
|
||||
|
||||
video_state = replace(video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx))
|
||||
audio_state = replace(audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx))
|
||||
if audio_state is not None and denoised_audio is not None:
|
||||
previous_audio_velocity, denoised_audio = update_velocity_and_sample(
|
||||
audio_state.latent, denoised_audio, sigmas[step_idx], previous_audio_velocity
|
||||
)
|
||||
audio_state = replace(
|
||||
audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx)
|
||||
)
|
||||
|
||||
return (video_state, audio_state)
|
||||
|
||||
@@ -146,6 +163,7 @@ def _inject_sde_noise(
|
||||
sigmas: torch.Tensor,
|
||||
step_idx: int,
|
||||
legacy_mode: bool = False,
|
||||
eta: float = 0.5,
|
||||
) -> torch.Tensor:
|
||||
sigmas_copy = sigmas.clone()
|
||||
new_noise = new_noise_fn(state.latent, step_noise_generator)
|
||||
@@ -160,6 +178,7 @@ def _inject_sde_noise(
|
||||
sigmas=sigmas,
|
||||
step_index=step_idx,
|
||||
noise=new_noise,
|
||||
eta=eta,
|
||||
)
|
||||
|
||||
if legacy_mode:
|
||||
@@ -168,20 +187,22 @@ def _inject_sde_noise(
|
||||
return x_next
|
||||
|
||||
|
||||
def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915
|
||||
def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
|
||||
sigmas: torch.Tensor,
|
||||
video_state: LatentState,
|
||||
audio_state: LatentState,
|
||||
video_state: LatentState | None,
|
||||
audio_state: LatentState | None,
|
||||
stepper: DiffusionStepProtocol,
|
||||
denoise_fn: DenoisingFunc,
|
||||
transformer: X0Model,
|
||||
denoiser: Denoiser,
|
||||
noise_seed: int = -1,
|
||||
noise_seed_substep: int | None = None,
|
||||
eta: float = 0.5,
|
||||
bongmath: bool = True,
|
||||
bongmath_max_iter: int = 100,
|
||||
new_noise_fn: Callable[[torch.Tensor, torch.Generator], torch.Tensor] = _get_new_noise,
|
||||
model_dtype: torch.dtype = torch.bfloat16,
|
||||
legacy_mode: bool = True,
|
||||
) -> tuple[LatentState, LatentState]:
|
||||
) -> tuple[LatentState | None, LatentState | None]:
|
||||
"""
|
||||
Joint audio-video denoising loop using the res_2s second-order sampler.
|
||||
Iterates over the diffusion schedule with a two-stage Runge-Kutta step:
|
||||
@@ -189,46 +210,48 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915
|
||||
noise), then combines both with RK coefficients. Supports anchor-point
|
||||
refinement (bong iteration) and optional SDE noise injection. Requires
|
||||
:class:`Res2sDiffusionStep` as ``stepper``.
|
||||
Either modality may be ``None`` (absent).
|
||||
### Parameters
|
||||
sigmas:
|
||||
A 1D tensor of noise levels defining the sampling schedule.
|
||||
video_state:
|
||||
Current video :class:`LatentState` (noisy latent, clean reference, mask).
|
||||
audio_state:
|
||||
Current audio :class:`LatentState`, same structure as ``video_state``.
|
||||
stepper:
|
||||
Must be an instance of :class:`Res2sDiffusionStep`; performs SDE step
|
||||
with noise injection.
|
||||
denoise_fn:
|
||||
Callable ``(video_state, audio_state, sigmas, step_index)`` returning
|
||||
``(denoised_video, denoised_audio)``.
|
||||
transformer:
|
||||
The diffusion model passed to the denoiser at each step.
|
||||
denoiser:
|
||||
Callable implementing :class:`Denoiser`.
|
||||
noise_seed:
|
||||
Seed for step-level SDE noise; substep seed defaults to ``noise_seed + 10000``.
|
||||
noise_seed_substep:
|
||||
Optional seed for substep SDE noise; if None, derived from ``noise_seed``.
|
||||
eta:
|
||||
Controls stochastic noise injection strength (0=deterministic, 1=maximum).
|
||||
Applies to main diffusion steps; substeps always use 0.5. Default 0.5.
|
||||
bongmath:
|
||||
Whether to run iterative anchor refinement (bong iteration) when step size is small.
|
||||
bongmath_max_iter:
|
||||
Max iterations for bong refinement when enabled.
|
||||
new_noise_fn:
|
||||
Callable ``(latent, generator) -> noise`` for SDE injection; default
|
||||
uses normalized channel-wise Gaussian noise.
|
||||
Callable ``(latent, generator) -> noise`` for SDE injection.
|
||||
model_dtype:
|
||||
Dtype for latent state updates (e.g. bfloat16).
|
||||
### Returns
|
||||
tuple[LatentState, LatentState]
|
||||
tuple[LatentState | None, LatentState | None]
|
||||
Final ``(video_state, audio_state)`` after the denoising loop.
|
||||
"""
|
||||
# Determine device from whichever state is present
|
||||
present_state = video_state or audio_state
|
||||
if present_state is None:
|
||||
raise ValueError("At least one of video_state or audio_state must be provided")
|
||||
state_device = present_state.latent.device
|
||||
|
||||
# Initialize noise generators with different seeds
|
||||
if noise_seed_substep is None:
|
||||
noise_seed_substep = noise_seed + 10000 # Offset to ensure different seeds
|
||||
step_noise_generator = torch.Generator(device=video_state.latent.device).manual_seed(noise_seed)
|
||||
substep_noise_generator = torch.Generator(device=video_state.latent.device).manual_seed(noise_seed_substep)
|
||||
step_noise_generator = torch.Generator(device=state_device).manual_seed(noise_seed)
|
||||
substep_noise_generator = torch.Generator(device=state_device).manual_seed(noise_seed_substep)
|
||||
sde_noise_injecting_fn = partial(
|
||||
_inject_sde_noise, stepper=stepper, new_noise_fn=new_noise_fn, legacy_mode=legacy_mode
|
||||
)
|
||||
step_noise_injecting_fn = partial(sde_noise_injecting_fn, step_noise_generator=step_noise_generator)
|
||||
substep_noise_injecting_fn = partial(sde_noise_injecting_fn, step_noise_generator=substep_noise_generator)
|
||||
step_noise_injecting_fn = partial(sde_noise_injecting_fn, step_noise_generator=step_noise_generator, eta=eta)
|
||||
# substep eta is always default 0.5 for compatibility with original implementation.
|
||||
substep_noise_injecting_fn = partial(sde_noise_injecting_fn, step_noise_generator=substep_noise_generator, eta=0.5)
|
||||
|
||||
if not isinstance(stepper, Res2sDiffusionStep):
|
||||
raise ValueError("stepper must be an instance of Res2sDiffusionStep")
|
||||
@@ -241,26 +264,25 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915
|
||||
hs = -torch.log(sigmas[1:].double().cpu() / (sigmas[:-1].double().cpu()))
|
||||
|
||||
# Initialize phi cache for reuse across loop iterations
|
||||
# Cache key: (j, neg_h) where j is phi order and neg_h is negative step value
|
||||
phi_cache = {}
|
||||
c2 = 0.5 # Midpoint for res_2s
|
||||
|
||||
# Progress bar shows only full two-stage steps; final (sigma_next==0) step is done silently
|
||||
|
||||
for step_idx in tqdm(range(n_full_steps)):
|
||||
sigma = sigmas[step_idx].double()
|
||||
sigma_next = sigmas[step_idx + 1].double()
|
||||
|
||||
# Initialize anchor point
|
||||
x_anchor_video = video_state.latent.clone().double()
|
||||
x_anchor_audio = audio_state.latent.clone().double()
|
||||
x_anchor_video = video_state.latent.clone().double() if video_state is not None else None
|
||||
x_anchor_audio = audio_state.latent.clone().double() if audio_state is not None else None
|
||||
|
||||
# ====================================================================
|
||||
# STAGE 1: Evaluate at current point
|
||||
# ====================================================================
|
||||
denoised_video_1, denoised_audio_1 = denoise_fn(video_state, audio_state, sigmas, step_idx)
|
||||
denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
|
||||
denoised_audio_1 = post_process_latent(denoised_audio_1, audio_state.denoise_mask, audio_state.clean_latent)
|
||||
denoised_video_1, denoised_audio_1 = denoiser(transformer, video_state, audio_state, sigmas, step_idx)
|
||||
if video_state is not None and denoised_video_1 is not None:
|
||||
denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
|
||||
if audio_state is not None and denoised_audio_1 is not None:
|
||||
denoised_audio_1 = post_process_latent(denoised_audio_1, audio_state.denoise_mask, audio_state.clean_latent)
|
||||
|
||||
h = hs[step_idx].item()
|
||||
|
||||
@@ -273,91 +295,127 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915
|
||||
# ====================================================================
|
||||
# Compute substep x using RK coefficient a21
|
||||
# ====================================================================
|
||||
eps_1_video = denoised_video_1.double() - x_anchor_video
|
||||
eps_1_audio = denoised_audio_1.double() - x_anchor_audio
|
||||
if x_anchor_video is not None and denoised_video_1 is not None:
|
||||
eps_1_video = denoised_video_1.double() - x_anchor_video
|
||||
x_mid_video = x_anchor_video.double() + h * a21 * eps_1_video
|
||||
else:
|
||||
eps_1_video = None
|
||||
x_mid_video = None
|
||||
|
||||
x_mid_video = x_anchor_video.double() + h * a21 * eps_1_video
|
||||
x_mid_audio = x_anchor_audio.double() + h * a21 * eps_1_audio
|
||||
if x_anchor_audio is not None and denoised_audio_1 is not None:
|
||||
eps_1_audio = denoised_audio_1.double() - x_anchor_audio
|
||||
x_mid_audio = x_anchor_audio.double() + h * a21 * eps_1_audio
|
||||
else:
|
||||
eps_1_audio = None
|
||||
x_mid_audio = None
|
||||
|
||||
# ====================================================================
|
||||
# SDE noise injection at substep
|
||||
# ====================================================================
|
||||
x_mid_video = substep_noise_injecting_fn(
|
||||
state=video_state,
|
||||
sample=x_anchor_video,
|
||||
denoised_sample=x_mid_video,
|
||||
sigmas=torch.stack([sigma, sub_sigma]),
|
||||
step_idx=0,
|
||||
)
|
||||
x_mid_audio = substep_noise_injecting_fn(
|
||||
state=audio_state,
|
||||
sample=x_anchor_audio,
|
||||
denoised_sample=x_mid_audio,
|
||||
sigmas=torch.stack([sigma, sub_sigma]),
|
||||
step_idx=0,
|
||||
)
|
||||
if x_mid_video is not None and video_state is not None:
|
||||
x_mid_video = substep_noise_injecting_fn(
|
||||
state=video_state,
|
||||
sample=x_anchor_video,
|
||||
denoised_sample=x_mid_video,
|
||||
sigmas=torch.stack([sigma, sub_sigma]),
|
||||
step_idx=0,
|
||||
)
|
||||
if x_mid_audio is not None and audio_state is not None:
|
||||
x_mid_audio = substep_noise_injecting_fn(
|
||||
state=audio_state,
|
||||
sample=x_anchor_audio,
|
||||
denoised_sample=x_mid_audio,
|
||||
sigmas=torch.stack([sigma, sub_sigma]),
|
||||
step_idx=0,
|
||||
)
|
||||
|
||||
# ====================================================================
|
||||
# ITERATIVE REFINEMENT (Bong Iteration) - Stabilize anchor point
|
||||
# ITERATIVE REFINEMENT (Bong Iteration)
|
||||
# ====================================================================
|
||||
if bongmath and h < 0.5 and sigma > 0.03:
|
||||
for _ in range(bongmath_max_iter):
|
||||
x_anchor_video = x_mid_video - h * a21 * eps_1_video
|
||||
eps_1_video = denoised_video_1.double() - x_anchor_video
|
||||
x_anchor_audio = x_mid_audio - h * a21 * eps_1_audio
|
||||
eps_1_audio = denoised_audio_1.double() - x_anchor_audio
|
||||
if x_mid_video is not None and eps_1_video is not None:
|
||||
x_anchor_video = x_mid_video - h * a21 * eps_1_video
|
||||
eps_1_video = denoised_video_1.double() - x_anchor_video
|
||||
if x_mid_audio is not None and eps_1_audio is not None:
|
||||
x_anchor_audio = x_mid_audio - h * a21 * eps_1_audio
|
||||
eps_1_audio = denoised_audio_1.double() - x_anchor_audio
|
||||
|
||||
# ====================================================================
|
||||
# STAGE 2: Evaluate at substep point (WITH NOISE)
|
||||
# ====================================================================
|
||||
mid_video_state = replace(video_state, latent=x_mid_video.to(model_dtype))
|
||||
mid_audio_state = replace(audio_state, latent=x_mid_audio.to(model_dtype))
|
||||
mid_video_state = (
|
||||
replace(video_state, latent=x_mid_video.to(model_dtype))
|
||||
if video_state is not None and x_mid_video is not None
|
||||
else None
|
||||
)
|
||||
mid_audio_state = (
|
||||
replace(audio_state, latent=x_mid_audio.to(model_dtype))
|
||||
if audio_state is not None and x_mid_audio is not None
|
||||
else None
|
||||
)
|
||||
|
||||
denoised_video_2, denoised_audio_2 = denoise_fn(
|
||||
denoised_video_2, denoised_audio_2 = denoiser(
|
||||
transformer,
|
||||
video_state=mid_video_state,
|
||||
audio_state=mid_audio_state,
|
||||
sigmas=torch.stack([sub_sigma]).to(sigmas.device),
|
||||
step_index=0,
|
||||
)
|
||||
denoised_video_2 = post_process_latent(denoised_video_2, video_state.denoise_mask, video_state.clean_latent)
|
||||
denoised_audio_2 = post_process_latent(denoised_audio_2, audio_state.denoise_mask, audio_state.clean_latent)
|
||||
if video_state is not None and denoised_video_2 is not None:
|
||||
denoised_video_2 = post_process_latent(denoised_video_2, video_state.denoise_mask, video_state.clean_latent)
|
||||
if audio_state is not None and denoised_audio_2 is not None:
|
||||
denoised_audio_2 = post_process_latent(denoised_audio_2, audio_state.denoise_mask, audio_state.clean_latent)
|
||||
|
||||
# ====================================================================
|
||||
# FINAL COMBINATION: Compute x_next using RK coefficients
|
||||
# ====================================================================
|
||||
eps_2_video = denoised_video_2.double() - x_anchor_video
|
||||
eps_2_audio = denoised_audio_2.double() - x_anchor_audio
|
||||
if x_anchor_video is not None and eps_1_video is not None and denoised_video_2 is not None:
|
||||
eps_2_video = denoised_video_2.double() - x_anchor_video
|
||||
x_next_video = x_anchor_video + h * (b1 * eps_1_video + b2 * eps_2_video)
|
||||
else:
|
||||
x_next_video = None
|
||||
|
||||
x_next_video = x_anchor_video + h * (b1 * eps_1_video + b2 * eps_2_video)
|
||||
x_next_audio = x_anchor_audio + h * (b1 * eps_1_audio + b2 * eps_2_audio)
|
||||
if x_anchor_audio is not None and eps_1_audio is not None and denoised_audio_2 is not None:
|
||||
eps_2_audio = denoised_audio_2.double() - x_anchor_audio
|
||||
x_next_audio = x_anchor_audio + h * (b1 * eps_1_audio + b2 * eps_2_audio)
|
||||
else:
|
||||
x_next_audio = None
|
||||
|
||||
# ====================================================================
|
||||
# SDE NOISE INJECTION AT STEP LEVEL
|
||||
# ====================================================================
|
||||
x_next_video = step_noise_injecting_fn(
|
||||
state=video_state,
|
||||
sample=x_anchor_video,
|
||||
denoised_sample=x_next_video,
|
||||
sigmas=sigmas,
|
||||
step_idx=step_idx,
|
||||
)
|
||||
x_next_audio = step_noise_injecting_fn(
|
||||
state=audio_state,
|
||||
sample=x_anchor_audio,
|
||||
denoised_sample=x_next_audio,
|
||||
sigmas=sigmas,
|
||||
step_idx=step_idx,
|
||||
)
|
||||
if x_next_video is not None and video_state is not None:
|
||||
x_next_video = step_noise_injecting_fn(
|
||||
state=video_state,
|
||||
sample=x_anchor_video,
|
||||
denoised_sample=x_next_video,
|
||||
sigmas=sigmas,
|
||||
step_idx=step_idx,
|
||||
)
|
||||
if x_next_audio is not None and audio_state is not None:
|
||||
x_next_audio = step_noise_injecting_fn(
|
||||
state=audio_state,
|
||||
sample=x_anchor_audio,
|
||||
denoised_sample=x_next_audio,
|
||||
sigmas=sigmas,
|
||||
step_idx=step_idx,
|
||||
)
|
||||
|
||||
# Update states
|
||||
video_state = replace(video_state, latent=x_next_video.to(model_dtype))
|
||||
audio_state = replace(audio_state, latent=x_next_audio.to(model_dtype))
|
||||
if video_state is not None and x_next_video is not None:
|
||||
video_state = replace(video_state, latent=x_next_video.to(model_dtype))
|
||||
if audio_state is not None and x_next_audio is not None:
|
||||
audio_state = replace(audio_state, latent=x_next_audio.to(model_dtype))
|
||||
|
||||
# Final step if we need to fully remove the noise
|
||||
if sigmas[-1] == 0:
|
||||
denoised_video_1, denoised_audio_1 = denoise_fn(video_state, audio_state, sigmas, n_full_steps)
|
||||
denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
|
||||
denoised_audio_1 = post_process_latent(denoised_audio_1, audio_state.denoise_mask, audio_state.clean_latent)
|
||||
video_state = replace(video_state, latent=denoised_video_1.to(model_dtype))
|
||||
audio_state = replace(audio_state, latent=denoised_audio_1.to(model_dtype))
|
||||
denoised_video_1, denoised_audio_1 = denoiser(transformer, video_state, audio_state, sigmas, n_full_steps)
|
||||
if video_state is not None and denoised_video_1 is not None:
|
||||
denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
|
||||
video_state = replace(video_state, latent=denoised_video_1.to(model_dtype))
|
||||
if audio_state is not None and denoised_audio_1 is not None:
|
||||
denoised_audio_1 = post_process_latent(denoised_audio_1, audio_state.denoise_mask, audio_state.clean_latent)
|
||||
audio_state = replace(audio_state, latent=denoised_audio_1.to(model_dtype))
|
||||
|
||||
return video_state, audio_state
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Protocol
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
|
||||
from ltx_core.components.protocols import DiffusionStepProtocol
|
||||
from ltx_core.conditioning import ConditioningItem
|
||||
from ltx_core.model.transformer import X0Model
|
||||
from ltx_core.types import LatentState
|
||||
from ltx_pipelines.utils.constants import VIDEO_LATENT_CHANNELS, VIDEO_SCALE_FACTORS
|
||||
|
||||
@@ -35,39 +37,40 @@ class PipelineComponents:
|
||||
self.audio_patchifier = AudioPatchifier(patch_size=1)
|
||||
|
||||
|
||||
class DenoisingFunc(Protocol):
|
||||
"""
|
||||
Protocol for a denoising function used in the LTX pipeline.
|
||||
class Denoiser(Protocol):
|
||||
"""Protocol for a denoiser that receives the transformer at call time.
|
||||
The transformer is not stored — it is passed as the first argument so the
|
||||
caller (a denoising loop or a pipeline block) controls its lifecycle.
|
||||
Args:
|
||||
video_state (LatentState): The current latent state for video.
|
||||
audio_state (LatentState): The current latent state for audio.
|
||||
sigmas (torch.Tensor): A 1D tensor of sigma values for each diffusion step.
|
||||
step_index (int): Index of the current denoising step.
|
||||
transformer: The diffusion model.
|
||||
video_state: Current video latent state, or ``None`` if absent.
|
||||
audio_state: Current audio latent state, or ``None`` if absent.
|
||||
sigmas: 1-D tensor of sigma values for each diffusion step.
|
||||
step_index: Index of the current denoising step.
|
||||
Returns:
|
||||
tuple[torch.Tensor, torch.Tensor]: The denoised video and audio tensors.
|
||||
"""
|
||||
|
||||
def __call__(
|
||||
self, video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
|
||||
) -> tuple[torch.Tensor, torch.Tensor]: ...
|
||||
|
||||
|
||||
class DenoisingLoopFunc(Protocol):
|
||||
"""
|
||||
Protocol for a denoising loop function used in the LTX pipeline.
|
||||
Args:
|
||||
sigmas (torch.Tensor): A 1D tensor of sigma values for each diffusion step.
|
||||
video_state (LatentState): The current latent state for video.
|
||||
audio_state (LatentState): The current latent state for audio.
|
||||
stepper (DiffusionStepProtocol): The diffusion step protocol to use.
|
||||
Returns:
|
||||
tuple[LatentState, LatentState]: The denoised video and audio latent states.
|
||||
``(denoised_video, denoised_audio)`` tensors (either may be ``None``).
|
||||
"""
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
transformer: X0Model,
|
||||
video_state: LatentState | None,
|
||||
audio_state: LatentState | None,
|
||||
sigmas: torch.Tensor,
|
||||
video_state: LatentState,
|
||||
audio_state: LatentState,
|
||||
stepper: DiffusionStepProtocol,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]: ...
|
||||
step_index: int,
|
||||
) -> tuple[torch.Tensor | None, torch.Tensor | None]: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModalitySpec:
|
||||
"""Specification for one modality passed to a diffusion stage.
|
||||
Carries everything needed to build the initial noised latent state
|
||||
and run the denoising loop for a single modality (video or audio).
|
||||
Tools are created by ``DiffusionStage`` from pixel-space dimensions.
|
||||
"""
|
||||
|
||||
context: torch.Tensor
|
||||
conditionings: list[ConditioningItem] = field(default_factory=list)
|
||||
noise_scale: float = 1.0
|
||||
frozen: bool = False
|
||||
initial_latent: torch.Tensor | None = None
|
||||
|
||||
Reference in New Issue
Block a user