482 lines
17 KiB
Python
482 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Shared denoise execution primitives for LTX pipelines."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import math
|
|
from collections.abc import Callable, Iterable
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any, Protocol
|
|
|
|
import torch
|
|
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
|
|
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
|
|
from .ltx2_guidance import euler_step_from_velocity
|
|
from .ltx2_latents import LTXAVState, unpack_audio_latents, unpad_audio_latents
|
|
|
|
if TYPE_CHECKING:
|
|
from .ltx2_conditioning import LTXPromptContext
|
|
from .ltx2_recipes import LTXPhaseRecipe
|
|
from .ltx2_request import LTXRequestInputs
|
|
|
|
|
|
@dataclass
|
|
class LTXForwardContext:
|
|
"""Immutable metadata and schedulers for one LTX denoise phase."""
|
|
|
|
req: DiffusionRequestBatch
|
|
request_inputs: LTXRequestInputs
|
|
prompt_context: LTXPromptContext
|
|
device: torch.device
|
|
cfg_parallel_ready: bool
|
|
attention_kwargs: dict[str, Any] | None
|
|
latent_num_frames: int
|
|
latent_height: int
|
|
latent_width: int
|
|
latent_mel_bins: int
|
|
original_audio_num_frames: int
|
|
padded_audio_num_frames: int
|
|
timesteps: torch.Tensor
|
|
audio_scheduler: Any
|
|
video_audio_step_adapter: Any
|
|
|
|
@property
|
|
def batch_size(self) -> int:
|
|
return self.prompt_context.batch_size
|
|
|
|
@property
|
|
def num_videos_per_prompt(self) -> int:
|
|
return self.request_inputs.num_videos_per_prompt
|
|
|
|
|
|
@dataclass
|
|
class LTXDenoiseContext:
|
|
"""Mutable AV state and positional metadata for a denoise phase."""
|
|
|
|
latents: torch.Tensor
|
|
audio_latents: torch.Tensor
|
|
video_coords: torch.Tensor
|
|
audio_coords: torch.Tensor
|
|
conditioning_mask: torch.Tensor | None = None
|
|
conditioning_mask_for_model: torch.Tensor | None = None
|
|
|
|
|
|
@dataclass
|
|
class LTXPhaseResult:
|
|
"""Denoised, unpacked AV latents and the context used to produce them."""
|
|
|
|
forward_context: LTXForwardContext
|
|
video: torch.Tensor
|
|
audio: torch.Tensor
|
|
audio_for_next_phase: torch.Tensor | None = None
|
|
|
|
|
|
class LTXDenoisePipeline(Protocol):
|
|
"""Pipeline state required by :class:`LTXDenoiseExecutor`."""
|
|
|
|
@property
|
|
def interrupt(self) -> bool: ...
|
|
|
|
def progress_bar(self, iterable=None, total=None): ...
|
|
|
|
|
|
LTXDenoiseStep = Callable[[int, torch.Tensor, LTXAVState], LTXAVState]
|
|
|
|
|
|
class LTXDenoiseExecutor:
|
|
"""Run the one shared LTX denoise loop.
|
|
|
|
Prediction and scheduler math remain injectable so structural refactors do
|
|
not change the existing LTX2/LTX2.3 numerical paths. Guidance will replace
|
|
that step policy independently.
|
|
"""
|
|
|
|
@staticmethod
|
|
def run(
|
|
pipeline: LTXDenoisePipeline,
|
|
state: LTXAVState,
|
|
timesteps: Iterable[torch.Tensor],
|
|
step: LTXDenoiseStep,
|
|
) -> LTXAVState:
|
|
timesteps = tuple(timesteps)
|
|
with pipeline.progress_bar(total=len(timesteps)) as progress_bar:
|
|
for index, timestep in enumerate(timesteps):
|
|
if pipeline.interrupt:
|
|
continue
|
|
state = step(index, timestep, state)
|
|
progress_bar.update()
|
|
return state
|
|
|
|
|
|
def calculate_shift(
|
|
image_seq_len: int,
|
|
base_seq_len: int = 256,
|
|
max_seq_len: int = 4096,
|
|
base_shift: float = 0.5,
|
|
max_shift: float = 1.15,
|
|
) -> float:
|
|
slope = (max_shift - base_shift) / (max_seq_len - base_seq_len)
|
|
intercept = base_shift - slope * base_seq_len
|
|
return image_seq_len * slope + intercept
|
|
|
|
|
|
class LTXVideoAudioStepAdapter:
|
|
"""Expose the shared LTX Euler update through the distributed scheduler API."""
|
|
|
|
def __init__(
|
|
self,
|
|
pipeline: Any,
|
|
audio_scheduler: Any,
|
|
latent_num_frames: int,
|
|
latent_height: int,
|
|
latent_width: int,
|
|
*,
|
|
image_conditioned: bool,
|
|
) -> None:
|
|
self._pipeline = pipeline
|
|
self._audio_scheduler = audio_scheduler
|
|
self._latent_num_frames = latent_num_frames
|
|
self._latent_height = latent_height
|
|
self._latent_width = latent_width
|
|
self._image_conditioned = image_conditioned
|
|
self._step_index = 0
|
|
|
|
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
|
|
del t, return_dict, generator
|
|
if self._image_conditioned:
|
|
video_out = self._pipeline._step_video_latents_i2v(
|
|
noise_pred[0],
|
|
latents[0],
|
|
self._step_index,
|
|
self._latent_num_frames,
|
|
self._latent_height,
|
|
self._latent_width,
|
|
)
|
|
else:
|
|
video_out = euler_step_from_velocity(
|
|
latents[0],
|
|
noise_pred[0],
|
|
self._pipeline.scheduler.sigmas,
|
|
self._step_index,
|
|
)
|
|
audio_out = euler_step_from_velocity(
|
|
latents[1],
|
|
noise_pred[1],
|
|
self._audio_scheduler.sigmas,
|
|
self._step_index,
|
|
)
|
|
self._step_index += 1
|
|
return ((video_out, audio_out),)
|
|
|
|
|
|
def _official_ltx_sigmas(scheduler: Any, steps: int, device: torch.device) -> torch.Tensor:
|
|
"""Build the official LTX one-stage sigma schedule in torch fp32."""
|
|
config = scheduler.config
|
|
base_anchor = config.get("base_image_seq_len", 1024)
|
|
max_anchor = config.get("max_image_seq_len", 4096)
|
|
base_shift = config.get("base_shift", 0.95)
|
|
max_shift = config.get("max_shift", 2.05)
|
|
|
|
sigmas = torch.linspace(1.0, 0.0, steps + 1)
|
|
slope = (max_shift - base_shift) / (max_anchor - base_anchor)
|
|
# Official LTX2 one-stage intentionally uses the max sequence anchor, so the shift stays at max_shift.
|
|
sigma_shift = max_anchor * slope + (base_shift - slope * base_anchor)
|
|
exp_shift = math.exp(sigma_shift)
|
|
sigmas = torch.where(sigmas != 0, exp_shift / (exp_shift + (1 / sigmas - 1)), 0)
|
|
|
|
# Official non-distilled checkpoints explicitly set this to 0.1. Missing
|
|
# or None means the loaded scheduler disabled terminal stretching.
|
|
terminal = config.get("shift_terminal")
|
|
if terminal is not None:
|
|
non_zero = sigmas != 0
|
|
one_minus_sigmas = 1.0 - sigmas[non_zero]
|
|
scale = one_minus_sigmas[-1] / (1.0 - terminal)
|
|
sigmas[non_zero] = 1.0 - one_minus_sigmas / scale
|
|
return sigmas.to(dtype=torch.float32, device=device)
|
|
|
|
|
|
def _set_scheduler_sigmas(scheduler: Any, sigmas: torch.Tensor) -> torch.Tensor:
|
|
sigmas = sigmas.to(torch.float32)
|
|
timesteps = sigmas[:-1] * scheduler.config.get("num_train_timesteps", 1000)
|
|
scheduler.sigmas = sigmas
|
|
scheduler.timesteps = timesteps
|
|
scheduler.num_inference_steps = len(timesteps)
|
|
scheduler._step_index = None
|
|
scheduler._begin_index = None
|
|
return timesteps
|
|
|
|
|
|
def prepare_scheduler_stage(
|
|
pipeline: Any,
|
|
request_inputs: LTXRequestInputs,
|
|
*,
|
|
device: torch.device,
|
|
sigmas: list[float] | None,
|
|
timesteps: list[int] | None,
|
|
latent_num_frames: int,
|
|
latent_height: int,
|
|
latent_width: int,
|
|
use_official_sigma_schedule: bool,
|
|
image_conditioned: bool = False,
|
|
) -> tuple[Any, Any, torch.Tensor]:
|
|
if sigmas is not None and timesteps is not None:
|
|
raise ValueError("Only one of `sigmas` or `timesteps` may be provided.")
|
|
|
|
audio_scheduler = copy.deepcopy(pipeline.scheduler)
|
|
video_audio_step_adapter = LTXVideoAudioStepAdapter(
|
|
pipeline,
|
|
audio_scheduler,
|
|
latent_num_frames,
|
|
latent_height,
|
|
latent_width,
|
|
image_conditioned=image_conditioned,
|
|
)
|
|
if sigmas is not None:
|
|
scheduler_sigmas = torch.as_tensor(sigmas, dtype=torch.float32, device=device)
|
|
if scheduler_sigmas.ndim != 1 or scheduler_sigmas.numel() < 2:
|
|
raise ValueError("An LTX custom sigma schedule must contain at least two boundary values.")
|
|
if scheduler_sigmas[-1] != 0:
|
|
scheduler_sigmas = torch.cat([scheduler_sigmas, scheduler_sigmas.new_zeros(1)])
|
|
timesteps_tensor = _set_scheduler_sigmas(pipeline.scheduler, scheduler_sigmas)
|
|
_set_scheduler_sigmas(audio_scheduler, scheduler_sigmas.clone())
|
|
return audio_scheduler, video_audio_step_adapter, timesteps_tensor
|
|
|
|
if sigmas is None and timesteps is None and use_official_sigma_schedule:
|
|
scheduler_sigmas = _official_ltx_sigmas(pipeline.scheduler, request_inputs.num_inference_steps, device)
|
|
timesteps_tensor = _set_scheduler_sigmas(pipeline.scheduler, scheduler_sigmas)
|
|
_set_scheduler_sigmas(audio_scheduler, scheduler_sigmas.clone())
|
|
return audio_scheduler, video_audio_step_adapter, timesteps_tensor
|
|
|
|
mu = calculate_shift(
|
|
pipeline.scheduler.config.get("max_image_seq_len", 4096),
|
|
pipeline.scheduler.config.get("base_image_seq_len", 1024),
|
|
pipeline.scheduler.config.get("max_image_seq_len", 4096),
|
|
pipeline.scheduler.config.get("base_shift", 0.95),
|
|
pipeline.scheduler.config.get("max_shift", 2.05),
|
|
)
|
|
retrieve_timesteps(
|
|
audio_scheduler,
|
|
request_inputs.num_inference_steps,
|
|
device,
|
|
timesteps,
|
|
sigmas=sigmas,
|
|
mu=mu,
|
|
)
|
|
timesteps_tensor, _ = retrieve_timesteps(
|
|
pipeline.scheduler,
|
|
request_inputs.num_inference_steps,
|
|
device,
|
|
timesteps,
|
|
sigmas=sigmas,
|
|
mu=mu,
|
|
)
|
|
return audio_scheduler, video_audio_step_adapter, timesteps_tensor
|
|
|
|
|
|
def prepare_rope_coords_stage(
|
|
pipeline: Any,
|
|
forward_ctx: LTXForwardContext,
|
|
latents: torch.Tensor,
|
|
audio_latents: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
video_coords = pipeline.transformer.rope.prepare_video_coords(
|
|
latents.shape[0],
|
|
forward_ctx.latent_num_frames,
|
|
forward_ctx.latent_height,
|
|
forward_ctx.latent_width,
|
|
latents.device,
|
|
fps=forward_ctx.request_inputs.frame_rate,
|
|
)
|
|
audio_coords = pipeline.transformer.audio_rope.prepare_audio_coords(
|
|
audio_latents.shape[0],
|
|
forward_ctx.padded_audio_num_frames,
|
|
audio_latents.device,
|
|
)
|
|
return video_coords, audio_coords
|
|
|
|
|
|
def build_transformer_kwargs(
|
|
pipeline: Any,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
*,
|
|
hidden_states: torch.Tensor,
|
|
audio_hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
audio_encoder_hidden_states: torch.Tensor,
|
|
encoder_attention_mask: torch.Tensor | None,
|
|
audio_encoder_attention_mask: torch.Tensor | None,
|
|
ts: torch.Tensor,
|
|
attention_kwargs: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
del encoder_attention_mask, audio_encoder_attention_mask
|
|
return {
|
|
"hidden_states": hidden_states,
|
|
"audio_hidden_states": audio_hidden_states,
|
|
"encoder_hidden_states": encoder_hidden_states,
|
|
"audio_encoder_hidden_states": audio_encoder_hidden_states,
|
|
**pipeline._denoise_timestep_kwargs(
|
|
ts,
|
|
forward_ctx,
|
|
denoise_ctx,
|
|
video_token_count=hidden_states.shape[1],
|
|
audio_token_count=audio_hidden_states.shape[1],
|
|
),
|
|
# This is valid only because LTX connectors replace every padding token
|
|
# with a learned register, making all output context tokens valid.
|
|
"encoder_attention_mask": None,
|
|
"audio_encoder_attention_mask": None,
|
|
"num_frames": forward_ctx.latent_num_frames,
|
|
"height": forward_ctx.latent_height,
|
|
"width": forward_ctx.latent_width,
|
|
"fps": forward_ctx.request_inputs.frame_rate,
|
|
"audio_num_frames": forward_ctx.padded_audio_num_frames,
|
|
"video_coords": denoise_ctx.video_coords,
|
|
"audio_coords": denoise_ctx.audio_coords,
|
|
"attention_kwargs": forward_ctx.attention_kwargs if attention_kwargs is None else attention_kwargs,
|
|
"return_dict": False,
|
|
}
|
|
|
|
|
|
def step_denoised_latents(
|
|
pipeline: Any,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
noise_pred_video: torch.Tensor,
|
|
noise_pred_audio: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
latents = pipeline.scheduler_step_maybe_with_cfg(
|
|
(noise_pred_video, noise_pred_audio),
|
|
(timestep, timestep),
|
|
(denoise_ctx.latents, denoise_ctx.audio_latents),
|
|
do_true_cfg=pipeline.do_classifier_free_guidance,
|
|
per_request_scheduler=forward_ctx.video_audio_step_adapter,
|
|
)
|
|
return pipeline._synchronize_cfg_parallel_step_output(
|
|
latents,
|
|
do_true_cfg=pipeline.do_classifier_free_guidance,
|
|
)
|
|
|
|
|
|
class LTXPhaseExecutor:
|
|
"""Prepare and execute one LTX phase without owning model modules."""
|
|
|
|
@staticmethod
|
|
def run(
|
|
pipeline: Any,
|
|
req: DiffusionRequestBatch,
|
|
request_inputs: LTXRequestInputs,
|
|
*,
|
|
noise_scale: float,
|
|
sigmas: list[float] | None,
|
|
timesteps: list[int] | None,
|
|
attention_kwargs: dict[str, Any] | None,
|
|
phase_recipe: LTXPhaseRecipe,
|
|
image: Any | None = None,
|
|
prompt_context: LTXPromptContext | None = None,
|
|
) -> LTXPhaseResult:
|
|
pipeline._check_forward_inputs(request_inputs, image=image)
|
|
cfg_parallel_ready = pipeline._setup_forward_runtime(req, request_inputs, attention_kwargs)
|
|
device = pipeline.device
|
|
if prompt_context is None:
|
|
prompt_context = pipeline._prepare_prompt_context(
|
|
prompt=request_inputs.prompt,
|
|
negative_prompt=request_inputs.negative_prompt,
|
|
prompt_embeds=request_inputs.prompt_embeds,
|
|
negative_prompt_embeds=request_inputs.negative_prompt_embeds,
|
|
prompt_attention_mask=request_inputs.prompt_attention_mask,
|
|
negative_prompt_attention_mask=request_inputs.negative_prompt_attention_mask,
|
|
num_videos_per_prompt=request_inputs.num_videos_per_prompt,
|
|
max_sequence_length=request_inputs.max_sequence_length,
|
|
)
|
|
|
|
latent_num_frames, latent_height, latent_width = pipeline._resolve_video_latent_dimensions(request_inputs)
|
|
latents, conditioning_mask = pipeline._prepare_video_latents_stage(
|
|
request_inputs,
|
|
prompt_context,
|
|
device=device,
|
|
noise_scale=noise_scale,
|
|
image=image,
|
|
)
|
|
audio_latents, original_audio_num_frames, padded_audio_num_frames, latent_mel_bins = (
|
|
pipeline._prepare_audio_latents_stage(
|
|
request_inputs,
|
|
prompt_context,
|
|
device=device,
|
|
noise_scale=noise_scale,
|
|
)
|
|
)
|
|
audio_scheduler, video_audio_step_adapter, timesteps_tensor = prepare_scheduler_stage(
|
|
pipeline,
|
|
request_inputs,
|
|
device=device,
|
|
sigmas=sigmas,
|
|
timesteps=timesteps,
|
|
latent_num_frames=latent_num_frames,
|
|
latent_height=latent_height,
|
|
latent_width=latent_width,
|
|
use_official_sigma_schedule=phase_recipe.use_official_sigma_schedule,
|
|
image_conditioned=conditioning_mask is not None,
|
|
)
|
|
forward_ctx = LTXForwardContext(
|
|
req=req,
|
|
request_inputs=request_inputs,
|
|
prompt_context=prompt_context,
|
|
device=device,
|
|
cfg_parallel_ready=cfg_parallel_ready,
|
|
attention_kwargs=attention_kwargs,
|
|
latent_num_frames=latent_num_frames,
|
|
latent_height=latent_height,
|
|
latent_width=latent_width,
|
|
latent_mel_bins=latent_mel_bins,
|
|
original_audio_num_frames=original_audio_num_frames,
|
|
padded_audio_num_frames=padded_audio_num_frames,
|
|
timesteps=timesteps_tensor,
|
|
audio_scheduler=audio_scheduler,
|
|
video_audio_step_adapter=video_audio_step_adapter,
|
|
)
|
|
video_coords, audio_coords = prepare_rope_coords_stage(pipeline, forward_ctx, latents, audio_latents)
|
|
denoise_ctx = LTXDenoiseContext(
|
|
latents=latents,
|
|
audio_latents=audio_latents,
|
|
video_coords=video_coords,
|
|
audio_coords=audio_coords,
|
|
conditioning_mask=conditioning_mask,
|
|
)
|
|
denoise_ctx = pipeline._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx)
|
|
state = LTXDenoiseExecutor.run(
|
|
pipeline,
|
|
LTXAVState(video=denoise_ctx.latents, audio=denoise_ctx.audio_latents),
|
|
forward_ctx.timesteps,
|
|
lambda index, timestep, state: pipeline._denoise_step(
|
|
index,
|
|
timestep,
|
|
state,
|
|
forward_ctx,
|
|
denoise_ctx,
|
|
),
|
|
)
|
|
denoise_ctx.latents = state.video
|
|
denoise_ctx.audio_latents = state.audio
|
|
latents, audio_latents = pipeline._unpack_and_denormalize_stage(
|
|
forward_ctx,
|
|
state.video,
|
|
state.audio,
|
|
)
|
|
normalized_audio = unpack_audio_latents(
|
|
unpad_audio_latents(state.audio, forward_ctx.original_audio_num_frames),
|
|
num_mel_bins=forward_ctx.latent_mel_bins,
|
|
)
|
|
return LTXPhaseResult(
|
|
forward_context=forward_ctx,
|
|
video=latents,
|
|
audio=audio_latents,
|
|
audio_for_next_phase=normalized_audio,
|
|
)
|