anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Model] Add two stages inference for model LTX-2 distilled. (#2260)

Browse Source 
Signed-off-by: Songrui625 <songrui625@gmail.com>
This commit is contained in:
Jerry Song authored and GitHub committed 2026-04-03 23:17:49 +08:00
1 parent 6dc61c9a20
commit cd71567b06
10 files changed
+1120 -303
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/models/supported_models.md
+2
View File
@@ -34,6 +34,8 @@ th {
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
examples/offline_inference/image_to_video/image_to_video.py
+1 -5
View File
@@ -490,10 +490,6 @@ def main():
if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
frames_np = frames_np[..., :3]
frames_np = np.clip(frames_np, 0.0, 1.0)
frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
video_tensor = torch.from_numpy(frames_u8)
audio_out = None
if audio is not None:
if isinstance(audio, list):
@@ -507,7 +503,7 @@ def main():
audio_out = audio_out.float().cpu()
encode_video(
video_tensor,
frames_np,
fps=fps,
audio=audio_out,
audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
examples/offline_inference/text_to_video/text_to_video.py
+8 -1
View File
@@ -56,8 +56,13 @@ def parse_args() -> argparse.Namespace:
"Examples: Wan-AI/Wan2.2-T2V-A14B-Diffusers, "
"hunyuanvideo-community/HunyuanVideo-1.5-480p_t2v",
)
parser.add_argument(
"--model-class-name",
default=None,
help="Override model class name (e.g., LTX2TwoStagesVideoPipeline).",
)
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
parser.add_argument("--negative-prompt", default="", help="Negative prompt (Wan2.2 only).")
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
parser.add_argument(
@@ -185,6 +190,7 @@ def parse_args() -> argparse.Namespace:
def main():
args = parse_args()
model_class_name = args.model_class_name
preset = _detect_preset(args.model)
for key, default_val in preset.items():
@@ -229,6 +235,7 @@ def main():
enable_cpu_offload=args.enable_cpu_offload,
parallel_config=parallel_config,
enforce_eager=args.enforce_eager,
model_class_name=model_class_name,
cache_backend=args.cache_backend,
cache_config=cache_config,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
vllm_omni/diffusion/cache/cache_dit_backend.py
Vendored
+12 -6
View File
@@ -24,6 +24,7 @@ from vllm.logger import init_logger
from vllm_omni.diffusion.cache.base import CacheBackend
from vllm_omni.diffusion.data import DiffusionCacheConfig, OmniDiffusionConfig
from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline
logger = init_logger(__name__)
@@ -533,7 +534,7 @@ def enable_cache_for_sd3(pipeline: Any, cache_config: Any) -> Callable[[int], No
def enable_cache_for_ltx2(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
"""Enable cache-dit for LTX2 pipelines (audio-video transformer blocks)."""
transformer = pipeline.transformer
transformer = get_transformer_from_pipeline(pipeline)
db_cache_config = _build_db_cache_config(cache_config)
@@ -566,11 +567,12 @@ def enable_cache_for_ltx2(pipeline: Any, cache_config: Any) -> Callable[[int], N
)
def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
transformer = get_transformer_from_pipeline(pipeline)
if cache_config.scm_steps_mask_policy is None:
cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
else:
cache_dit.refresh_context(
pipeline.transformer,
transformer,
cache_config=DBCacheConfig().reset(
num_inference_steps=num_inference_steps,
steps_computation_mask=cache_dit.steps_mask(
@@ -613,8 +615,9 @@ def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], No
)
# Enable cache-dit on the transformer
transformer = get_transformer_from_pipeline(pipeline)
cache_dit.enable_cache(
pipeline.transformer,
transformer,
cache_config=db_cache_config,
calibrator_config=calibrator_config,
)
@@ -626,11 +629,12 @@ def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], No
pipeline: The diffusion pipeline instance.
num_inference_steps: New number of inference steps.
"""
transformer = get_transformer_from_pipeline(pipeline)
if cache_config.scm_steps_mask_policy is None:
cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
else:
cache_dit.refresh_context(
pipeline.transformer,
transformer,
cache_config=DBCacheConfig().reset(
num_inference_steps=num_inference_steps,
steps_computation_mask=cache_dit.steps_mask(
@@ -1211,6 +1215,8 @@ CUSTOM_DIT_ENABLERS.update(
"StableDiffusion3Pipeline": enable_cache_for_sd3,
"LTX2Pipeline": enable_cache_for_ltx2,
"LTX2ImageToVideoPipeline": enable_cache_for_ltx2,
"LTX2TwoStagesPipeline": enable_cache_for_ltx2,
"LTX2ImageToVideoTwoStagesPipeline": enable_cache_for_ltx2,
"BagelPipeline": enable_cache_for_bagel,
"GlmImagePipeline": enable_cache_for_glm_image,
"Flux2Pipeline": enable_cache_for_flux2,
vllm_omni/diffusion/models/ltx2/__init__.py
+9 -1
View File
@@ -4,15 +4,23 @@
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2Pipeline,
LTX2TwoStagesPipeline,
create_transformer_from_config,
get_ltx2_post_process_func,
load_transformer_config,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import (
LTX2ImageToVideoPipeline,
LTX2ImageToVideoTwoStagesPipeline,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
__all__ = [
"LTX2Pipeline",
"LTX2ImageToVideoPipeline",
"LTX2LatentUpsamplePipeline",
"LTX2TwoStagesPipeline",
"LTX2ImageToVideoTwoStagesPipeline",
"get_ltx2_post_process_func",
"load_transformer_config",
"create_transformer_from_config",
vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py
+397 -147
View File
@@ -15,12 +15,14 @@ import numpy as np
import torch
from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.ltx2 import LTX2TextConnectors
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg, retrieve_timesteps
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
from torch import nn
from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
from vllm.logger import init_logger
from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -31,10 +33,16 @@ from vllm_omni.diffusion.distributed.parallel_state import (
get_classifier_free_guidance_world_size,
)
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.lora.request import LoRARequest
from .ltx2_transformer import LTX2VideoTransformer3DModel
from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
logger = init_logger(__name__)
def load_transformer_config(model_path: str, subfolder: str = "transformer", local_files_only: bool = True) -> dict:
@@ -114,7 +122,7 @@ def calculate_shift(
return mu
class LTX2Pipeline(nn.Module, CFGParallelMixin):
class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
def __init__(
self,
*,
@@ -145,12 +153,15 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin):
subfolder="tokenizer",
local_files_only=local_files_only,
)
self.text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
model,
subfolder="text_encoder",
torch_dtype=dtype,
local_files_only=local_files_only,
).to(self.device)
# prefer mmap loading as default device is cuda, and the output of text encoder
# could be deterministic.
with torch.device("cpu"):
self.text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
model,
subfolder="text_encoder",
torch_dtype=dtype,
local_files_only=local_files_only,
).to(self.device)
self.connectors = LTX2TextConnectors.from_pretrained(
model,
subfolder="connectors",
@@ -460,6 +471,22 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin):
latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
return latents
@staticmethod
def _normalize_latents(
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
) -> torch.Tensor:
# Normalize latents across the channel dimension [B, C, F, H, W]
latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents = (latents - latents_mean) * scaling_factor / latents_std
return latents
@staticmethod
def _normalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor):
latents_mean = latents_mean.to(latents.device, latents.dtype)
latents_std = latents_std.to(latents.device, latents.dtype)
return (latents - latents_mean) / latents_std
@staticmethod
def _denormalize_latents(
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
@@ -475,6 +502,14 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin):
latents_std = latents_std.to(latents.device, latents.dtype)
return (latents * latents_std) + latents_mean
@staticmethod
def _create_noised_state(
latents: torch.Tensor, noise_scale: float | torch.Tensor, generator: torch.Generator | None = None
):
noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype)
noised_latents = noise_scale * noise + (1 - noise_scale) * latents
return noised_latents
@staticmethod
def _pack_audio_latents(
latents: torch.Tensor, patch_size: int | None = None, patch_size_t: int | None = None
@@ -514,12 +549,26 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin):
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 | None = None,
latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
if latents.ndim == 5:
latents = self._normalize_latents(
latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
)
# latents are of shape [B, C, F, H, W], need to be packed
latents = self._pack_latents(
latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
)
if latents.ndim != 3:
raise ValueError(
f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is [batch_size, num_seq, num_features]." # noqa
)
latents = self._create_noised_state(latents, noise_scale, generator)
return latents.to(device=device, dtype=dtype)
height = height // self.vae_spatial_compression_ratio
@@ -543,37 +592,30 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin):
self,
batch_size: int = 1,
num_channels_latents: int = 8,
audio_latent_length: int = 1, # 1 is just a dummy value
num_mel_bins: int = 64,
num_frames: int = 121,
frame_rate: float = 25.0,
sampling_rate: int = 16000,
hop_length: int = 160,
noise_scale: float = 0.0,
dtype: torch.dtype | None = None,
device: torch.device | None = None,
generator: torch.Generator | None = None,
latents: torch.Tensor | None = None,
) -> tuple[torch.Tensor, int]:
duration_s = num_frames / frame_rate
latents_per_second = float(sampling_rate) / float(hop_length) / float(self.audio_vae_temporal_compression_ratio)
latent_length = round(duration_s * latents_per_second)
if latents is not None:
if latents.ndim == 4:
# latents are of shape [B, C, L, M], need to be packed
latents = self._pack_audio_latents(latents)
if latents.ndim != 3:
raise ValueError(
f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is [batch_size, num_seq, num_features]." # noqa
)
latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std)
latents = self._create_noised_state(latents, noise_scale, generator)
return latents.to(device=device, dtype=dtype)
# TODO: confirm whether this logic is correct
latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1)
if sp_size > 1:
pad_len = (sp_size - (latent_length % sp_size)) % sp_size
if pad_len > 0:
if latents is not None:
pad_shape = list(latents.shape)
pad_shape[2] = pad_len
padding = torch.zeros(pad_shape, dtype=latents.dtype, device=latents.device)
latents = torch.cat([latents, padding], dim=2)
latent_length += pad_len
if latents is not None:
return latents.to(device=device, dtype=dtype), latent_length
shape = (batch_size, num_channels_latents, latent_length, latent_mel_bins)
shape = (batch_size, num_channels_latents, audio_latent_length, latent_mel_bins)
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
@@ -583,7 +625,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin):
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
latents = self._pack_audio_latents(latents)
return latents, latent_length
return latents
@property
def guidance_scale(self):
@@ -766,9 +808,11 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin):
num_frames: int | None = None,
frame_rate: float | None = None,
num_inference_steps: int | None = None,
sigmas: list[float] | None = None,
timesteps: list[int] | None = None,
guidance_scale: float = 4.0,
guidance_rescale: float = 0.0,
noise_scale: float = 0.0,
num_videos_per_prompt: int | None = 1,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
@@ -925,6 +969,21 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin):
latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
latent_height = height // self.vae_spatial_compression_ratio
latent_width = width // self.vae_spatial_compression_ratio
if latents is not None:
if latents.ndim == 5:
logger.info(
"Got latents of shape [batch_size, latent_dim, latent_frames, latent_height, latent_width], `latent_num_frames`, `latent_height`, `latent_width` will be inferred." # noqa
)
_, _, latent_num_frames, latent_height, latent_width = latents.shape # [B, C, F, H, W]
elif latents.ndim == 3:
logger.warning(
f"You have supplied packed `latents` of shape {latents.shape}, so the latent dims cannot be"
f" inferred. Make sure the supplied `height`, `width`, and `num_frames` are correct."
)
else:
raise ValueError(
f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, latent_dim, latent_frames, latent_height, latent_width]." # noqa
)
video_sequence_length = latent_num_frames * latent_height * latent_width
num_channels_latents = self.transformer.config.in_channels
@@ -934,33 +993,66 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin):
height,
width,
num_frames,
noise_scale,
torch.float32,
device,
generator,
latents,
)
duration_s = num_frames / 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)
if audio_latents is not None:
if audio_latents.ndim == 4:
logger.info(
"Got audio_latents of shape [batch_size, num_channels, audio_length, mel_bins], `audio_num_frames` will be inferred." # noqa
)
_, _, audio_num_frames, _ = audio_latents.shape # [B, C, L, M]
elif audio_latents.ndim == 3:
logger.warning(
f"You have supplied packed `audio_latents` of shape {audio_latents.shape}, so the latent dims"
f" cannot be inferred. Make sure the supplied `num_frames` and `frame_rate` are correct."
)
else:
raise ValueError(
f"Provided `audio_latents` tensor has shape {audio_latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, num_channels, audio_length, mel_bins]." # noqa
)
num_mel_bins = self.audio_vae.config.mel_bins if getattr(self, "audio_vae", None) is not None else 64
latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
num_channels_latents_audio = (
self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8
)
audio_latents, audio_num_frames = self.prepare_audio_latents(
# padding audio_latents if needed
sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1)
if sp_size > 1:
pad_len = (sp_size - (audio_num_frames % sp_size)) % sp_size
if pad_len > 0:
if audio_latents is not None:
pad_shape = list(audio_latents.shape)
pad_shape[2] = pad_len
padding = torch.zeros(pad_shape, dtype=audio_latents.dtype, device=audio_latents.device)
audio_latents = torch.cat([audio_latents, padding], dim=2)
audio_num_frames += pad_len
audio_latents = self.prepare_audio_latents(
batch_size * num_videos_per_prompt,
num_channels_latents=num_channels_latents_audio,
audio_latent_length=audio_num_frames,
num_mel_bins=num_mel_bins,
num_frames=num_frames,
frame_rate=frame_rate,
sampling_rate=self.audio_sampling_rate,
hop_length=self.audio_hop_length,
noise_scale=noise_scale,
dtype=torch.float32,
device=device,
generator=generator,
latents=audio_latents,
)
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
mu = calculate_shift(
video_sequence_length,
self.scheduler.config.get("base_image_seq_len", 1024),
@@ -985,7 +1077,6 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin):
sigmas=sigmas,
mu=mu,
)
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
self._num_timesteps = len(timesteps)
video_coords = self.transformer.rope.prepare_video_coords(
@@ -994,129 +1085,133 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin):
audio_coords = self.transformer.audio_rope.prepare_audio_coords(
audio_latents.shape[0], audio_num_frames, audio_latents.device
)
# Duplicate the positional ids as well if using CFG
if self.do_classifier_free_guidance and not cfg_parallel_ready:
video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) # Repeat twice in batch dim
audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
for i, t in enumerate(timesteps):
if self.interrupt:
continue
with self.progress_bar(total=len(timesteps)) as pbar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
self._current_timestep = t
self._current_timestep = t
if cfg_parallel_ready:
latent_model_input = latents.to(prompt_embeds.dtype)
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
timestep = t.expand(latent_model_input.shape[0])
if cfg_parallel_ready:
latent_model_input = latents.to(prompt_embeds.dtype)
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
timestep = t.expand(latent_model_input.shape[0])
positive_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": connector_prompt_embeds,
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
"timestep": timestep,
"encoder_attention_mask": connector_attention_mask,
"audio_encoder_attention_mask": connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
negative_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": negative_connector_prompt_embeds,
"audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
"timestep": timestep,
"encoder_attention_mask": negative_connector_attention_mask,
"audio_encoder_attention_mask": negative_connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
positive_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": connector_prompt_embeds,
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
"timestep": timestep,
"encoder_attention_mask": connector_attention_mask,
"audio_encoder_attention_mask": connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
negative_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": negative_connector_prompt_embeds,
"audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
"timestep": timestep,
"encoder_attention_mask": negative_connector_attention_mask,
"audio_encoder_attention_mask": negative_connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
noise_pred_video, noise_pred_audio = self.predict_noise_av_maybe_with_cfg(
do_true_cfg=True,
true_cfg_scale=guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
guidance_rescale=guidance_rescale,
cfg_normalize=False,
)
latents, audio_latents = self._scheduler_step_video_audio_maybe_with_cfg(
noise_pred_video,
noise_pred_audio,
t,
latents,
audio_latents,
audio_scheduler,
do_true_cfg=True,
)
else:
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
latent_model_input = latent_model_input.to(prompt_embeds.dtype)
audio_latent_model_input = (
torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents
)
audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype)
timestep = t.expand(latent_model_input.shape[0])
with self._transformer_cache_context("cond_uncond"):
noise_pred_video, noise_pred_audio = self.transformer(
hidden_states=latent_model_input,
audio_hidden_states=audio_latent_model_input,
encoder_hidden_states=connector_prompt_embeds,
audio_encoder_hidden_states=connector_audio_prompt_embeds,
timestep=timestep,
encoder_attention_mask=connector_attention_mask,
audio_encoder_attention_mask=connector_attention_mask,
num_frames=latent_num_frames,
height=latent_height,
width=latent_width,
fps=frame_rate,
audio_num_frames=audio_num_frames,
video_coords=video_coords,
audio_coords=audio_coords,
attention_kwargs=attention_kwargs,
return_dict=False,
)
noise_pred_video = noise_pred_video.float()
noise_pred_audio = noise_pred_audio.float()
if self.do_classifier_free_guidance:
noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
noise_pred_video = noise_pred_video_uncond + guidance_scale * (
noise_pred_video_text - noise_pred_video_uncond
noise_pred_video, noise_pred_audio = self.predict_noise_av_maybe_with_cfg(
do_true_cfg=True,
true_cfg_scale=guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
guidance_rescale=guidance_rescale,
cfg_normalize=False,
)
noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (
noise_pred_audio_text - noise_pred_audio_uncond
latents, audio_latents = self._scheduler_step_video_audio_maybe_with_cfg(
noise_pred_video,
noise_pred_audio,
t,
latents,
audio_latents,
audio_scheduler,
do_true_cfg=True,
)
else:
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
latent_model_input = latent_model_input.to(prompt_embeds.dtype)
audio_latent_model_input = (
torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents
)
audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype)
if guidance_rescale > 0:
noise_pred_video = rescale_noise_cfg(
noise_pred_video, noise_pred_video_text, guidance_rescale=guidance_rescale
timestep = t.expand(latent_model_input.shape[0])
with self._transformer_cache_context("cond_uncond"):
noise_pred_video, noise_pred_audio = self.transformer(
hidden_states=latent_model_input,
audio_hidden_states=audio_latent_model_input,
encoder_hidden_states=connector_prompt_embeds,
audio_encoder_hidden_states=connector_audio_prompt_embeds,
timestep=timestep,
encoder_attention_mask=connector_attention_mask,
audio_encoder_attention_mask=connector_attention_mask,
num_frames=latent_num_frames,
height=latent_height,
width=latent_width,
fps=frame_rate,
audio_num_frames=audio_num_frames,
video_coords=video_coords,
audio_coords=audio_coords,
attention_kwargs=attention_kwargs,
return_dict=False,
)
noise_pred_audio = rescale_noise_cfg(
noise_pred_audio, noise_pred_audio_text, guidance_rescale=guidance_rescale
noise_pred_video = noise_pred_video.float()
noise_pred_audio = noise_pred_audio.float()
if self.do_classifier_free_guidance:
noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
noise_pred_video = noise_pred_video_uncond + guidance_scale * (
noise_pred_video_text - noise_pred_video_uncond
)
latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (
noise_pred_audio_text - noise_pred_audio_uncond
)
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
pass
if guidance_rescale > 0:
noise_pred_video = rescale_noise_cfg(
noise_pred_video, noise_pred_video_text, guidance_rescale=guidance_rescale
)
noise_pred_audio = rescale_noise_cfg(
noise_pred_audio, noise_pred_audio_text, guidance_rescale=guidance_rescale
)
latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
pbar.update()
latents = self._unpack_latents(
latents,
@@ -1174,3 +1269,158 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self)
return loader.load_weights(weights)
class LTX2TwoStagesPipeline(nn.Module):
"""LTX2TwoStagesPipeline is for two stages image to video generation"""
def __init__(
self,
*,
od_config: OmniDiffusionConfig,
prefix: str = "",
):
super().__init__()
self.device = get_local_device()
self.dtype = getattr(od_config, "dtype", torch.bfloat16)
self.model_path = od_config.model
self.distilled = False
# User provided model path may contain '/' in the end and basename function
# will not return the expected directory name, so we need to remove it by normpath
if "distilled" in os.path.basename(os.path.normpath(self.model_path)):
self.distilled = True
else:
raise NotImplementedError(f"{self.model_path} is not supported for {self.__class__.__name__}.")
self.pipe = LTX2Pipeline(od_config=od_config, prefix=prefix)
self.upsample_pipe = LTX2LatentUpsamplePipeline(
vae=self.pipe.vae,
od_config=od_config,
)
self.lora_manager = DiffusionLoRAManager(
pipeline=self.pipe,
device=self.device,
dtype=self.dtype,
max_cached_adapters=od_config.max_cpu_loras,
)
self.weights_sources = [
DiffusersPipelineLoader.ComponentSource(
model_or_path=od_config.model,
subfolder="transformer",
revision=None,
prefix="pipe.transformer.",
fall_back_to_pt=True,
),
]
def forward(
self,
req: OmniDiffusionRequest,
prompt: str | list[str] | None = None,
negative_prompt: str | list[str] | None = None,
height: int | None = None,
width: int | None = None,
num_frames: int | None = None,
frame_rate: float | None = None,
num_inference_steps: int | None = None,
timesteps: list[int] | None = None,
guidance_scale: float = 4.0,
guidance_rescale: float = 0.0,
noise_scale: float = 0.0,
num_videos_per_prompt: int | None = 1,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
audio_latents: torch.Tensor | None = None,
prompt_embeds: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
prompt_attention_mask: torch.Tensor | None = None,
negative_prompt_attention_mask: torch.Tensor | None = None,
decode_timestep: float | list[float] = 0.0,
decode_noise_scale: float | list[float] | None = None,
output_type: str = "np",
return_dict: bool = True,
attention_kwargs: dict[str, Any] | None = None,
max_sequence_length: int | None = None,
):
video_latent, audio_latent = self.pipe(
req=req,
prompt=prompt,
negative_prompt=negative_prompt,
height=height,
width=width,
num_frames=num_frames,
frame_rate=frame_rate,
num_inference_steps=num_inference_steps,
sigmas=DISTILLED_SIGMA_VALUES if self.distilled else None,
timesteps=timesteps,
guidance_scale=guidance_scale,
guidance_rescale=guidance_rescale,
noise_scale=noise_scale,
num_videos_per_prompt=num_videos_per_prompt,
generator=generator,
latents=latents,
audio_latents=audio_latents,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
decode_timestep=decode_timestep,
decode_noise_scale=decode_noise_scale,
output_type="latent",
return_dict=return_dict,
attention_kwargs=attention_kwargs,
max_sequence_length=max_sequence_length,
).output
upscaled_video_latent = self.upsample_pipe(
latents=video_latent,
output_type="latent",
return_dict=False,
)[0]
if not self.distilled:
# Load Stage 2 distilled LoRA
lora_path = f"{self.model_path}/ltx-2-19b-distilled-lora-384.safetensors"
lora_request = LoRARequest(
lora_name="stage_2_distilled",
lora_int_id=1,
lora_path=lora_path,
)
self.lora_manager.set_active_adapter(lora_request, lora_scale=1.0)
# Change scheduler to use Stage 2 distilled sigmas as is
new_scheduler = FlowMatchEulerDiscreteScheduler.from_config(
self.pipe.scheduler.config,
use_dynamic_shifting=False,
shift_terminal=None,
)
self.pipe.scheduler = new_scheduler
# We only want to change num_inference_steps here, so no need
# to deep copy the whole request
stage_2_req = copy.copy(req)
stage_2_req.sampling_params = req.sampling_params.clone()
stage_2_req.sampling_params.num_inference_steps = 3
video, audio = self.pipe(
req=stage_2_req,
latents=upscaled_video_latent,
audio_latents=audio_latent,
prompt=prompt,
negative_prompt=negative_prompt,
noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
sigmas=STAGE_2_DISTILLED_SIGMA_VALUES,
guidance_scale=1.0,
generator=generator,
output_type="np",
return_dict=False,
).output
return DiffusionOutput(output=(video, audio))
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self)
return loader.load_weights(weights)
vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py
+386 -142
View File
@@ -4,19 +4,30 @@
from __future__ import annotations
import copy
import os
from collections.abc import Iterable
from typing import Any
import numpy as np
import PIL.Image
import torch
import torch.nn as nn
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg, retrieve_timesteps
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
from vllm.logger import init_logger
from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.parallel_state import get_cfg_group, get_classifier_free_guidance_rank
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.lora.request import LoRARequest
from .pipeline_ltx2 import (
LTX2Pipeline,
@@ -26,6 +37,9 @@ from .pipeline_ltx2 import (
from .pipeline_ltx2 import (
get_ltx2_post_process_func as _get_ltx2_post_process_func,
)
from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
logger = init_logger(__name__)
def get_ltx2_post_process_func(od_config: OmniDiffusionConfig):
@@ -61,6 +75,7 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
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,
@@ -74,11 +89,29 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
mask_shape = (batch_size, 1, num_frames, height, width)
if latents is not None:
conditioning_mask = latents.new_zeros(mask_shape)
conditioning_mask[:, :, 0] = 1.0
if latents.ndim == 5:
# conditioning_mask needs to the same shape as latents in two stages generation.
batch_size, _, num_frames, height, width = latents.shape
mask_shape = (batch_size, 1, num_frames, height, width)
conditioning_mask = latents.new_zeros(mask_shape)
conditioning_mask[:, :, 0] = 1.0
latents = self._normalize_latents(
latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
)
latents = self._create_noised_state(latents, noise_scale * (1 - conditioning_mask), generator)
# latents are of shape [B, C, F, H, W], need to be packed
latents = self._pack_latents(
latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
)
else:
conditioning_mask = latents.new_zeros(mask_shape)
conditioning_mask[:, :, 0] = 1.0
conditioning_mask = self._pack_latents(
conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
).squeeze(-1)
if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape:
raise ValueError(
"Provided `latents` tensor has shape"
@@ -234,9 +267,11 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
num_frames: int | None = None,
frame_rate: float | None = None,
num_inference_steps: int | None = None,
sigmas: list[float] | None = None,
timesteps: list[int] | None = None,
guidance_scale: float = 4.0,
guidance_rescale: float = 0.0,
noise_scale: float = 0.0,
num_videos_per_prompt: int | None = 1,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
@@ -421,6 +456,26 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
additive_mask=True,
)
latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
latent_height = height // self.vae_spatial_compression_ratio
latent_width = width // self.vae_spatial_compression_ratio
if latents is not None:
if latents.ndim == 5:
logger.info(
"Got latents of shape [batch_size, latent_dim, latent_frames, latent_height, latent_width], `latent_num_frames`, `latent_height`, `latent_width` will be inferred." # noqa
)
_, _, latent_num_frames, latent_height, latent_width = latents.shape # [B, C, F, H, W]
elif latents.ndim == 3:
logger.warning(
f"You have supplied packed `latents` of shape {latents.shape}, so the latent dims cannot be"
f" inferred. Make sure the supplied `height`, `width`, and `num_frames` are correct."
)
else:
raise ValueError(
f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, latent_dim, latent_frames, latent_height, latent_width]." # noqa
)
video_sequence_length = latent_num_frames * latent_height * latent_width
if latents is None:
if isinstance(image, torch.Tensor):
if image.ndim == 3:
@@ -439,6 +494,7 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
height,
width,
num_frames,
noise_scale,
torch.float32,
device,
generator,
@@ -447,32 +503,58 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
if self.do_classifier_free_guidance and not cfg_parallel_ready:
conditioning_mask = torch.cat([conditioning_mask, conditioning_mask])
duration_s = num_frames / 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)
if audio_latents is not None:
if audio_latents.ndim == 4:
logger.info(
"Got audio_latents of shape [batch_size, num_channels, audio_length, mel_bins], `audio_num_frames` will be inferred." # noqa
)
_, _, audio_num_frames, _ = audio_latents.shape # [B, C, L, M]
elif audio_latents.ndim == 3:
logger.warning(
f"You have supplied packed `audio_latents` of shape {audio_latents.shape}, so the latent dims"
f" cannot be inferred. Make sure the supplied `num_frames` and `frame_rate` are correct."
)
else:
raise ValueError(
f"Provided `audio_latents` tensor has shape {audio_latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, num_channels, audio_length, mel_bins]." # noqa
)
num_mel_bins = self.audio_vae.config.mel_bins if getattr(self, "audio_vae", None) is not None else 64
latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
num_channels_latents_audio = (
self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8
)
audio_latents, audio_num_frames = self.prepare_audio_latents(
# padding audio_latents if needed
sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1)
if sp_size > 1:
pad_len = (sp_size - (audio_num_frames % sp_size)) % sp_size
if pad_len > 0:
if audio_latents is not None:
pad_shape = list(audio_latents.shape)
pad_shape[2] = pad_len
padding = torch.zeros(pad_shape, dtype=audio_latents.dtype, device=audio_latents.device)
audio_latents = torch.cat([audio_latents, padding], dim=2)
audio_num_frames += pad_len
audio_latents = self.prepare_audio_latents(
batch_size * num_videos_per_prompt,
num_channels_latents=num_channels_latents_audio,
audio_latent_length=audio_num_frames,
num_mel_bins=num_mel_bins,
num_frames=num_frames,
frame_rate=frame_rate,
sampling_rate=self.audio_sampling_rate,
hop_length=self.audio_hop_length,
noise_scale=noise_scale,
dtype=torch.float32,
device=device,
generator=generator,
latents=audio_latents,
)
latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
latent_height = height // self.vae_spatial_compression_ratio
latent_width = width // self.vae_spatial_compression_ratio
video_sequence_length = latent_num_frames * latent_height * latent_width
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
mu = calculate_shift(
video_sequence_length,
self.scheduler.config.get("base_image_seq_len", 1024),
@@ -497,7 +579,6 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
sigmas=sigmas,
mu=mu,
)
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
self._num_timesteps = len(timesteps)
video_coords = self.transformer.rope.prepare_video_coords(
@@ -506,69 +587,142 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
audio_coords = self.transformer.audio_rope.prepare_audio_coords(
audio_latents.shape[0], audio_num_frames, audio_latents.device
)
# Duplicate the positional ids as well if using CFG
if self.do_classifier_free_guidance and not cfg_parallel_ready:
video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) # Repeat twice in batch dim
audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
for i, t in enumerate(timesteps):
if self.interrupt:
continue
with self.progress_bar(total=len(timesteps)) as pbar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
self._current_timestep = t
self._current_timestep = t
if cfg_parallel_ready:
latent_model_input = latents.to(prompt_embeds.dtype)
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
if cfg_parallel_ready:
latent_model_input = latents.to(prompt_embeds.dtype)
audio_latent_model_input = audio_latents.to(prompt_embeds.dtype)
timestep = t.expand(latent_model_input.shape[0])
video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
timestep = t.expand(latent_model_input.shape[0])
video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
positive_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": connector_prompt_embeds,
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
"timestep": video_timestep,
"audio_timestep": timestep,
"encoder_attention_mask": connector_attention_mask,
"audio_encoder_attention_mask": connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
negative_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": negative_connector_prompt_embeds,
"audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
"timestep": video_timestep,
"audio_timestep": timestep,
"encoder_attention_mask": negative_connector_attention_mask,
"audio_encoder_attention_mask": negative_connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
positive_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": connector_prompt_embeds,
"audio_encoder_hidden_states": connector_audio_prompt_embeds,
"timestep": video_timestep,
"audio_timestep": timestep,
"encoder_attention_mask": connector_attention_mask,
"audio_encoder_attention_mask": connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
negative_kwargs = {
"hidden_states": latent_model_input,
"audio_hidden_states": audio_latent_model_input,
"encoder_hidden_states": negative_connector_prompt_embeds,
"audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
"timestep": video_timestep,
"audio_timestep": timestep,
"encoder_attention_mask": negative_connector_attention_mask,
"audio_encoder_attention_mask": negative_connector_attention_mask,
"num_frames": latent_num_frames,
"height": latent_height,
"width": latent_width,
"fps": frame_rate,
"audio_num_frames": audio_num_frames,
"video_coords": video_coords,
"audio_coords": audio_coords,
"attention_kwargs": attention_kwargs,
"return_dict": False,
}
noise_pred_video, noise_pred_audio = self.predict_noise_av_maybe_with_cfg(
do_true_cfg=True,
true_cfg_scale=guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
guidance_rescale=guidance_rescale,
cfg_normalize=False,
)
noise_pred_video, noise_pred_audio = self.predict_noise_av_maybe_with_cfg(
do_true_cfg=True,
true_cfg_scale=guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
guidance_rescale=guidance_rescale,
cfg_normalize=False,
)
if get_classifier_free_guidance_rank() == 0:
latents = self._step_video_latents_i2v(
noise_pred_video,
latents,
t,
latent_num_frames,
latent_height,
latent_width,
)
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
cfg_group = get_cfg_group()
latents = latents.contiguous()
audio_latents = audio_latents.contiguous()
cfg_group.broadcast(latents, src=0)
cfg_group.broadcast(audio_latents, src=0)
else:
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
latent_model_input = latent_model_input.to(prompt_embeds.dtype)
audio_latent_model_input = (
torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents
)
audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype)
timestep = t.expand(latent_model_input.shape[0])
video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
with self._transformer_cache_context("cond_uncond"):
noise_pred_video, noise_pred_audio = self.transformer(
hidden_states=latent_model_input,
audio_hidden_states=audio_latent_model_input,
encoder_hidden_states=connector_prompt_embeds,
audio_encoder_hidden_states=connector_audio_prompt_embeds,
timestep=video_timestep,
audio_timestep=timestep,
encoder_attention_mask=connector_attention_mask,
audio_encoder_attention_mask=connector_attention_mask,
num_frames=latent_num_frames,
height=latent_height,
width=latent_width,
fps=frame_rate,
audio_num_frames=audio_num_frames,
video_coords=video_coords,
audio_coords=audio_coords,
attention_kwargs=attention_kwargs,
return_dict=False,
)
noise_pred_video = noise_pred_video.float()
noise_pred_audio = noise_pred_audio.float()
if self.do_classifier_free_guidance:
noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
noise_pred_video = noise_pred_video_uncond + guidance_scale * (
noise_pred_video_text - noise_pred_video_uncond
)
noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (
noise_pred_audio_text - noise_pred_audio_uncond
)
if guidance_rescale > 0:
noise_pred_video = rescale_noise_cfg(
noise_pred_video, noise_pred_video_text, guidance_rescale=guidance_rescale
)
noise_pred_audio = rescale_noise_cfg(
noise_pred_audio, noise_pred_audio_text, guidance_rescale=guidance_rescale
)
if get_classifier_free_guidance_rank() == 0:
latents = self._step_video_latents_i2v(
noise_pred_video,
latents,
@@ -577,79 +731,10 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
latent_height,
latent_width,
)
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
cfg_group = get_cfg_group()
latents = latents.contiguous()
audio_latents = audio_latents.contiguous()
cfg_group.broadcast(latents, src=0)
cfg_group.broadcast(audio_latents, src=0)
else:
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
latent_model_input = latent_model_input.to(prompt_embeds.dtype)
audio_latent_model_input = (
torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents
)
audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype)
timestep = t.expand(latent_model_input.shape[0])
video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
with self._transformer_cache_context("cond_uncond"):
noise_pred_video, noise_pred_audio = self.transformer(
hidden_states=latent_model_input,
audio_hidden_states=audio_latent_model_input,
encoder_hidden_states=connector_prompt_embeds,
audio_encoder_hidden_states=connector_audio_prompt_embeds,
timestep=video_timestep,
audio_timestep=timestep,
encoder_attention_mask=connector_attention_mask,
audio_encoder_attention_mask=connector_attention_mask,
num_frames=latent_num_frames,
height=latent_height,
width=latent_width,
fps=frame_rate,
audio_num_frames=audio_num_frames,
video_coords=video_coords,
audio_coords=audio_coords,
attention_kwargs=attention_kwargs,
return_dict=False,
)
noise_pred_video = noise_pred_video.float()
noise_pred_audio = noise_pred_audio.float()
if self.do_classifier_free_guidance:
noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
noise_pred_video = noise_pred_video_uncond + guidance_scale * (
noise_pred_video_text - noise_pred_video_uncond
)
noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (
noise_pred_audio_text - noise_pred_audio_uncond
)
if guidance_rescale > 0:
noise_pred_video = rescale_noise_cfg(
noise_pred_video, noise_pred_video_text, guidance_rescale=guidance_rescale
)
noise_pred_audio = rescale_noise_cfg(
noise_pred_audio, noise_pred_audio_text, guidance_rescale=guidance_rescale
)
latents = self._step_video_latents_i2v(
noise_pred_video,
latents,
t,
latent_num_frames,
latent_height,
latent_width,
)
audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
pass
pbar.update()
latents = self._unpack_latents(
latents,
@@ -703,3 +788,162 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
return DiffusionOutput(output=(video, audio))
return DiffusionOutput(output=(video, audio))
class LTX2ImageToVideoTwoStagesPipeline(nn.Module):
"""LTXImageToVideoTwoStagesPipeline is for two stages image to video generation"""
support_image_input = True
def __init__(
self,
*,
od_config: OmniDiffusionConfig,
prefix: str = "",
):
super().__init__()
self.device = get_local_device()
self.dtype = getattr(od_config, "dtype", torch.bfloat16)
self.model_path = od_config.model
self.distilled = False
# User provided model path may contain '/' in the end and basename function
# will not return the expected directory name, so we need to remove it by normpath
if "distilled" in os.path.basename(os.path.normpath(self.model_path)):
self.distilled = True
else:
raise NotImplementedError(f"{self.model_path} is not supported for {self.__class__.__name__}.")
self.pipe = LTX2ImageToVideoPipeline(od_config=od_config, prefix=prefix)
self.upsample_pipe = LTX2LatentUpsamplePipeline(
vae=self.pipe.vae,
od_config=od_config,
)
self.lora_manager = DiffusionLoRAManager(
pipeline=self.pipe,
device=self.device,
dtype=self.dtype,
max_cached_adapters=od_config.max_cpu_loras,
)
self.weights_sources = [
DiffusersPipelineLoader.ComponentSource(
model_or_path=od_config.model,
subfolder="transformer",
revision=None,
prefix="pipe.transformer.",
fall_back_to_pt=True,
),
]
@torch.no_grad()
def forward(
self,
req: OmniDiffusionRequest,
image: PIL.Image.Image | torch.Tensor | None = None,
prompt: str | list[str] | None = None,
negative_prompt: str | list[str] | None = None,
height: int | None = None,
width: int | None = None,
num_frames: int | None = None,
frame_rate: float | None = None,
num_inference_steps: int | None = None,
sigmas: list[float] | None = None,
timesteps: list[int] | None = None,
guidance_scale: float = 4.0,
guidance_rescale: float = 0.0,
noise_scale: float = 0.0,
num_videos_per_prompt: int | None = 1,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
audio_latents: torch.Tensor | None = None,
prompt_embeds: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
prompt_attention_mask: torch.Tensor | None = None,
negative_prompt_attention_mask: torch.Tensor | None = None,
decode_timestep: float | list[float] = 0.0,
decode_noise_scale: float | list[float] | None = None,
output_type: str = "np",
return_dict: bool = True,
attention_kwargs: dict[str, Any] | None = None,
max_sequence_length: int | None = None,
):
video_latent, audio_latent = self.pipe(
req=req,
image=image,
prompt=prompt,
negative_prompt=negative_prompt,
height=height,
width=width,
num_frames=num_frames,
frame_rate=frame_rate,
num_inference_steps=num_inference_steps,
sigmas=DISTILLED_SIGMA_VALUES if self.distilled else None,
timesteps=timesteps,
guidance_scale=guidance_scale,
guidance_rescale=guidance_rescale,
noise_scale=noise_scale,
num_videos_per_prompt=num_videos_per_prompt,
generator=generator,
latents=latents,
audio_latents=audio_latents,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
decode_timestep=decode_timestep,
decode_noise_scale=decode_noise_scale,
output_type="latent",
return_dict=return_dict,
attention_kwargs=attention_kwargs,
max_sequence_length=max_sequence_length,
).output
upscaled_video_latent = self.upsample_pipe(
latents=video_latent,
output_type="latent",
return_dict=False,
)[0]
if not self.distilled:
# Load Stage 2 distilled LoRA
lora_path = f"{self.model_path}/ltx-2-19b-distilled-lora-384.safetensors"
lora_request = LoRARequest(
lora_name="stage_2_distilled",
lora_int_id=1,
lora_path=lora_path,
)
self.lora_manager.set_active_adapter(lora_request, lora_scale=1.0)
# Change scheduler to use Stage 2 distilled sigmas as is
new_scheduler = FlowMatchEulerDiscreteScheduler.from_config(
self.pipe.scheduler.config,
use_dynamic_shifting=False,
shift_terminal=None,
)
self.pipe.scheduler = new_scheduler
stage_2_req = copy.copy(req)
stage_2_req.sampling_params = req.sampling_params.clone()
stage_2_req.sampling_params.num_inference_steps = 3
video, audio = self.pipe(
req=stage_2_req,
latents=upscaled_video_latent,
audio_latents=audio_latent,
prompt=prompt,
negative_prompt=negative_prompt,
noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
sigmas=STAGE_2_DISTILLED_SIGMA_VALUES,
guidance_scale=1.0,
generator=generator,
output_type="np",
return_dict=False,
).output
return DiffusionOutput(output=(video, audio))
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self)
return loader.load_weights(weights)
vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py
+262
View File
@@ -0,0 +1,262 @@
import os
import torch
import torch.nn as nn
from diffusers import AutoencoderKLLTX2Video
from diffusers.image_processor import PipelineImageInput
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
from vllm.logger import init_logger
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device
logger = init_logger(__name__)
class LTX2LatentUpsamplePipeline(nn.Module):
def __init__(
self,
od_config: OmniDiffusionConfig,
vae: AutoencoderKLLTX2Video,
latent_upsampler: LTX2LatentUpsamplerModel = None,
) -> None:
super().__init__()
if vae is None:
raise ValueError("vae must be provided")
self.vae = vae
self.device = get_local_device()
model = od_config.model
local_files_only = os.path.exists(model)
if latent_upsampler is None:
# Use cpu context to create latent upsampler. The code k[:, None] @ k[None, :] in
# diffuser's BlurDownsample is not supported on GPU as k is type of torch.Int64
with torch.device("cpu"):
latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained(
model,
subfolder="latent_upsampler",
torch_dtype=torch.bfloat16,
local_files_only=local_files_only,
).to(self.device)
self.latent_upsampler = latent_upsampler
self.vae_spatial_compression_ratio = (
self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32
)
self.vae_temporal_compression_ratio = (
self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8
)
self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio)
def prepare_latents(
self,
video: torch.Tensor | None = None,
batch_size: int = 1,
num_frames: int = 121,
height: int = 512,
width: int = 768,
spatial_patch_size: int = 1,
temporal_patch_size: int = 1,
dtype: torch.dtype | None = None,
device: torch.device | None = None,
generator: torch.Generator | None = None,
latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
if latents.ndim == 3:
# Convert token seq [B, S, D] to latent video [B, C, F, H, W]
latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
latent_height = height // self.vae_spatial_compression_ratio
latent_width = width // self.vae_spatial_compression_ratio
latents = self._unpack_latents(
latents, latent_num_frames, latent_height, latent_width, spatial_patch_size, temporal_patch_size
)
return latents.to(device=device, dtype=dtype)
video = video.to(device=device, dtype=self.vae.dtype)
if isinstance(generator, list):
if len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
init_latents = [
retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), generator[i]) for i in range(batch_size)
]
else:
init_latents = [retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator) for vid in video]
init_latents = torch.cat(init_latents, dim=0).to(dtype)
# NOTE: latent upsampler operates on the unnormalized latents, so don't normalize here
# init_latents = self._normalize_latents(init_latents, self.vae.latents_mean, self.vae.latents_std)
return init_latents
def adain_filter_latent(self, latents: torch.Tensor, reference_latents: torch.Tensor, factor: float = 1.0):
result = latents.clone()
for i in range(latents.size(0)):
for c in range(latents.size(1)):
r_sd, r_mean = torch.std_mean(reference_latents[i, c], dim=None) # index by original dim order
i_sd, i_mean = torch.std_mean(result[i, c], dim=None)
result[i, c] = ((result[i, c] - i_mean) / i_sd) * r_sd + r_mean
result = torch.lerp(latents, result, factor)
return result
def tone_map_latents(self, latents: torch.Tensor, compression: float) -> torch.Tensor:
# Remap [0-1] to [0-0.75] and apply sigmoid compression in one shot
scale_factor = compression * 0.75
abs_latents = torch.abs(latents)
# Sigmoid compression: sigmoid shifts large values toward 0.2, small values stay ~1.0
# When scale_factor=0, sigmoid term vanishes, when scale_factor=0.75, full effect
sigmoid_term = torch.sigmoid(4.0 * scale_factor * (abs_latents - 1.0))
scales = 1.0 - 0.8 * scale_factor * sigmoid_term
filtered = latents * scales
return filtered
@staticmethod
# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._denormalize_latents
def _denormalize_latents(
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
) -> torch.Tensor:
# Denormalize latents across the channel dimension [B, C, F, H, W]
latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents = latents * latents_std / scaling_factor + latents_mean
return latents
@staticmethod
# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._unpack_latents
def _unpack_latents(
latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1
) -> torch.Tensor:
# Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) # noqa
# are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of
# what happens in the `_pack_latents` method.
batch_size = latents.size(0)
latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size)
latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
return latents
def check_inputs(self, video, height, width, latents, tone_map_compression_ratio):
if height % self.vae_spatial_compression_ratio != 0 or width % self.vae_spatial_compression_ratio != 0:
raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.")
if video is not None and latents is not None:
raise ValueError("Only one of `video` or `latents` can be provided.")
if video is None and latents is None:
raise ValueError("One of `video` or `latents` has to be provided.")
if not (0 <= tone_map_compression_ratio <= 1):
raise ValueError("`tone_map_compression_ratio` must be in the range [0, 1]")
def forward(
self,
video: list[PipelineImageInput] | None = None,
height: int = 512,
width: int = 768,
num_frames: int = 121,
spatial_patch_size: int = 1,
temporal_patch_size: int = 1,
latents: torch.Tensor | None = None,
latents_normalized: bool = False,
decode_timestep: float | list[float] = 0.0,
decode_noise_scale: float | list[float] | None = None,
adain_factor: float = 0.0,
tone_map_compression_ratio: float = 0.0,
generator: torch.Generator | list[torch.Generator] | None = None,
output_type: str | None = "pil",
return_dict: bool = True,
):
self.check_inputs(
video=video,
height=height,
width=width,
latents=latents,
tone_map_compression_ratio=tone_map_compression_ratio,
)
if video is not None:
# Batched video input is not yet tested/supported. TODO: take a look later
batch_size = 1
else:
batch_size = latents.shape[0]
device = self.device
if video is not None:
num_frames = len(video)
if num_frames % self.vae_temporal_compression_ratio != 1:
num_frames = num_frames // self.vae_temporal_compression_ratio * self.vae_temporal_compression_ratio + 1
video = video[:num_frames]
logger.warning(
f"Video length expected to be of the form `k * {self.vae_temporal_compression_ratio} + 1` but is {len(video)}. Truncating to {num_frames} frames." # noqa
)
video = self.video_processor.preprocess_video(video, height=height, width=width)
video = video.to(device=device, dtype=torch.float32)
latents_supplied = latents is not None
latents = self.prepare_latents(
video=video,
batch_size=batch_size,
num_frames=num_frames,
height=height,
width=width,
spatial_patch_size=spatial_patch_size,
temporal_patch_size=temporal_patch_size,
dtype=torch.float32,
device=device,
generator=generator,
latents=latents,
)
if latents_supplied and latents_normalized:
latents = self._denormalize_latents(
latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
)
latents = latents.to(self.latent_upsampler.dtype)
latents_upsampled = self.latent_upsampler(latents)
if adain_factor > 0.0:
latents = self.adain_filter_latent(latents_upsampled, latents, adain_factor)
else:
latents = latents_upsampled
if tone_map_compression_ratio > 0.0:
latents = self.tone_map_latents(latents, tone_map_compression_ratio)
if output_type == "latent":
video = latents
else:
if not self.vae.config.timestep_conditioning:
timestep = None
else:
noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype)
if not isinstance(decode_timestep, list):
decode_timestep = [decode_timestep] * batch_size
if decode_noise_scale is None:
decode_noise_scale = decode_timestep
elif not isinstance(decode_noise_scale, list):
decode_noise_scale = [decode_noise_scale] * batch_size
timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype)
decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[
:, None, None, None, None
]
latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
video = self.vae.decode(latents, timestep, return_dict=False)[0]
video = self.video_processor.postprocess_video(video, output_type=output_type)
if not return_dict:
return (video,)
return DiffusionOutput(output=(video,))
vllm_omni/diffusion/registry.py
+19 -1
View File
@@ -12,6 +12,7 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig, get_sp_plan_from_model
from vllm_omni.diffusion.forward_context import get_forward_context
from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
from vllm_omni.diffusion.utils.tf_utils import find_module_with_attr
logger = init_logger(__name__)
@@ -72,6 +73,16 @@ _DIFFUSION_MODELS = {
"pipeline_ltx2_image2video",
"LTX2ImageToVideoPipeline",
),
"LTX2TwoStagesPipeline": (
"ltx2",
"pipeline_ltx2",
"LTX2TwoStagesPipeline",
),
"LTX2ImageToVideoTwoStagesPipeline": (
"ltx2",
"pipeline_ltx2_image2video",
"LTX2ImageToVideoTwoStagesPipeline",
),
"StableAudioPipeline": (
"stable_audio",
"pipeline_stable_audio",
@@ -266,7 +277,12 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -
for attr in transformer_attrs:
if not hasattr(model, attr):
continue
# Some pipeline like LTX2TwoStagesPipeline have recursive
# modules that have the transformer
module = find_module_with_attr(model, attr)
if module is None:
continue
model = module
transformer = getattr(model, attr)
if transformer is None:
@@ -323,7 +339,9 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
"WanPipeline": "get_wan22_post_process_func",
"WanVACEPipeline": "get_wan22_vace_post_process_func",
"LTX2Pipeline": "get_ltx2_post_process_func",
"LTX2TwoStagesPipeline": "get_ltx2_post_process_func",
"LTX2ImageToVideoPipeline": "get_ltx2_post_process_func",
"LTX2ImageToVideoTwoStagesPipeline": "get_ltx2_post_process_func",
"StableAudioPipeline": "get_stable_audio_post_process_func",
"WanImageToVideoPipeline": "get_wan22_i2v_post_process_func",
"LongCatImagePipeline": "get_longcat_image_post_process_func",
vllm_omni/diffusion/utils/tf_utils.py
+24
View File
@@ -52,3 +52,27 @@ def get_transformer_config_kwargs(
pass
return filtered_params
def find_module_with_attr(model, attr_name="transformer"):
"""
This function searches for a module in the model that has the specified attribute.
If the model itself has the attribute, it returns the model.
If none of the modules have the attribute, it returns None.
"""
if hasattr(model, attr_name):
return model
for _, child in model.named_children():
if hasattr(child, attr_name):
return child
return None
def get_transformer_from_pipeline(pipeline: Any):
pipe = find_module_with_attr(pipeline, attr_name="transformer")
if pipe is not None:
return pipe.transformer
return None