620 lines
23 KiB
Python
620 lines
23 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Shared runtime surface for LTX pipeline variants."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterable
|
|
from contextlib import nullcontext
|
|
from typing import Any, ClassVar
|
|
|
|
import torch
|
|
from diffusers.utils.torch_utils import randn_tensor
|
|
from torch import nn
|
|
from vllm.model_executor.models.utils import AutoWeightsLoader
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
|
from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
|
|
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
|
|
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
|
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
|
|
|
|
from . import ltx2_latents as latent_ops
|
|
from .ltx2_components import LTXComponentProfile, initialize_pipeline_components
|
|
from .ltx2_conditioning import LTXPromptContext, LTXTextConditioningMixin
|
|
from .ltx2_denoise import (
|
|
LTXDenoiseContext,
|
|
LTXForwardContext,
|
|
LTXPhaseExecutor,
|
|
LTXPhaseResult,
|
|
VideoAudioScheduler,
|
|
build_transformer_kwargs,
|
|
step_denoised_latents,
|
|
)
|
|
from .ltx2_guidance import LTXGuidanceStrategy
|
|
from .ltx2_recipes import LTXOneStageRecipe
|
|
from .ltx2_request import LTXRequestInputs, LTXRequestMixin
|
|
|
|
|
|
def _expand_per_prompt_decode_value(
|
|
value: float | list[float],
|
|
*,
|
|
prompt_batch_size: int,
|
|
effective_batch_size: int,
|
|
field_name: str,
|
|
) -> list[float]:
|
|
if not isinstance(value, list):
|
|
return [value] * effective_batch_size
|
|
if len(value) == 1:
|
|
return value * effective_batch_size
|
|
if len(value) == effective_batch_size:
|
|
return value
|
|
if prompt_batch_size > 0 and len(value) == prompt_batch_size and effective_batch_size % prompt_batch_size == 0:
|
|
repeats = effective_batch_size // prompt_batch_size
|
|
return [item for item in value for _ in range(repeats)]
|
|
raise ValueError(
|
|
f"`{field_name}` must have length 1, prompt batch size ({prompt_batch_size}), or effective batch size"
|
|
f" ({effective_batch_size}); got {len(value)}."
|
|
)
|
|
|
|
|
|
def _prepare_decode_timestep_conditioning(
|
|
*,
|
|
decode_timestep: float | list[float],
|
|
decode_noise_scale: float | list[float] | None,
|
|
prompt_batch_size: int,
|
|
effective_batch_size: int,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
decode_timestep_values = _expand_per_prompt_decode_value(
|
|
decode_timestep,
|
|
prompt_batch_size=prompt_batch_size,
|
|
effective_batch_size=effective_batch_size,
|
|
field_name="decode_timestep",
|
|
)
|
|
decode_noise_scale_values = (
|
|
decode_timestep_values
|
|
if decode_noise_scale is None
|
|
else _expand_per_prompt_decode_value(
|
|
decode_noise_scale,
|
|
prompt_batch_size=prompt_batch_size,
|
|
effective_batch_size=effective_batch_size,
|
|
field_name="decode_noise_scale",
|
|
)
|
|
)
|
|
return (
|
|
torch.tensor(decode_timestep_values, device=device, dtype=dtype),
|
|
torch.tensor(decode_noise_scale_values, device=device, dtype=dtype)[:, None, None, None, None],
|
|
)
|
|
|
|
|
|
class LTXPipelineRuntime(
|
|
LTXRequestMixin,
|
|
LTXTextConditioningMixin,
|
|
nn.Module,
|
|
CFGParallelMixin,
|
|
ProgressBarMixin,
|
|
SupportsComponentDiscovery,
|
|
DiffusionPipelineProfilerMixin,
|
|
):
|
|
"""Shared Omni runtime for explicitly composed LTX denoise phases."""
|
|
|
|
component_profile: ClassVar[LTXComponentProfile]
|
|
guidance_strategy: ClassVar[LTXGuidanceStrategy]
|
|
one_stage_recipe: ClassVar[LTXOneStageRecipe]
|
|
supports_request_batch = False
|
|
connector_batches_cfg = False
|
|
distributed_video_decode = True
|
|
support_image_input = False
|
|
dummy_run_num_frames = 2
|
|
preserve_sp_padded_audio_duration = False
|
|
reports_stage_durations = False
|
|
|
|
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
|
|
del prefix
|
|
super().__init__()
|
|
initialize_pipeline_components(self, od_config)
|
|
self.setup_diffusion_pipeline_profiler(
|
|
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
|
|
)
|
|
|
|
def prepare_latents(
|
|
self,
|
|
batch_size: int = 1,
|
|
num_channels_latents: int = 128,
|
|
height: int = 512,
|
|
width: int = 768,
|
|
num_frames: int = 121,
|
|
noise_scale: float = 0.0,
|
|
dtype: torch.dtype | None = None,
|
|
device: torch.device | None = None,
|
|
generator: torch.Generator | list[torch.Generator] | None = None,
|
|
latents: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
return latent_ops.prepare_video_latents(
|
|
self,
|
|
batch_size,
|
|
num_channels_latents,
|
|
height,
|
|
width,
|
|
num_frames,
|
|
noise_scale,
|
|
dtype,
|
|
device,
|
|
generator,
|
|
latents,
|
|
)
|
|
|
|
def prepare_audio_latents(
|
|
self,
|
|
batch_size: int = 1,
|
|
num_channels_latents: int = 8,
|
|
audio_latent_length: int = 1,
|
|
num_mel_bins: int = 64,
|
|
noise_scale: float = 0.0,
|
|
dtype: torch.dtype | None = None,
|
|
device: torch.device | None = None,
|
|
generator: torch.Generator | list[torch.Generator] | None = None,
|
|
latents: torch.Tensor | None = None,
|
|
) -> tuple[torch.Tensor, int, int]:
|
|
return latent_ops.prepare_audio_latents(
|
|
self,
|
|
batch_size,
|
|
num_channels_latents,
|
|
audio_latent_length,
|
|
num_mel_bins,
|
|
noise_scale,
|
|
dtype,
|
|
device,
|
|
generator,
|
|
latents,
|
|
)
|
|
|
|
@property
|
|
def guidance_scale(self):
|
|
return self._guidance_scale
|
|
|
|
@property
|
|
def guidance_rescale(self):
|
|
return self._guidance_rescale
|
|
|
|
@property
|
|
def do_classifier_free_guidance(self):
|
|
return self._guidance_scale is not None and self._guidance_scale > 1.0
|
|
|
|
@property
|
|
def interrupt(self):
|
|
return self._interrupt
|
|
|
|
def _transformer_cache_context(self, context_name: str):
|
|
cache_context = getattr(self.transformer, "cache_context", None)
|
|
if callable(cache_context):
|
|
return cache_context(context_name)
|
|
return nullcontext()
|
|
|
|
def predict_noise(self, **kwargs):
|
|
with self._transformer_cache_context("cond_uncond"):
|
|
noise_pred_video, noise_pred_audio = self.transformer(**kwargs)
|
|
return noise_pred_video.float(), noise_pred_audio.float()
|
|
|
|
def combine_cfg_noise(
|
|
self,
|
|
positive_noise_pred,
|
|
negative_noise_pred,
|
|
true_cfg_scale,
|
|
cfg_normalize=False,
|
|
kwargs: dict[str, Any] | None = None,
|
|
**context: Any,
|
|
):
|
|
if kwargs is not None:
|
|
context = {**kwargs, **context}
|
|
return self.guidance_strategy.combine_cfg_noise(
|
|
self,
|
|
positive_noise_pred,
|
|
negative_noise_pred,
|
|
true_cfg_scale,
|
|
cfg_normalize,
|
|
context,
|
|
)
|
|
|
|
def predict_noise_with_parallel_cfg(self, *args, **kwargs):
|
|
predict_parallel_cfg = getattr(self.guidance_strategy, "predict_parallel_cfg", None)
|
|
if predict_parallel_cfg is None:
|
|
raise NotImplementedError("The selected LTX guidance strategy does not implement parallel CFG.")
|
|
return predict_parallel_cfg(self, *args, **kwargs)
|
|
|
|
def _synchronize_cfg_parallel_step_output(
|
|
self,
|
|
latents: tuple[torch.Tensor, torch.Tensor],
|
|
do_true_cfg: bool,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1):
|
|
return latents
|
|
|
|
# CUDA async execution otherwise permits numerical drift to accumulate
|
|
# across CFG-parallel denoise steps.
|
|
latents = tuple(tensor.contiguous() for tensor in latents)
|
|
device = next((tensor.device for tensor in latents if tensor.is_cuda), None)
|
|
if device is not None:
|
|
torch.cuda.current_stream(device).synchronize()
|
|
return latents
|
|
|
|
def _setup_forward_runtime(
|
|
self,
|
|
request_inputs: LTXRequestInputs,
|
|
attention_kwargs: dict[str, Any] | None,
|
|
) -> bool:
|
|
self._guidance_scale = request_inputs.guidance_scale
|
|
self._guidance_rescale = request_inputs.guidance_rescale
|
|
del attention_kwargs
|
|
self._interrupt = False
|
|
cfg_world_size = get_classifier_free_guidance_world_size()
|
|
if self.do_classifier_free_guidance:
|
|
self.guidance_strategy.validate_cfg_world_size(cfg_world_size)
|
|
return self.do_classifier_free_guidance and cfg_world_size > 1
|
|
|
|
def _check_forward_inputs(
|
|
self,
|
|
request_inputs: LTXRequestInputs,
|
|
image: Any | None = None,
|
|
) -> None:
|
|
self.check_inputs(
|
|
prompt=request_inputs.prompt,
|
|
height=request_inputs.height,
|
|
width=request_inputs.width,
|
|
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,
|
|
)
|
|
|
|
def _resolve_request_image(
|
|
self,
|
|
req: DiffusionRequestBatch,
|
|
image: Any | None,
|
|
request_inputs: LTXRequestInputs,
|
|
) -> Any | None:
|
|
del req, request_inputs
|
|
return image
|
|
|
|
def _make_output(self, output: tuple[torch.Tensor, torch.Tensor]) -> DiffusionOutput:
|
|
if self.reports_stage_durations:
|
|
return DiffusionOutput(
|
|
output=output,
|
|
stage_durations=getattr(self, "stage_durations", None),
|
|
)
|
|
return DiffusionOutput(output=output)
|
|
|
|
def _decode_output(
|
|
self,
|
|
*,
|
|
latents: torch.Tensor,
|
|
audio_latents: torch.Tensor,
|
|
output_type: str,
|
|
connector_prompt_embeds: torch.Tensor,
|
|
generator: torch.Generator | list[torch.Generator] | None,
|
|
device: torch.device,
|
|
decode_timestep: float | list[float],
|
|
decode_noise_scale: float | list[float] | None,
|
|
prompt_batch_size: int,
|
|
) -> DiffusionOutput:
|
|
if output_type == "latent":
|
|
return self._make_output((latents, audio_latents))
|
|
|
|
latents = latents.to(connector_prompt_embeds.dtype)
|
|
if not self.vae.config.timestep_conditioning:
|
|
timestep_decode = None
|
|
else:
|
|
noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype)
|
|
timestep_decode, decode_noise_scale_t = _prepare_decode_timestep_conditioning(
|
|
decode_timestep=decode_timestep,
|
|
decode_noise_scale=decode_noise_scale,
|
|
prompt_batch_size=prompt_batch_size,
|
|
effective_batch_size=latents.shape[0],
|
|
device=device,
|
|
dtype=latents.dtype,
|
|
)
|
|
latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise
|
|
|
|
dist_initialized = torch.distributed.is_initialized()
|
|
is_output_rank = not dist_initialized or torch.distributed.get_rank() == 0
|
|
vae_decode_needs_all_ranks = False
|
|
is_distributed_vae_enabled = getattr(self.vae, "is_distributed_enabled", None)
|
|
if self.distributed_video_decode and dist_initialized and callable(is_distributed_vae_enabled):
|
|
try:
|
|
# Distributed tiled decode is collective, so every rank must enter it.
|
|
vae_decode_needs_all_ranks = bool(is_distributed_vae_enabled())
|
|
except Exception:
|
|
pass
|
|
|
|
should_decode_video = not self.distributed_video_decode or is_output_rank or vae_decode_needs_all_ranks
|
|
if should_decode_video:
|
|
video = self.vae.decode(latents.to(self.vae.dtype), timestep_decode, return_dict=False)[0]
|
|
else:
|
|
video = torch.empty(0, device=latents.device, dtype=latents.dtype)
|
|
|
|
if self.distributed_video_decode and not is_output_rank:
|
|
return self._make_output(
|
|
(
|
|
torch.empty(0, device=video.device, dtype=video.dtype),
|
|
torch.empty(0, device=audio_latents.device, dtype=audio_latents.dtype),
|
|
)
|
|
)
|
|
|
|
if video.numel() > 0:
|
|
video = self.video_processor.postprocess_video(video, output_type=output_type)
|
|
generated_mel = self.audio_vae.decode(audio_latents.to(self.audio_vae.dtype), return_dict=False)[0]
|
|
audio = self.vocoder(generated_mel)
|
|
return self._make_output((video, audio))
|
|
|
|
def _prepare_video_latents_stage(
|
|
self,
|
|
request_inputs: LTXRequestInputs,
|
|
prompt_context: LTXPromptContext,
|
|
*,
|
|
device: torch.device,
|
|
noise_scale: float,
|
|
image: Any | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|
latents = self.prepare_latents(
|
|
prompt_context.batch_size * request_inputs.num_videos_per_prompt,
|
|
self.transformer.config.in_channels,
|
|
request_inputs.height,
|
|
request_inputs.width,
|
|
request_inputs.num_frames,
|
|
noise_scale,
|
|
torch.float32,
|
|
device,
|
|
request_inputs.generator,
|
|
request_inputs.latents,
|
|
)
|
|
return latents, None
|
|
|
|
def _resolve_video_latent_dimensions(self, request_inputs: LTXRequestInputs) -> tuple[int, int, int]:
|
|
latent_num_frames, latent_height, latent_width = latent_ops.resolve_video_latent_shape(
|
|
request_inputs.height,
|
|
request_inputs.width,
|
|
request_inputs.num_frames,
|
|
vae_spatial_compression_ratio=self.vae_spatial_compression_ratio,
|
|
vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
|
|
)
|
|
latents = request_inputs.latents
|
|
if latents is not None:
|
|
if latents.ndim == 5:
|
|
_, _, latent_num_frames, latent_height, latent_width = latents.shape
|
|
elif latents.ndim != 3:
|
|
raise ValueError(
|
|
f"Provided `latents` tensor has shape {latents.shape}, expected a packed 3D or unpacked 5D tensor."
|
|
)
|
|
return latent_num_frames, latent_height, latent_width
|
|
|
|
def _prepare_audio_latents_stage(
|
|
self,
|
|
request_inputs: LTXRequestInputs,
|
|
prompt_context: LTXPromptContext,
|
|
*,
|
|
device: torch.device,
|
|
noise_scale: float,
|
|
) -> tuple[torch.Tensor, int, int, int]:
|
|
duration_s = request_inputs.num_frames / request_inputs.frame_rate
|
|
audio_latents_per_second = (
|
|
self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio)
|
|
)
|
|
audio_num_frames = round(duration_s * audio_latents_per_second)
|
|
audio_num_frames = self._resolve_audio_latent_length(audio_num_frames, request_inputs.audio_latents)
|
|
|
|
num_mel_bins = self.audio_vae.config.mel_bins if self.audio_vae is not None else 64
|
|
latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
|
|
num_channels = self.audio_vae.config.latent_channels if self.audio_vae is not None else 8
|
|
audio_latents, original_num_frames, padded_num_frames = self.prepare_audio_latents(
|
|
prompt_context.batch_size * request_inputs.num_videos_per_prompt,
|
|
num_channels_latents=num_channels,
|
|
audio_latent_length=audio_num_frames,
|
|
num_mel_bins=num_mel_bins,
|
|
noise_scale=noise_scale,
|
|
dtype=torch.float32,
|
|
device=device,
|
|
generator=request_inputs.generator,
|
|
latents=request_inputs.audio_latents,
|
|
)
|
|
return audio_latents, original_num_frames, padded_num_frames, latent_mel_bins
|
|
|
|
def _resolve_audio_latent_length(
|
|
self,
|
|
requested_length: int,
|
|
audio_latents: torch.Tensor | None,
|
|
) -> int:
|
|
if audio_latents is None or audio_latents.ndim != 4:
|
|
return requested_length
|
|
|
|
provided_length = audio_latents.shape[2]
|
|
if not self.preserve_sp_padded_audio_duration:
|
|
return provided_length
|
|
|
|
sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) or 1
|
|
padded_length = latent_ops.get_sp_padded_audio_latent_length(requested_length, int(sp_size))
|
|
return requested_length if provided_length in {requested_length, padded_length} else provided_length
|
|
|
|
def _make_video_audio_scheduler(
|
|
self,
|
|
audio_scheduler: Any,
|
|
latent_num_frames: int,
|
|
latent_height: int,
|
|
latent_width: int,
|
|
) -> Any:
|
|
return VideoAudioScheduler(self.scheduler, audio_scheduler)
|
|
|
|
def _prepare_denoise_context_for_cfg(
|
|
self,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> LTXDenoiseContext:
|
|
return self.guidance_strategy.prepare_denoise_context(self, forward_ctx, denoise_ctx)
|
|
|
|
def _denoise_timestep_kwargs(
|
|
self,
|
|
ts: torch.Tensor,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> dict[str, torch.Tensor]:
|
|
return self.guidance_strategy.timestep_kwargs(ts, forward_ctx, denoise_ctx)
|
|
|
|
def _build_transformer_kwargs(
|
|
self,
|
|
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 build_transformer_kwargs(
|
|
self,
|
|
forward_ctx,
|
|
denoise_ctx,
|
|
hidden_states=hidden_states,
|
|
audio_hidden_states=audio_hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
audio_encoder_hidden_states=audio_encoder_hidden_states,
|
|
encoder_attention_mask=encoder_attention_mask,
|
|
audio_encoder_attention_mask=audio_encoder_attention_mask,
|
|
ts=ts,
|
|
)
|
|
|
|
def _predict_noise_for_step(
|
|
self,
|
|
index: int,
|
|
timestep: torch.Tensor,
|
|
state: latent_ops.LTXAVState,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
return self.guidance_strategy.predict_noise(
|
|
self,
|
|
index,
|
|
timestep,
|
|
state,
|
|
forward_ctx,
|
|
denoise_ctx,
|
|
)
|
|
|
|
def _denoise_step(
|
|
self,
|
|
index: int,
|
|
timestep: torch.Tensor,
|
|
state: latent_ops.LTXAVState,
|
|
forward_ctx: LTXForwardContext,
|
|
denoise_ctx: LTXDenoiseContext,
|
|
) -> latent_ops.LTXAVState:
|
|
denoise_ctx.latents = state.video
|
|
denoise_ctx.audio_latents = state.audio
|
|
noise_pred_video, noise_pred_audio = self._predict_noise_for_step(
|
|
index,
|
|
timestep,
|
|
state,
|
|
forward_ctx,
|
|
denoise_ctx,
|
|
)
|
|
video, audio = step_denoised_latents(
|
|
self,
|
|
forward_ctx,
|
|
denoise_ctx,
|
|
noise_pred_video,
|
|
noise_pred_audio,
|
|
timestep,
|
|
)
|
|
return latent_ops.LTXAVState(video=video, audio=audio)
|
|
|
|
def _unpack_and_denormalize_stage(
|
|
self,
|
|
forward_ctx: LTXForwardContext,
|
|
latents: torch.Tensor,
|
|
audio_latents: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
latents = latent_ops.unpack_latents(
|
|
latents,
|
|
forward_ctx.latent_num_frames,
|
|
forward_ctx.latent_height,
|
|
forward_ctx.latent_width,
|
|
self.transformer_spatial_patch_size,
|
|
self.transformer_temporal_patch_size,
|
|
)
|
|
latents = latent_ops.denormalize_latents(
|
|
latents,
|
|
self.vae.latents_mean,
|
|
self.vae.latents_std,
|
|
self.vae.config.scaling_factor,
|
|
)
|
|
|
|
audio_latents = latent_ops.unpad_audio_latents(audio_latents, forward_ctx.original_audio_num_frames)
|
|
audio_latents = latent_ops.denormalize_audio_latents(
|
|
audio_latents,
|
|
self.audio_vae.latents_mean,
|
|
self.audio_vae.latents_std,
|
|
)
|
|
audio_latents = latent_ops.unpack_audio_latents(
|
|
audio_latents,
|
|
forward_ctx.original_audio_num_frames,
|
|
num_mel_bins=forward_ctx.latent_mel_bins,
|
|
)
|
|
return latents, audio_latents
|
|
|
|
def run_phase(
|
|
self,
|
|
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:
|
|
"""Prepare and execute one phase without decoding its output."""
|
|
return LTXPhaseExecutor.run(
|
|
self,
|
|
req,
|
|
request_inputs,
|
|
noise_scale=noise_scale,
|
|
sigmas=sigmas,
|
|
timesteps=timesteps,
|
|
attention_kwargs=attention_kwargs,
|
|
image=image,
|
|
prompt_context=prompt_context,
|
|
)
|
|
|
|
def decode_phase(self, phase: LTXPhaseResult) -> DiffusionOutput | list[DiffusionOutput]:
|
|
"""Decode one completed phase and restore per-request outputs."""
|
|
forward_ctx = phase.forward_context
|
|
request_inputs = forward_ctx.request_inputs
|
|
output = self._decode_output(
|
|
latents=phase.video,
|
|
audio_latents=phase.audio,
|
|
output_type=request_inputs.output_type,
|
|
connector_prompt_embeds=forward_ctx.prompt_context.connector_prompt_embeds,
|
|
generator=request_inputs.generator,
|
|
device=forward_ctx.device,
|
|
decode_timestep=request_inputs.decode_timestep,
|
|
decode_noise_scale=request_inputs.decode_noise_scale,
|
|
prompt_batch_size=forward_ctx.batch_size,
|
|
)
|
|
if not self.supports_request_batch:
|
|
return output
|
|
return split_diffusion_output_by_request(
|
|
output,
|
|
forward_ctx.req,
|
|
num_outputs_per_prompt=forward_ctx.num_videos_per_prompt,
|
|
)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
return AutoWeightsLoader(self).load_weights(weights)
|