380 lines
15 KiB
Python
380 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Guidance strategies for the LTX model family."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Any, Protocol
|
|
|
|
import torch
|
|
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg
|
|
|
|
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
|
get_cfg_group,
|
|
get_classifier_free_guidance_rank,
|
|
get_classifier_free_guidance_world_size,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from .ltx2_denoise import LTXDenoiseContext, LTXForwardContext
|
|
from .ltx2_latents import LTXAVState
|
|
|
|
|
|
class LTXGuidanceStrategy(Protocol):
|
|
"""Model-version guidance behavior consumed by the shared pipeline."""
|
|
|
|
def validate_cfg_world_size(self, cfg_world_size: int) -> None: ...
|
|
|
|
def prepare_denoise_context(
|
|
self,
|
|
pipeline: Any,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> LTXDenoiseContext: ...
|
|
|
|
def timestep_kwargs(
|
|
self,
|
|
ts: torch.Tensor,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> dict[str, torch.Tensor]: ...
|
|
|
|
def combine_cfg_noise(
|
|
self,
|
|
pipeline: Any,
|
|
positive_noise_pred: tuple[torch.Tensor, torch.Tensor],
|
|
negative_noise_pred: tuple[torch.Tensor, torch.Tensor],
|
|
true_cfg_scale: float,
|
|
cfg_normalize: bool,
|
|
context: dict[str, Any],
|
|
) -> tuple[torch.Tensor, torch.Tensor]: ...
|
|
|
|
def predict_noise(
|
|
self,
|
|
pipeline: Any,
|
|
index: int,
|
|
timestep: torch.Tensor,
|
|
state: LTXAVState,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> tuple[torch.Tensor, torch.Tensor]: ...
|
|
|
|
|
|
class _LTXGuidanceStrategyBase:
|
|
def validate_cfg_world_size(self, cfg_world_size: int) -> None:
|
|
del cfg_world_size
|
|
|
|
def prepare_denoise_context(
|
|
self,
|
|
pipeline: Any,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> LTXDenoiseContext:
|
|
del pipeline, forward_ctx
|
|
return denoise_ctx
|
|
|
|
def timestep_kwargs(
|
|
self,
|
|
ts: torch.Tensor,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> dict[str, torch.Tensor]:
|
|
del forward_ctx, denoise_ctx
|
|
return {"timestep": ts}
|
|
|
|
|
|
def combine_velocity_via_x0(
|
|
sample: torch.Tensor,
|
|
positive_velocity: torch.Tensor,
|
|
negative_velocity: torch.Tensor,
|
|
sigma: torch.Tensor,
|
|
guidance_scale: float,
|
|
) -> torch.Tensor:
|
|
"""Apply CFG in x0 space and convert the result back to velocity."""
|
|
x0_cond = sample - positive_velocity * sigma
|
|
x0_uncond = sample - negative_velocity * sigma
|
|
x0_guided = x0_cond + (guidance_scale - 1) * (x0_cond - x0_uncond)
|
|
return (sample - x0_guided) / sigma
|
|
|
|
|
|
class LTXLegacyVelocityGuidance(_LTXGuidanceStrategyBase):
|
|
"""Existing LTX2 Diffusers-style velocity-space CFG behavior."""
|
|
|
|
def combine_cfg_noise(
|
|
self,
|
|
pipeline: Any,
|
|
positive_noise_pred: tuple[torch.Tensor, torch.Tensor],
|
|
negative_noise_pred: tuple[torch.Tensor, torch.Tensor],
|
|
true_cfg_scale: float,
|
|
cfg_normalize: bool,
|
|
context: dict[str, Any],
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
del context
|
|
video_pos, audio_pos = positive_noise_pred
|
|
video_neg, audio_neg = negative_noise_pred
|
|
video_combined = CFGParallelMixin.combine_cfg_noise(
|
|
pipeline,
|
|
video_pos,
|
|
video_neg,
|
|
true_cfg_scale,
|
|
cfg_normalize,
|
|
)
|
|
audio_combined = CFGParallelMixin.combine_cfg_noise(
|
|
pipeline,
|
|
audio_pos,
|
|
audio_neg,
|
|
true_cfg_scale,
|
|
cfg_normalize,
|
|
)
|
|
if pipeline.guidance_rescale and pipeline.guidance_rescale > 0:
|
|
video_combined = rescale_noise_cfg(
|
|
video_combined,
|
|
video_pos,
|
|
guidance_rescale=pipeline.guidance_rescale,
|
|
)
|
|
audio_combined = rescale_noise_cfg(
|
|
audio_combined,
|
|
audio_pos,
|
|
guidance_rescale=pipeline.guidance_rescale,
|
|
)
|
|
return video_combined, audio_combined
|
|
|
|
def predict_noise(
|
|
self,
|
|
pipeline: Any,
|
|
index: int,
|
|
timestep: torch.Tensor,
|
|
state: LTXAVState,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
del index
|
|
prompt_context = forward_ctx.prompt_context
|
|
video_input = state.video.to(prompt_context.positive_connector_prompt_embeds.dtype)
|
|
audio_input = state.audio.to(prompt_context.positive_connector_prompt_embeds.dtype)
|
|
expanded_timestep = timestep.expand(video_input.shape[0])
|
|
positive_kwargs = pipeline._build_transformer_kwargs(
|
|
forward_ctx,
|
|
denoise_ctx,
|
|
hidden_states=video_input,
|
|
audio_hidden_states=audio_input,
|
|
encoder_hidden_states=prompt_context.positive_connector_prompt_embeds,
|
|
audio_encoder_hidden_states=prompt_context.positive_connector_audio_prompt_embeds,
|
|
encoder_attention_mask=prompt_context.positive_connector_attention_mask,
|
|
audio_encoder_attention_mask=prompt_context.positive_connector_attention_mask,
|
|
ts=expanded_timestep,
|
|
)
|
|
negative_kwargs = (
|
|
{
|
|
**positive_kwargs,
|
|
"encoder_hidden_states": prompt_context.negative_connector_prompt_embeds,
|
|
"audio_encoder_hidden_states": prompt_context.negative_connector_audio_prompt_embeds,
|
|
"encoder_attention_mask": prompt_context.negative_connector_attention_mask,
|
|
"audio_encoder_attention_mask": prompt_context.negative_connector_attention_mask,
|
|
}
|
|
if pipeline.do_classifier_free_guidance
|
|
else None
|
|
)
|
|
return pipeline.predict_noise_maybe_with_cfg(
|
|
do_true_cfg=pipeline.do_classifier_free_guidance,
|
|
true_cfg_scale=forward_ctx.request_inputs.guidance_scale,
|
|
positive_kwargs=positive_kwargs,
|
|
negative_kwargs=negative_kwargs,
|
|
cfg_normalize=False,
|
|
)
|
|
|
|
|
|
class LTXOfficialX0Guidance(_LTXGuidanceStrategyBase):
|
|
"""Official LTX x0-space CFG behavior currently used by LTX2.3."""
|
|
|
|
def validate_cfg_world_size(self, cfg_world_size: int) -> None:
|
|
if cfg_world_size not in (1, 2):
|
|
raise ValueError(
|
|
f"LTX x0 guidance supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}."
|
|
)
|
|
|
|
def prepare_denoise_context(
|
|
self,
|
|
pipeline: Any,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> LTXDenoiseContext:
|
|
if pipeline.do_classifier_free_guidance and not forward_ctx.cfg_parallel_ready:
|
|
denoise_ctx.video_coords = denoise_ctx.video_coords.repeat(
|
|
(2,) + (1,) * (denoise_ctx.video_coords.ndim - 1)
|
|
)
|
|
denoise_ctx.audio_coords = denoise_ctx.audio_coords.repeat(
|
|
(2,) + (1,) * (denoise_ctx.audio_coords.ndim - 1)
|
|
)
|
|
return denoise_ctx
|
|
|
|
def timestep_kwargs(
|
|
self,
|
|
ts: torch.Tensor,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> dict[str, torch.Tensor]:
|
|
del forward_ctx, denoise_ctx
|
|
return {"timestep": ts, "sigma": ts}
|
|
|
|
def combine_cfg_noise(
|
|
self,
|
|
pipeline: Any,
|
|
positive_noise_pred: tuple[torch.Tensor, torch.Tensor],
|
|
negative_noise_pred: tuple[torch.Tensor, torch.Tensor],
|
|
true_cfg_scale: float,
|
|
cfg_normalize: bool,
|
|
context: dict[str, Any],
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
required = ("video_latents", "audio_latents", "video_sigma", "audio_sigma")
|
|
if any(context.get(name) is None for name in required):
|
|
raise ValueError("LTX x0-space CFG requires video/audio latents and sigmas.")
|
|
|
|
video_pos, audio_pos = positive_noise_pred
|
|
video_neg, audio_neg = negative_noise_pred
|
|
video_combined = combine_velocity_via_x0(
|
|
context["video_latents"],
|
|
video_pos,
|
|
video_neg,
|
|
context["video_sigma"],
|
|
true_cfg_scale,
|
|
)
|
|
audio_combined = combine_velocity_via_x0(
|
|
context["audio_latents"],
|
|
audio_pos,
|
|
audio_neg,
|
|
context["audio_sigma"],
|
|
true_cfg_scale,
|
|
)
|
|
if cfg_normalize:
|
|
video_combined = pipeline.cfg_normalize_function(video_pos, video_combined)
|
|
audio_combined = pipeline.cfg_normalize_function(audio_pos, audio_combined)
|
|
return video_combined, audio_combined
|
|
|
|
def predict_parallel_cfg(
|
|
self,
|
|
pipeline: Any,
|
|
true_cfg_scale: float,
|
|
positive_kwargs: dict[str, Any],
|
|
negative_kwargs: dict[str, Any],
|
|
cfg_normalize: bool = True,
|
|
output_slice: int | None = None,
|
|
**context: Any,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
def maybe_slice(pred: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if output_slice is None:
|
|
return pred
|
|
return pred[0][:, :output_slice], pred[1][:, :output_slice]
|
|
|
|
cfg_world_size = get_classifier_free_guidance_world_size()
|
|
self.validate_cfg_world_size(cfg_world_size)
|
|
if cfg_world_size != 2:
|
|
raise ValueError(f"Parallel LTX x0 guidance requires cfg_parallel_size 2, but got {cfg_world_size}.")
|
|
|
|
cfg_rank = get_classifier_free_guidance_rank()
|
|
branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs
|
|
local_video_pred, local_audio_pred = maybe_slice(pipeline.predict_noise(**branch_kwargs))
|
|
cfg_group = get_cfg_group()
|
|
gathered_video = cfg_group.all_gather(local_video_pred, separate_tensors=True)
|
|
gathered_audio = cfg_group.all_gather(local_audio_pred, separate_tensors=True)
|
|
return self.combine_cfg_noise(
|
|
pipeline,
|
|
(gathered_video[0], gathered_audio[0]),
|
|
(gathered_video[1], gathered_audio[1]),
|
|
true_cfg_scale,
|
|
cfg_normalize,
|
|
context,
|
|
)
|
|
|
|
def predict_noise(
|
|
self,
|
|
pipeline: Any,
|
|
index: int,
|
|
timestep: torch.Tensor,
|
|
state: LTXAVState,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
prompt_context = forward_ctx.prompt_context
|
|
guidance_scale = forward_ctx.request_inputs.guidance_scale
|
|
audio_scheduler = forward_ctx.audio_scheduler
|
|
if forward_ctx.cfg_parallel_ready:
|
|
video_input = state.video.to(prompt_context.positive_connector_prompt_embeds.dtype)
|
|
audio_input = state.audio.to(prompt_context.positive_connector_prompt_embeds.dtype)
|
|
ts = timestep.expand(video_input.shape[0])
|
|
positive_kwargs = pipeline._build_transformer_kwargs(
|
|
forward_ctx,
|
|
denoise_ctx,
|
|
hidden_states=video_input,
|
|
audio_hidden_states=audio_input,
|
|
encoder_hidden_states=prompt_context.positive_connector_prompt_embeds,
|
|
audio_encoder_hidden_states=prompt_context.positive_connector_audio_prompt_embeds,
|
|
encoder_attention_mask=prompt_context.positive_connector_attention_mask,
|
|
audio_encoder_attention_mask=prompt_context.positive_connector_attention_mask,
|
|
ts=ts,
|
|
)
|
|
negative_kwargs = {
|
|
**positive_kwargs,
|
|
"encoder_hidden_states": prompt_context.negative_connector_prompt_embeds,
|
|
"audio_encoder_hidden_states": prompt_context.negative_connector_audio_prompt_embeds,
|
|
"encoder_attention_mask": prompt_context.negative_connector_attention_mask,
|
|
"audio_encoder_attention_mask": prompt_context.negative_connector_attention_mask,
|
|
}
|
|
return self.predict_parallel_cfg(
|
|
pipeline,
|
|
true_cfg_scale=guidance_scale,
|
|
positive_kwargs=positive_kwargs,
|
|
negative_kwargs=negative_kwargs,
|
|
cfg_normalize=False,
|
|
video_latents=state.video,
|
|
audio_latents=state.audio,
|
|
video_sigma=pipeline.scheduler.sigmas[index],
|
|
audio_sigma=audio_scheduler.sigmas[index],
|
|
)
|
|
|
|
video_input = torch.cat([state.video] * 2) if pipeline.do_classifier_free_guidance else state.video
|
|
video_input = video_input.to(prompt_context.connector_prompt_embeds.dtype)
|
|
audio_input = torch.cat([state.audio] * 2) if pipeline.do_classifier_free_guidance else state.audio
|
|
audio_input = audio_input.to(prompt_context.connector_prompt_embeds.dtype)
|
|
ts = timestep.expand(video_input.shape[0])
|
|
transformer_kwargs = pipeline._build_transformer_kwargs(
|
|
forward_ctx,
|
|
denoise_ctx,
|
|
hidden_states=video_input,
|
|
audio_hidden_states=audio_input,
|
|
encoder_hidden_states=prompt_context.connector_prompt_embeds,
|
|
audio_encoder_hidden_states=prompt_context.connector_audio_prompt_embeds,
|
|
encoder_attention_mask=prompt_context.connector_attention_mask,
|
|
audio_encoder_attention_mask=prompt_context.connector_attention_mask,
|
|
ts=ts,
|
|
)
|
|
with pipeline._transformer_cache_context("cond_uncond"):
|
|
video_pred, audio_pred = pipeline.transformer(**transformer_kwargs)
|
|
video_pred = video_pred.float()
|
|
audio_pred = audio_pred.float()
|
|
|
|
if pipeline.do_classifier_free_guidance:
|
|
video_uncond, video_cond = video_pred.chunk(2)
|
|
video_pred = combine_velocity_via_x0(
|
|
state.video,
|
|
video_cond,
|
|
video_uncond,
|
|
pipeline.scheduler.sigmas[index],
|
|
guidance_scale,
|
|
)
|
|
audio_uncond, audio_cond = audio_pred.chunk(2)
|
|
audio_pred = combine_velocity_via_x0(
|
|
state.audio,
|
|
audio_cond,
|
|
audio_uncond,
|
|
audio_scheduler.sigmas[index],
|
|
guidance_scale,
|
|
)
|
|
return video_pred, audio_pred
|
|
|
|
|
|
LTX_LEGACY_VELOCITY_GUIDANCE = LTXLegacyVelocityGuidance()
|
|
LTX_OFFICIAL_X0_GUIDANCE = LTXOfficialX0Guidance()
|