431 lines
15 KiB
Python
431 lines
15 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
|
|
from collections.abc import Callable, Iterable
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any, Protocol
|
|
|
|
import numpy as np
|
|
import torch
|
|
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
|
|
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
|
|
from .ltx2_latents import LTXAVState
|
|
|
|
if TYPE_CHECKING:
|
|
from .ltx2_conditioning import LTXPromptContext
|
|
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_scheduler: 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
|
|
|
|
|
|
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 VideoAudioScheduler:
|
|
"""Composite scheduler dispatching video and audio updates."""
|
|
|
|
def __init__(self, video_scheduler, audio_scheduler):
|
|
self.video_scheduler = video_scheduler
|
|
self.audio_scheduler = audio_scheduler
|
|
|
|
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
|
|
video_out = self.video_scheduler.step(
|
|
noise_pred[0],
|
|
t[0],
|
|
latents[0],
|
|
return_dict=False,
|
|
generator=generator,
|
|
)[0]
|
|
audio_out = self.audio_scheduler.step(
|
|
noise_pred[1],
|
|
t[1],
|
|
latents[1],
|
|
return_dict=False,
|
|
generator=generator,
|
|
)[0]
|
|
return ((video_out, audio_out),)
|
|
|
|
|
|
class I2VVideoAudioScheduler:
|
|
"""Update the unconditioned video frames and the full audio state."""
|
|
|
|
def __init__(self, pipeline, audio_scheduler, latent_num_frames, latent_height, latent_width):
|
|
self.video_scheduler = pipeline.scheduler
|
|
self.audio_scheduler = audio_scheduler
|
|
self._pipeline = pipeline
|
|
self._latent_num_frames = latent_num_frames
|
|
self._latent_height = latent_height
|
|
self._latent_width = latent_width
|
|
|
|
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
|
|
video_out = self._pipeline._step_video_latents_i2v(
|
|
noise_pred[0],
|
|
latents[0],
|
|
t[0],
|
|
self._latent_num_frames,
|
|
self._latent_height,
|
|
self._latent_width,
|
|
)
|
|
audio_out = self.audio_scheduler.step(
|
|
noise_pred[1],
|
|
t[1],
|
|
latents[1],
|
|
return_dict=False,
|
|
generator=generator,
|
|
)[0]
|
|
return ((video_out, audio_out),)
|
|
|
|
|
|
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,
|
|
) -> 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_scheduler = pipeline._make_video_audio_scheduler(
|
|
audio_scheduler,
|
|
latent_num_frames,
|
|
latent_height,
|
|
latent_width,
|
|
)
|
|
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_scheduler, timesteps_tensor
|
|
|
|
sigmas = np.linspace(1.0, 1 / request_inputs.num_inference_steps, request_inputs.num_inference_steps)
|
|
# Official LTX2 one-stage omits the latent, so mu stays at max_shift across resolutions.
|
|
# Mu varies with latent token count only when the scheduler receives a latent, as in official HQ two-stage;
|
|
# explicit sigma schedules bypass mu entirely.
|
|
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_scheduler, 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,
|
|
) -> dict[str, Any]:
|
|
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),
|
|
"encoder_attention_mask": encoder_attention_mask,
|
|
"audio_encoder_attention_mask": audio_encoder_attention_mask,
|
|
"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,
|
|
"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_scheduler,
|
|
)
|
|
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,
|
|
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(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_scheduler, 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,
|
|
)
|
|
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_scheduler=video_audio_scheduler,
|
|
)
|
|
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,
|
|
)
|
|
return LTXPhaseResult(forward_context=forward_ctx, video=latents, audio=audio_latents)
|