From cd71567b0686968d378486c38844e1fa5fc92998 Mon Sep 17 00:00:00 2001 From: Jerry Song <46962917+Songrui625@users.noreply.github.com> Date: Fri, 3 Apr 2026 23:17:49 +0800 Subject: [PATCH] [Model] Add two stages inference for model LTX-2 distilled. (#2260) Signed-off-by: Songrui625 --- docs/models/supported_models.md | 2 + .../image_to_video/image_to_video.py | 6 +- .../text_to_video/text_to_video.py | 9 +- .../diffusion/cache/cache_dit_backend.py | 18 +- vllm_omni/diffusion/models/ltx2/__init__.py | 10 +- .../diffusion/models/ltx2/pipeline_ltx2.py | 544 +++++++++++++----- .../models/ltx2/pipeline_ltx2_image2video.py | 528 ++++++++++++----- .../ltx2/pipeline_ltx2_latent_upsample.py | 262 +++++++++ vllm_omni/diffusion/registry.py | 20 +- vllm_omni/diffusion/utils/tf_utils.py | 24 + 10 files changed, 1120 insertions(+), 303 deletions(-) create mode 100644 vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 68024e18b..d611c0311 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 04e056691..c8c55c485 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -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, diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index a3aa818d2..322911c99 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -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, diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index e5337be12..a5055a068 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -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, diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py index 0a92d4f24..9f9d70f01 100644 --- a/vllm_omni/diffusion/models/ltx2/__init__.py +++ b/vllm_omni/diffusion/models/ltx2/__init__.py @@ -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", diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index 34263e217..efc342e93 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -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) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py index 5fa9cc797..11091518b 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py @@ -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) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py new file mode 100644 index 000000000..0c72a41d5 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py @@ -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,)) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index dcd227237..db8805722 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -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", diff --git a/vllm_omni/diffusion/utils/tf_utils.py b/vllm_omni/diffusion/utils/tf_utils.py index 44a788044..745410ec2 100644 --- a/vllm_omni/diffusion/utils/tf_utils.py +++ b/vllm_omni/diffusion/utils/tf_utils.py @@ -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