632 lines
24 KiB
Python
632 lines
24 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
|
|
guidance_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
|
|
sampler: str
|
|
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
|
|
audio_attention_mask: torch.Tensor | None = None
|
|
conditioning_mask: torch.Tensor | None = None
|
|
conditioning_mask_for_model: torch.Tensor | None = None
|
|
|
|
|
|
@dataclass
|
|
class LTXPhaseResult:
|
|
"""Denoised 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]
|
|
|
|
LTX_ANCESTRAL_ETA = 1.0
|
|
LTX_ANCESTRAL_S_NOISE = 1.0
|
|
LTX_ANCESTRAL_NOISE_SEED_OFFSET = 10000
|
|
LTX_OFFICIAL_DEFAULT_SEED = 10
|
|
|
|
|
|
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,
|
|
sampler: str = "euler",
|
|
generator: torch.Generator | list[torch.Generator] | None = None,
|
|
conditioning_mask: torch.Tensor | None = None,
|
|
) -> 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._sampler = sampler
|
|
self._conditioning_mask = conditioning_mask
|
|
self._ancestral_generators = (
|
|
_make_ancestral_generators(generator, pipeline.device) if sampler == "euler_ancestral" else None
|
|
)
|
|
self._step_index = 0
|
|
|
|
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
|
|
del t, return_dict, generator
|
|
if self._sampler == "euler_ancestral":
|
|
video_out = _ancestral_euler_step_from_velocity(
|
|
latents[0],
|
|
noise_pred[0],
|
|
self._pipeline.scheduler.sigmas,
|
|
self._step_index,
|
|
self._ancestral_generators,
|
|
)
|
|
if self._conditioning_mask is not None:
|
|
mask = self._conditioning_mask.unsqueeze(-1).to(video_out.dtype)
|
|
video_out = torch.lerp(video_out, latents[0], mask)
|
|
elif 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,
|
|
)
|
|
if self._sampler == "euler_ancestral":
|
|
audio_out = _ancestral_euler_step_from_velocity(
|
|
latents[1],
|
|
noise_pred[1],
|
|
self._audio_scheduler.sigmas,
|
|
self._step_index,
|
|
self._ancestral_generators,
|
|
)
|
|
else:
|
|
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 _make_ancestral_generators(
|
|
generator: torch.Generator | list[torch.Generator] | None,
|
|
device: torch.device,
|
|
) -> torch.Generator | list[torch.Generator]:
|
|
"""Create the official independent ``seed + 10000`` noise stream."""
|
|
|
|
def _offset(source: torch.Generator | None) -> torch.Generator:
|
|
seed = LTX_OFFICIAL_DEFAULT_SEED if source is None else source.initial_seed()
|
|
return torch.Generator(device=device).manual_seed(seed + LTX_ANCESTRAL_NOISE_SEED_OFFSET)
|
|
|
|
if isinstance(generator, list):
|
|
return [_offset(item) for item in generator]
|
|
return _offset(generator)
|
|
|
|
|
|
def _randn_like_with_generators(
|
|
sample: torch.Tensor,
|
|
generators: torch.Generator | list[torch.Generator],
|
|
) -> torch.Tensor:
|
|
if not isinstance(generators, list):
|
|
return torch.randn(sample.shape, generator=generators, dtype=sample.dtype, device=sample.device)
|
|
if len(generators) != sample.shape[0]:
|
|
raise ValueError(
|
|
f"LTX ancestral sampling received {len(generators)} generators for batch size {sample.shape[0]}."
|
|
)
|
|
return torch.cat(
|
|
[
|
|
torch.randn((1, *sample.shape[1:]), generator=item, dtype=sample.dtype, device=sample.device)
|
|
for item in generators
|
|
]
|
|
)
|
|
|
|
|
|
def _ancestral_euler_step_from_velocity(
|
|
sample: torch.Tensor,
|
|
velocity: torch.Tensor,
|
|
sigmas: torch.Tensor,
|
|
step_index: int,
|
|
generators: torch.Generator | list[torch.Generator] | None,
|
|
) -> torch.Tensor:
|
|
"""Apply the official LTX-2.5 rectified-flow ancestral Euler step."""
|
|
sigma = sigmas[step_index].to(torch.float32)
|
|
sigma_next = sigmas[step_index + 1].to(torch.float32)
|
|
# Official X0Model materializes the denoised prediction in the model
|
|
# dtype before the ancestral step promotes it back to float32.
|
|
denoised_model_dtype = (sample.float() - velocity.float() * sigma).to(sample.dtype)
|
|
denoised = denoised_model_dtype.float()
|
|
if sigma_next == 0:
|
|
return denoised.to(sample.dtype)
|
|
if generators is None:
|
|
raise ValueError("LTX ancestral Euler sampling requires a noise generator.")
|
|
|
|
downstep_ratio = 1.0 + (sigma_next / sigma - 1.0) * LTX_ANCESTRAL_ETA
|
|
sigma_down = sigma_next * downstep_ratio
|
|
sigma_down_ratio = sigma_down / sigma
|
|
x_next = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised
|
|
|
|
alpha_next = 1.0 - sigma_next
|
|
alpha_down = 1.0 - sigma_down
|
|
renoise_coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).clamp(min=0).sqrt()
|
|
noise = _randn_like_with_generators(sample, generators)
|
|
x_next = (alpha_next / alpha_down) * x_next + noise.float() * LTX_ANCESTRAL_S_NOISE * renoise_coeff
|
|
result = x_next.to(sample.dtype)
|
|
return result
|
|
|
|
|
|
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 LTX one-stage pipelines omit the latent when constructing
|
|
# this schedule, so the shift stays at the max sequence anchor.
|
|
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,
|
|
sampler: str = "euler",
|
|
generator: torch.Generator | list[torch.Generator] | None = None,
|
|
conditioning_mask: torch.Tensor | None = None,
|
|
) -> 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,
|
|
sampler=sampler,
|
|
generator=generator,
|
|
conditioning_mask=conditioning_mask,
|
|
)
|
|
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]:
|
|
transformer = pipeline.transformer
|
|
video_coords = 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 = 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 _first_frame_keyframes_mask(reference: torch.Tensor, latent_num_frames: int) -> torch.Tensor:
|
|
"""Mark the first causal latent frame, matching official LTX-2.5."""
|
|
if latent_num_frames <= 0:
|
|
raise ValueError(f"LTX latent frame count must be positive, got {latent_num_frames}.")
|
|
tokens_per_frame, remainder = divmod(reference.shape[1], latent_num_frames)
|
|
if remainder:
|
|
raise ValueError(
|
|
f"LTX video token count {reference.shape[1]} is not divisible by {latent_num_frames} latent frames."
|
|
)
|
|
mask = reference.new_zeros((reference.shape[0], reference.shape[1], 1))
|
|
mask[:, :tokens_per_frame] = 1
|
|
return mask
|
|
|
|
|
|
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
|
|
keyframes_mask = None
|
|
if getattr(pipeline.transformer.config, "use_keyframes_abs_pos_embedding", False):
|
|
keyframes_mask = _first_frame_keyframes_mask(hidden_states, forward_ctx.latent_num_frames)
|
|
return {
|
|
"hidden_states": hidden_states,
|
|
"audio_hidden_states": audio_hidden_states,
|
|
"keyframes_mask": keyframes_mask,
|
|
"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,
|
|
"audio_attention_mask": denoise_ctx.audio_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 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_guidance_parallel_step_output(
|
|
latents,
|
|
guidance_parallel_ready=forward_ctx.guidance_parallel_ready,
|
|
)
|
|
|
|
|
|
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)
|
|
guidance_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,
|
|
sampler=phase_recipe.sampler,
|
|
generator=request_inputs.generator,
|
|
conditioning_mask=conditioning_mask,
|
|
)
|
|
forward_ctx = LTXForwardContext(
|
|
req=req,
|
|
request_inputs=request_inputs,
|
|
prompt_context=prompt_context,
|
|
device=device,
|
|
guidance_parallel_ready=guidance_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,
|
|
sampler=phase_recipe.sampler,
|
|
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)
|
|
ring_degree = getattr(pipeline.od_config.parallel_config, "ring_degree", 1) or 1
|
|
if padded_audio_num_frames > original_audio_num_frames and ring_degree > 1:
|
|
raise ValueError(
|
|
"LTX audio padding requires an attention mask, which Ring sequence parallelism does not support. "
|
|
"Use Ulysses-only SP or choose a request whose audio latent length is divisible by the SP size."
|
|
)
|
|
denoise_ctx = LTXDenoiseContext(
|
|
latents=latents,
|
|
audio_latents=audio_latents,
|
|
video_coords=video_coords,
|
|
audio_coords=audio_coords,
|
|
audio_attention_mask=(
|
|
torch.arange(padded_audio_num_frames, device=audio_latents.device)
|
|
.lt(original_audio_num_frames)
|
|
.unsqueeze(0)
|
|
.expand(audio_latents.shape[0], -1)
|
|
if padded_audio_num_frames > original_audio_num_frames
|
|
else None
|
|
),
|
|
conditioning_mask=conditioning_mask,
|
|
)
|
|
denoise_ctx = pipeline._prepare_denoise_context_for_guidance(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,
|
|
)
|