427 lines
17 KiB
Python
427 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from collections.abc import Iterable
|
|
|
|
import torch
|
|
from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler
|
|
from diffusers.utils.torch_utils import randn_tensor
|
|
from torch import nn
|
|
from transformers import AutoTokenizer, UMT5EncoderModel
|
|
from vllm.model_executor.models.utils import AutoWeightsLoader
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.distributed.utils import get_local_device
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
|
|
|
|
def get_wan22_post_process_func(
|
|
od_config: OmniDiffusionConfig,
|
|
):
|
|
from diffusers.video_processor import VideoProcessor
|
|
|
|
video_processor = VideoProcessor(vae_scale_factor=8)
|
|
|
|
def post_process_func(
|
|
video: torch.Tensor,
|
|
output_type: str = "np",
|
|
):
|
|
if output_type == "latent":
|
|
return video
|
|
return video_processor.postprocess_video(video, output_type=output_type)
|
|
|
|
return post_process_func
|
|
|
|
|
|
class Wan22Pipeline(nn.Module):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
od_config: OmniDiffusionConfig,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
self.od_config = od_config
|
|
|
|
self.device = get_local_device()
|
|
dtype = getattr(od_config, "dtype", torch.bfloat16)
|
|
|
|
model = od_config.model
|
|
local_files_only = os.path.exists(model)
|
|
|
|
# Set up weights sources for transformer(s)
|
|
self.weights_sources = [
|
|
DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path=od_config.model,
|
|
subfolder="transformer",
|
|
revision=None,
|
|
prefix="transformer.",
|
|
fall_back_to_pt=True,
|
|
),
|
|
DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path=od_config.model,
|
|
subfolder="transformer_2",
|
|
revision=None,
|
|
prefix="transformer_2.",
|
|
fall_back_to_pt=True,
|
|
),
|
|
]
|
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
|
|
self.text_encoder = UMT5EncoderModel.from_pretrained(
|
|
model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
|
|
).to(self.device)
|
|
self.vae = AutoencoderKLWan.from_pretrained(
|
|
model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
|
|
).to(self.device)
|
|
|
|
# Initialize transformers without from_pretrained (weights loaded via load_weights)
|
|
self.transformer = WanTransformer3DModel()
|
|
self.transformer_2 = WanTransformer3DModel()
|
|
|
|
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
|
|
model, subfolder="scheduler", local_files_only=local_files_only
|
|
)
|
|
# Apply flow_shift if specified (12.0 for 480p, 5.0 for 720p recommended for Wan2.2)
|
|
if od_config.flow_shift is not None:
|
|
self.scheduler.config.flow_shift = od_config.flow_shift
|
|
|
|
self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4
|
|
self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8
|
|
self.boundary_ratio = od_config.boundary_ratio
|
|
|
|
self._guidance_scale = None
|
|
self._guidance_scale_2 = None
|
|
self._num_timesteps = None
|
|
self._current_timestep = None
|
|
|
|
@property
|
|
def guidance_scale(self):
|
|
return self._guidance_scale
|
|
|
|
@property
|
|
def do_classifier_free_guidance(self):
|
|
return self._guidance_scale is not None and self._guidance_scale > 1.0
|
|
|
|
@property
|
|
def num_timesteps(self):
|
|
return self._num_timesteps
|
|
|
|
@property
|
|
def current_timestep(self):
|
|
return self._current_timestep
|
|
|
|
@torch.no_grad()
|
|
def forward(
|
|
self,
|
|
req: OmniDiffusionRequest,
|
|
prompt: str | None = None,
|
|
negative_prompt: str | None = None,
|
|
height: int | None = None,
|
|
width: int | None = None,
|
|
num_inference_steps: int | None = None,
|
|
guidance_scale: float | tuple[float, float] = 4.0,
|
|
frame_num: int | None = None,
|
|
output_type: str | None = "np",
|
|
generator: torch.Generator | None = None,
|
|
prompt_embeds: torch.Tensor | None = None,
|
|
negative_prompt_embeds: torch.Tensor | None = None,
|
|
attention_kwargs: dict | None = None,
|
|
**kwargs,
|
|
) -> DiffusionOutput:
|
|
prompt = req.prompt if req.prompt is not None else prompt
|
|
negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt
|
|
if prompt is None and prompt_embeds is None:
|
|
raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
|
|
|
|
height = req.height or height or 720
|
|
width = req.width or width or 1280
|
|
num_frames = req.num_frames if req.num_frames else frame_num or 81
|
|
num_steps = req.num_inference_steps or num_inference_steps or 40
|
|
|
|
guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
|
|
guidance_high = (
|
|
req.guidance_scale_2
|
|
if req.guidance_scale_2 is not None
|
|
else (
|
|
guidance_scale[1]
|
|
if isinstance(guidance_scale, (list, tuple)) and len(guidance_scale) > 1
|
|
else guidance_low
|
|
)
|
|
)
|
|
|
|
# record guidance for properties
|
|
self._guidance_scale = guidance_low
|
|
self._guidance_scale_2 = guidance_high
|
|
|
|
# validate shapes
|
|
self.check_inputs(
|
|
prompt=prompt,
|
|
negative_prompt=negative_prompt,
|
|
height=height,
|
|
width=width,
|
|
prompt_embeds=prompt_embeds,
|
|
negative_prompt_embeds=negative_prompt_embeds,
|
|
guidance_scale_2=guidance_high if self.boundary_ratio is not None else None,
|
|
)
|
|
|
|
if num_frames % self.vae_scale_factor_temporal != 1:
|
|
num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
|
|
num_frames = max(num_frames, 1)
|
|
|
|
device = self.device
|
|
dtype = self.transformer.dtype
|
|
|
|
# Seed / generator
|
|
if generator is None:
|
|
generator = req.generator
|
|
if generator is None and req.seed is not None:
|
|
generator = torch.Generator(device=device).manual_seed(req.seed)
|
|
|
|
# Encode prompts
|
|
if prompt_embeds is None:
|
|
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
|
|
prompt=prompt,
|
|
negative_prompt=negative_prompt,
|
|
do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
|
|
num_videos_per_prompt=req.num_outputs_per_prompt or 1,
|
|
max_sequence_length=req.max_sequence_length or 512,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
else:
|
|
prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
|
|
if negative_prompt_embeds is not None:
|
|
negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
|
|
elif guidance_low > 1.0 or guidance_high > 1.0:
|
|
raise ValueError(
|
|
"negative_prompt_embeds must be provided when prompt_embeds are given and guidance > 1."
|
|
)
|
|
|
|
# Timesteps
|
|
self.scheduler.set_timesteps(num_steps, device=device)
|
|
timesteps = self.scheduler.timesteps
|
|
self._num_timesteps = len(timesteps)
|
|
boundary_timestep = None
|
|
if self.boundary_ratio is not None:
|
|
boundary_timestep = self.boundary_ratio * self.scheduler.config.num_train_timesteps
|
|
|
|
# Latents
|
|
num_channels_latents = self.transformer.config.in_channels
|
|
latents = self.prepare_latents(
|
|
batch_size=prompt_embeds.shape[0],
|
|
num_channels_latents=num_channels_latents,
|
|
height=height,
|
|
width=width,
|
|
num_frames=num_frames,
|
|
dtype=torch.float32,
|
|
device=device,
|
|
generator=generator,
|
|
latents=req.latents,
|
|
)
|
|
|
|
if attention_kwargs is None:
|
|
attention_kwargs = {}
|
|
|
|
# Denoising
|
|
for t in timesteps:
|
|
self._current_timestep = t
|
|
current_model = self.transformer
|
|
current_guidance_scale = guidance_low
|
|
if boundary_timestep is not None and t < boundary_timestep:
|
|
current_model = self.transformer_2
|
|
current_guidance_scale = guidance_high
|
|
|
|
latent_model_input = latents.to(dtype)
|
|
timestep = t.expand(latents.shape[0])
|
|
|
|
noise_pred = current_model(
|
|
hidden_states=latent_model_input,
|
|
timestep=timestep,
|
|
encoder_hidden_states=prompt_embeds,
|
|
attention_kwargs=attention_kwargs,
|
|
return_dict=False,
|
|
)[0]
|
|
|
|
if current_guidance_scale > 1.0 and negative_prompt_embeds is not None:
|
|
noise_uncond = current_model(
|
|
hidden_states=latent_model_input,
|
|
timestep=timestep,
|
|
encoder_hidden_states=negative_prompt_embeds,
|
|
attention_kwargs=attention_kwargs,
|
|
return_dict=False,
|
|
)[0]
|
|
noise_pred = noise_uncond + current_guidance_scale * (noise_pred - noise_uncond)
|
|
|
|
latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
|
|
|
|
self._current_timestep = None
|
|
|
|
# Decode
|
|
if output_type == "latent":
|
|
output = latents
|
|
else:
|
|
latents = latents.to(self.vae.dtype)
|
|
latents_mean = (
|
|
torch.tensor(self.vae.config.latents_mean)
|
|
.view(1, self.vae.config.z_dim, 1, 1, 1)
|
|
.to(latents.device, latents.dtype)
|
|
)
|
|
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
|
|
latents.device, latents.dtype
|
|
)
|
|
latents = latents / latents_std + latents_mean
|
|
output = self.vae.decode(latents, return_dict=False)[0]
|
|
|
|
return DiffusionOutput(output=output)
|
|
|
|
def encode_prompt(
|
|
self,
|
|
prompt: str | list[str],
|
|
negative_prompt: str | list[str] | None = None,
|
|
do_classifier_free_guidance: bool = True,
|
|
num_videos_per_prompt: int = 1,
|
|
max_sequence_length: int = 512,
|
|
device: torch.device | None = None,
|
|
dtype: torch.dtype | None = None,
|
|
):
|
|
device = device or self.device
|
|
dtype = dtype or self.text_encoder.dtype
|
|
|
|
prompt = [prompt] if isinstance(prompt, str) else prompt
|
|
prompt_clean = [self._prompt_clean(p) for p in prompt]
|
|
batch_size = len(prompt_clean)
|
|
|
|
text_inputs = self.tokenizer(
|
|
prompt_clean,
|
|
padding="max_length",
|
|
max_length=max_sequence_length,
|
|
truncation=True,
|
|
add_special_tokens=True,
|
|
return_attention_mask=True,
|
|
return_tensors="pt",
|
|
)
|
|
ids, mask = text_inputs.input_ids, text_inputs.attention_mask
|
|
seq_lens = mask.gt(0).sum(dim=1).long()
|
|
|
|
prompt_embeds = self.text_encoder(ids.to(device), mask.to(device)).last_hidden_state
|
|
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
|
|
prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
|
|
prompt_embeds = torch.stack(
|
|
[torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
|
|
)
|
|
|
|
_, seq_len, _ = prompt_embeds.shape
|
|
prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
|
|
prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
|
|
|
|
negative_prompt_embeds = None
|
|
if do_classifier_free_guidance:
|
|
negative_prompt = negative_prompt or ""
|
|
negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
|
|
neg_text_inputs = self.tokenizer(
|
|
[self._prompt_clean(p) for p in negative_prompt],
|
|
padding="max_length",
|
|
max_length=max_sequence_length,
|
|
truncation=True,
|
|
add_special_tokens=True,
|
|
return_attention_mask=True,
|
|
return_tensors="pt",
|
|
)
|
|
ids_neg, mask_neg = neg_text_inputs.input_ids, neg_text_inputs.attention_mask
|
|
seq_lens_neg = mask_neg.gt(0).sum(dim=1).long()
|
|
negative_prompt_embeds = self.text_encoder(ids_neg.to(device), mask_neg.to(device)).last_hidden_state
|
|
negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
|
|
negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)]
|
|
negative_prompt_embeds = torch.stack(
|
|
[
|
|
torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))])
|
|
for u in negative_prompt_embeds
|
|
],
|
|
dim=0,
|
|
)
|
|
negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1)
|
|
negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
|
|
|
|
return prompt_embeds, negative_prompt_embeds
|
|
|
|
@staticmethod
|
|
def _prompt_clean(text: str) -> str:
|
|
return " ".join(text.strip().split())
|
|
|
|
def prepare_latents(
|
|
self,
|
|
batch_size: int,
|
|
num_channels_latents: int,
|
|
height: int,
|
|
width: int,
|
|
num_frames: int,
|
|
dtype: torch.dtype | None,
|
|
device: torch.device | None,
|
|
generator: torch.Generator | list[torch.Generator] | None,
|
|
latents: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
if latents is not None:
|
|
return latents.to(device=device, dtype=dtype)
|
|
|
|
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
|
|
shape = (
|
|
batch_size,
|
|
num_channels_latents,
|
|
num_latent_frames,
|
|
int(height) // self.vae_scale_factor_spatial,
|
|
int(width) // self.vae_scale_factor_spatial,
|
|
)
|
|
if isinstance(generator, list) and len(generator) != batch_size:
|
|
raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.")
|
|
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
|
return latents
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
"""Load weights using AutoWeightsLoader for vLLM integration."""
|
|
loader = AutoWeightsLoader(self)
|
|
return loader.load_weights(weights)
|
|
|
|
def check_inputs(
|
|
self,
|
|
prompt,
|
|
negative_prompt,
|
|
height,
|
|
width,
|
|
prompt_embeds=None,
|
|
negative_prompt_embeds=None,
|
|
guidance_scale_2=None,
|
|
):
|
|
if height % 16 != 0 or width % 16 != 0:
|
|
raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
|
|
|
|
if prompt is not None and prompt_embeds is not None:
|
|
raise ValueError(
|
|
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
|
|
" only forward one of the two."
|
|
)
|
|
elif negative_prompt is not None and negative_prompt_embeds is not None:
|
|
raise ValueError(
|
|
f"Cannot forward both `negative_prompt`: {negative_prompt} and "
|
|
f"`negative_prompt_embeds`: {negative_prompt_embeds}. "
|
|
"Please make sure to only forward one of the two."
|
|
)
|
|
elif prompt is None and prompt_embeds is None:
|
|
raise ValueError(
|
|
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
|
|
)
|
|
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
|
|
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
|
|
elif negative_prompt is not None and (
|
|
not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list)
|
|
):
|
|
raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
|
|
|
|
if self.boundary_ratio is None and guidance_scale_2 is not None:
|
|
raise ValueError("`guidance_scale_2` is only supported when `boundary_ratio` is set.")
|