diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0ba642c92..0fbbe6eb7 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -20,6 +20,7 @@ th { | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | +| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | diff --git a/pyproject.toml b/pyproject.toml index 483f7625a..0748b7eab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -188,6 +188,7 @@ extend-ignore-identifiers-re = [ ".*arange", ".*MoBA", ".*temperal_downsample", + ".*_nd", ".*nothink.*", ".*NOTHINK.*", ".*nin.*", diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 9fa0a5bfc..12da9c43b 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -757,6 +757,7 @@ def initialize_model_parallel( backend=backend, parallel_mode="data", ) + vllm_parallel_state._DP = _DP global _CFG assert _CFG is None, "classifier_free_guidance group is already initialized" @@ -774,6 +775,7 @@ def initialize_model_parallel( backend=backend, parallel_mode="pipeline", ) + vllm_parallel_state._PP = _PP global _SP assert _SP is None, "sequence parallel group is already initialized" diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/__init__.py b/vllm_omni/diffusion/models/hunyuan_image_3/__init__.py new file mode 100644 index 000000000..98a3ac07b --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image_3/__init__.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Hunyuan Image 3 diffusion model components.""" + +from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_image_3_transformer import ( + HunyuanImage3Model, + HunyuanImage3Text2ImagePipeline, +) +from vllm_omni.diffusion.models.hunyuan_image_3.pipeline_hunyuan_image_3 import ( + HunyuanImage3Pipeline, +) + +__all__ = [ + "HunyuanImage3Pipeline", + "HunyuanImage3Model", + "HunyuanImage3Text2ImagePipeline", +] diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/autoencoder.py b/vllm_omni/diffusion/models/hunyuan_image_3/autoencoder.py new file mode 100644 index 000000000..efba2f274 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image_3/autoencoder.py @@ -0,0 +1,668 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils import BaseOutput +from diffusers.utils.torch_utils import randn_tensor +from einops import rearrange +from torch import Tensor, nn + + +class DiagonalGaussianDistribution: + def __init__(self, parameters: torch.Tensor, deterministic: bool = False): + if parameters.ndim == 3: + dim = 2 # (B, L, C) + elif parameters.ndim == 5 or parameters.ndim == 4: + dim = 1 # (B, C, T, H ,W) / (B, C, H, W) + else: + raise NotImplementedError + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like( + self.mean, device=self.parameters.device, dtype=self.parameters.dtype + ) + + def sample(self, generator: torch.Generator | None = None) -> torch.FloatTensor: + # make sure sample is on the same device as the parameters and has same dtype + sample = randn_tensor( + self.mean.shape, + generator=generator, + device=self.parameters.device, + dtype=self.parameters.dtype, + ) + x = self.mean + self.std * sample + return x + + +@dataclass +class DecoderOutput(BaseOutput): + sample: torch.FloatTensor + posterior: DiagonalGaussianDistribution | None = None + + +def swish(x: Tensor) -> Tensor: + return x * torch.sigmoid(x) + + +def forward_with_checkpointing(module, *inputs, use_checkpointing=False): + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if use_checkpointing: + return torch.utils.checkpoint.checkpoint(create_custom_forward(module), *inputs, use_reentrant=False) + else: + return module(*inputs) + + +class Conv3d(nn.Conv3d): + """ + Perform Conv3d on patches with numerical differences from nn.Conv3d within 1e-5. + Only symmetric padding is supported. + """ + + def forward(self, input): + B, C, T, H, W = input.shape + memory_count = (C * T * H * W) * 2 / 1024**3 + if memory_count > 2: + n_split = math.ceil(memory_count / 2) + assert n_split >= 2 + chunks = torch.chunk(input, chunks=n_split, dim=-3) + padded_chunks = [] + for i in range(len(chunks)): + if self.padding[0] > 0: + padded_chunk = F.pad( + chunks[i], + (0, 0, 0, 0, self.padding[0], self.padding[0]), + mode="constant" if self.padding_mode == "zeros" else self.padding_mode, + value=0, + ) + if i > 0: + padded_chunk[:, :, : self.padding[0]] = chunks[i - 1][:, :, -self.padding[0] :] + if i < len(chunks) - 1: + padded_chunk[:, :, -self.padding[0] :] = chunks[i + 1][:, :, : self.padding[0]] + else: + padded_chunk = chunks[i] + padded_chunks.append(padded_chunk) + padding_bak = self.padding + self.padding = (0, self.padding[1], self.padding[2]) + outputs = [] + for i in range(len(padded_chunks)): + outputs.append(super().forward(padded_chunks[i])) + self.padding = padding_bak + return torch.cat(outputs, dim=-3) + else: + return super().forward(input) + + +class AttnBlock(nn.Module): + """Attention with torch sdpa implementation.""" + + def __init__(self, in_channels: int): + super().__init__() + self.in_channels = in_channels + + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + + self.q = Conv3d(in_channels, in_channels, kernel_size=1) + self.k = Conv3d(in_channels, in_channels, kernel_size=1) + self.v = Conv3d(in_channels, in_channels, kernel_size=1) + self.proj_out = Conv3d(in_channels, in_channels, kernel_size=1) + + def attention(self, h_: Tensor) -> Tensor: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + b, c, f, h, w = q.shape + q = rearrange(q, "b c f h w -> b 1 (f h w) c").contiguous() + k = rearrange(k, "b c f h w -> b 1 (f h w) c").contiguous() + v = rearrange(v, "b c f h w -> b 1 (f h w) c").contiguous() + h_ = nn.functional.scaled_dot_product_attention(q, k, v) + + return rearrange(h_, "b 1 (f h w) c -> b c f h w", f=f, h=h, w=w, c=c, b=b) + + def forward(self, x: Tensor) -> Tensor: + return x + self.proj_out(self.attention(x)) + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + self.nin_shortcut = Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x): + h = x + h = self.norm1(h) + h = swish(h) + h = self.conv1(h) + + h = self.norm2(h) + h = swish(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +class DownsampleDCAE(nn.Module): + def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample: bool = True): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2 + assert out_channels % factor == 0 + self.conv = Conv3d(in_channels, out_channels // factor, kernel_size=3, stride=1, padding=1) + + self.add_temporal_downsample = add_temporal_downsample + self.group_size = factor * in_channels // out_channels + + def forward(self, x: Tensor): + r1 = 2 if self.add_temporal_downsample else 1 + h = self.conv(x) + h = rearrange(h, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + shortcut = rearrange(x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2) + return h + shortcut + + +class UpsampleDCAE(nn.Module): + def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: bool = True): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2 + self.conv = Conv3d(in_channels, out_channels * factor, kernel_size=3, stride=1, padding=1) + + self.add_temporal_upsample = add_temporal_upsample + self.repeats = factor * out_channels // in_channels + + def forward(self, x: Tensor): + r1 = 2 if self.add_temporal_upsample else 1 + h = self.conv(x) + h = rearrange(h, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + shortcut = x.repeat_interleave(repeats=self.repeats, dim=1) + shortcut = rearrange(shortcut, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + return h + shortcut + + +class Encoder(nn.Module): + """ + The encoder network of AutoencoderKLConv3D. + """ + + def __init__( + self, + in_channels: int, + z_channels: int, + block_out_channels: tuple[int, ...], + num_res_blocks: int, + ffactor_spatial: int, + ffactor_temporal: int, + downsample_match_channel: bool = True, + ): + super().__init__() + assert block_out_channels[-1] % (2 * z_channels) == 0 + + self.z_channels = z_channels + self.block_out_channels = block_out_channels + self.num_res_blocks = num_res_blocks + + # downsampling + self.conv_in = Conv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) + + self.down = nn.ModuleList() + block_in = block_out_channels[0] + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + down = nn.Module() + down.block = block + + add_spatial_downsample = bool(i_level < np.log2(ffactor_spatial)) + add_temporal_downsample = add_spatial_downsample and bool( + i_level >= np.log2(ffactor_spatial // ffactor_temporal) + ) + if add_spatial_downsample or add_temporal_downsample: + assert i_level < len(block_out_channels) - 1 + block_out = block_out_channels[i_level + 1] if downsample_match_channel else block_in + down.downsample = DownsampleDCAE(block_in, block_out, add_temporal_downsample) + block_in = block_out + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = Conv3d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x: Tensor) -> Tensor: + use_checkpointing = bool(self.training and self.gradient_checkpointing) + + # downsampling + h = self.conv_in(x) + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks): + h = forward_with_checkpointing( + self.down[i_level].block[i_block], h, use_checkpointing=use_checkpointing + ) + if hasattr(self.down[i_level], "downsample"): + h = forward_with_checkpointing(self.down[i_level].downsample, h, use_checkpointing=use_checkpointing) + + # middle + h = forward_with_checkpointing(self.mid.block_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.attn_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.block_2, h, use_checkpointing=use_checkpointing) + + # end + group_size = self.block_out_channels[-1] // (2 * self.z_channels) + shortcut = rearrange(h, "b (c r) f h w -> b c r f h w", r=group_size).mean(dim=2) + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + h += shortcut + return h + + +class Decoder(nn.Module): + """ + The decoder network of AutoencoderKLConv3D. + """ + + def __init__( + self, + z_channels: int, + out_channels: int, + block_out_channels: tuple[int, ...], + num_res_blocks: int, + ffactor_spatial: int, + ffactor_temporal: int, + upsample_match_channel: bool = True, + ): + super().__init__() + assert block_out_channels[0] % z_channels == 0 + + self.z_channels = z_channels + self.block_out_channels = block_out_channels + self.num_res_blocks = num_res_blocks + + # z to block_in + block_in = block_out_channels[0] + self.conv_in = Conv3d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + # upsampling + self.up = nn.ModuleList() + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + + add_spatial_upsample = bool(i_level < np.log2(ffactor_spatial)) + add_temporal_upsample = bool(i_level < np.log2(ffactor_temporal)) + if add_spatial_upsample or add_temporal_upsample: + assert i_level < len(block_out_channels) - 1 + block_out = block_out_channels[i_level + 1] if upsample_match_channel else block_in + up.upsample = UpsampleDCAE(block_in, block_out, add_temporal_upsample) + block_in = block_out + self.up.append(up) + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = Conv3d(block_in, out_channels, kernel_size=3, stride=1, padding=1) + + self.gradient_checkpointing = False + + def forward(self, z: Tensor) -> Tensor: + use_checkpointing = bool(self.training and self.gradient_checkpointing) + + # z to block_in + repeats = self.block_out_channels[0] // (self.z_channels) + h = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1) + + # middle + h = forward_with_checkpointing(self.mid.block_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.attn_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.block_2, h, use_checkpointing=use_checkpointing) + + # upsampling + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks + 1): + h = forward_with_checkpointing(self.up[i_level].block[i_block], h, use_checkpointing=use_checkpointing) + if hasattr(self.up[i_level], "upsample"): + h = forward_with_checkpointing(self.up[i_level].upsample, h, use_checkpointing=use_checkpointing) + + # end + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + return h + + +class AutoencoderKLConv3D(ModelMixin, ConfigMixin): + """ + Autoencoder model with KL-regularized latent space based on 3D convolutions. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int, + out_channels: int, + latent_channels: int, + block_out_channels: tuple[int, ...], + layers_per_block: int, + ffactor_spatial: int, + ffactor_temporal: int, + sample_size: int, + sample_tsize: int, + scaling_factor: float | None = None, + shift_factor: float | None = None, + downsample_match_channel: bool = True, + upsample_match_channel: bool = True, + only_encoder: bool = False, # only build encoder for saving memory + only_decoder: bool = False, # only build decoder for saving memory + ): + super().__init__() + self.ffactor_spatial = ffactor_spatial + self.ffactor_temporal = ffactor_temporal + self.scaling_factor = scaling_factor + self.shift_factor = shift_factor + + # build model + if not only_decoder: + self.encoder = Encoder( + in_channels=in_channels, + z_channels=latent_channels, + block_out_channels=block_out_channels, + num_res_blocks=layers_per_block, + ffactor_spatial=ffactor_spatial, + ffactor_temporal=ffactor_temporal, + downsample_match_channel=downsample_match_channel, + ) + if not only_encoder: + self.decoder = Decoder( + z_channels=latent_channels, + out_channels=out_channels, + block_out_channels=list(reversed(block_out_channels)), + num_res_blocks=layers_per_block, + ffactor_spatial=ffactor_spatial, + ffactor_temporal=ffactor_temporal, + upsample_match_channel=upsample_match_channel, + ) + + # slicing and tiling related + self.use_slicing = False + self.slicing_bsz = 1 + self.use_spatial_tiling = False + self.use_temporal_tiling = False + self.use_tiling_during_training = False + + # only relevant if vae tiling is enabled + self.tile_sample_min_size = sample_size + self.tile_latent_min_size = sample_size // ffactor_spatial + self.tile_sample_min_tsize = sample_tsize + self.tile_latent_min_tsize = sample_tsize // ffactor_temporal + self.tile_overlap_factor = 0.25 + + # use torch.compile for faster encode speed + self.use_compile = False + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (Encoder, Decoder)): + module.gradient_checkpointing = value + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * ( + x / blend_extent + ) + return b + + def spatial_tiled_encode(self, x: torch.Tensor): + """spatial tailing for frames""" + B, C, T, H, W = x.shape + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) # 256 * (1 - 0.25) = 192 + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) # 8 * 0.25 = 2 + row_limit = self.tile_latent_min_size - blend_extent # 8 - 2 = 6 + + rows = [] + for i in range(0, H, overlap_size): + row = [] + for j in range(0, W, overlap_size): + tile = x[:, :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size] + tile = self.encoder(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + moments = torch.cat(result_rows, dim=-2) + return moments + + def temporal_tiled_encode(self, x: torch.Tensor): + """temporal tailing for frames""" + B, C, T, H, W = x.shape + overlap_size = int(self.tile_sample_min_tsize * (1 - self.tile_overlap_factor)) # 64 * (1 - 0.25) = 48 + blend_extent = int(self.tile_latent_min_tsize * self.tile_overlap_factor) # 8 * 0.25 = 2 + t_limit = self.tile_latent_min_tsize - blend_extent # 8 - 2 = 6 + + row = [] + for i in range(0, T, overlap_size): + tile = x[:, :, i : i + self.tile_sample_min_tsize, :, :] + if self.use_spatial_tiling and ( + tile.shape[-1] > self.tile_sample_min_size or tile.shape[-2] > self.tile_sample_min_size + ): + tile = self.spatial_tiled_encode(tile) + else: + tile = self.encoder(tile) + row.append(tile) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + moments = torch.cat(result_row, dim=-3) + return moments + + def spatial_tiled_decode(self, z: torch.Tensor): + """spatial tailing for frames""" + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) # 8 * (1 - 0.25) = 6 + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) # 256 * 0.25 = 64 + row_limit = self.tile_sample_min_size - blend_extent # 256 - 64 = 192 + + rows = [] + for i in range(0, H, overlap_size): + row = [] + for j in range(0, W, overlap_size): + tile = z[:, :, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size] + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=-2) + return dec + + def temporal_tiled_decode(self, z: torch.Tensor): + """temporal tailing for frames""" + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor)) # 8 * (1 - 0.25) = 6 + blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor) # 64 * 0.25 = 16 + t_limit = self.tile_sample_min_tsize - blend_extent # 64 - 16 = 48 + assert 0 < overlap_size < self.tile_latent_min_tsize + + row = [] + for i in range(0, T, overlap_size): + tile = z[:, :, i : i + self.tile_latent_min_tsize, :, :] + if self.use_spatial_tiling and ( + tile.shape[-1] > self.tile_latent_min_size or tile.shape[-2] > self.tile_latent_min_size + ): + decoded = self.spatial_tiled_decode(tile) + else: + decoded = self.decoder(tile) + row.append(decoded) + + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + dec = torch.cat(result_row, dim=-3) + return dec + + def encode(self, x: Tensor, return_dict: bool = True): + """ + Encodes the input by passing through the encoder network. + Support slicing and tiling for memory efficiency. + """ + + def _encode(x): + if self.use_temporal_tiling and x.shape[-3] > self.tile_sample_min_tsize: + return self.temporal_tiled_encode(x) + if self.use_spatial_tiling and ( + x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size + ): + return self.spatial_tiled_encode(x) + + if self.use_compile: + + @torch.compile + def encoder(x): + return self.encoder(x) + + return encoder(x) + return self.encoder(x) + + if len(x.shape) != 5: # (B, C, T, H, W) + x = x[:, :, None] + assert len(x.shape) == 5 # (B, C, T, H, W) + if x.shape[2] == 1: + x = x.expand(-1, -1, self.ffactor_temporal, -1, -1) + else: + assert x.shape[2] != self.ffactor_temporal and x.shape[2] % self.ffactor_temporal == 0 + + if self.use_slicing and x.shape[0] > 1: + if self.slicing_bsz == 1: + encoded_slices = [_encode(x_slice) for x_slice in x.split(1)] + else: + sections = [self.slicing_bsz] * (x.shape[0] // self.slicing_bsz) + if x.shape[0] % self.slicing_bsz != 0: + sections.append(x.shape[0] % self.slicing_bsz) + encoded_slices = [_encode(x_slice) for x_slice in x.split(sections)] + h = torch.cat(encoded_slices) + else: + h = _encode(x) + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + + return AutoencoderKLOutput(latent_dist=posterior) + + def decode(self, z: Tensor, return_dict: bool = True, generator=None): + """ + Decodes the input by passing through the decoder network. + Support slicing and tiling for memory efficiency. + """ + + def _decode(z): + if self.use_temporal_tiling and z.shape[-3] > self.tile_latent_min_tsize: + return self.temporal_tiled_decode(z) + if self.use_spatial_tiling and ( + z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size + ): + return self.spatial_tiled_decode(z) + return self.decoder(z) + + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [_decode(z_slice) for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = _decode(z) + + if z.shape[-3] == 1: + decoded = decoded[:, :, -1:] + + if not return_dict: + return (decoded,) + + return DecoderOutput(sample=decoded) diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py new file mode 100644 index 000000000..360904b5e --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py @@ -0,0 +1,1155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import random +import warnings +from collections import defaultdict +from copy import deepcopy +from typing import Any + +import torch +import torch.nn.functional as F +from diffusers.utils.outputs import BaseOutput +from transformers import AutoTokenizer + +from .hunyuan_image_3_transformer import ImageInfo, JointImageInfo, default + + +class TokenizerEncodeOutput(BaseOutput): + tokens: torch.Tensor | None = None + timestep_scatter_index: torch.Tensor | None = None + guidance_scatter_index: torch.Tensor | None = None + text_slices: list[slice] | None = None + gen_image_slices: list[slice] | None = None + joint_image_slices: list[slice] | None = None + cond_vae_image_slices: list[slice] | None = None + cond_vit_image_slices: list[slice] | None = None + text_mask: torch.Tensor | None = None + gen_image_mask: torch.Tensor | None = None + cond_vae_image_mask: torch.Tensor | None = None + cond_vit_image_mask: torch.Tensor | None = None + real_pos: torch.Tensor | None = None + all_image_slices: list[slice] | None = None + cond_timestep_scatter_index: torch.Tensor | None = None + gen_timestep_scatter_index: torch.Tensor | None = None + + +class Conversation: + roles: list[str] = ["User", "Assistant"] + sep: str = "\n\n" + + +class TokenizerWrapper: + def __init__(self, tokenizer): + if isinstance(tokenizer, str): + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer) + else: + self.tokenizer = tokenizer + + # Define short names + self.bos_token_id = self.tokenizer.bos_token_id + self.eos_token_id = self.tokenizer.eos_token_id + self.pad_token_id = self.tokenizer.pad_token_id + self.boi_token_id = self.tokenizer.convert_tokens_to_ids("") + self.eoi_token_id = self.tokenizer.convert_tokens_to_ids("") + self.img_token_id = self.tokenizer.convert_tokens_to_ids("") + self.cfg_token_id = self.tokenizer.convert_tokens_to_ids("") + self.end_answer_token_id = self.tokenizer.convert_tokens_to_ids("") + self.end_recaption_token_id = self.tokenizer.convert_tokens_to_ids("") + self.ratio_token_offset = self.tokenizer.convert_tokens_to_ids("") + self.special_token_map = self.tokenizer.added_tokens_encoder + + def pad(self, tensor_list, dim=0, pad_val=None): + if pad_val is None: + pad_val = self.pad_token_id + max_len = max([t.shape[dim] for t in tensor_list]) + padded_tensor_list = [] + for t in tensor_list: + if t.shape[dim] < max_len: + assert pad_val is not False, "Not allowed pad." + t = F.pad(t, (0, max_len - t.shape[dim]), value=pad_val) + padded_tensor_list.append(t) + return padded_tensor_list + + def encode(self, *args, **kwargs): + return self.tokenizer.encode(*args, **kwargs) + + def decode(self, *args, **kwargs): + return self.tokenizer.decode(*args, **kwargs) + + def encode_text( + self, + *texts, + uncond_enabled: bool | list[bool] | None = None, + uncond_p: float | None = None, + max_length: int | None = None, + pad: str | None = None, + return_lengths: bool = False, + ) -> list[torch.Tensor]: + r""" + Encode text and image for AR-like model training of the text-to-image/instruction tuning tasks. + Support encode multiple texts at once. Each text can be separately conditioned or unconditioned + based on the uncond_flags and a uniform uncond_p. + ** token is always prepended to the text tokens.** + + Args: + texts (`str` or `List[str]`): + List of texts to be encoded. + uncond_enabled (`bool` or `List[bool]`): + List of flags to indicate whether the text should be unconditioned. + If False, the text will never be unconditioned. + If True, the text will be unconditioned with uncond_p. + uncond_p (`float`): + Probability to the unconditional text. Only works when uncond_enabled is True. + max_length (`int`): + Maximum length of the encoded text. + pad (`str`, *optional*): + Padding method. Can be 'left' or 'right'. + return_lengths (`bool`): + Whether to return the length of each encoded text. + + Returns: + `tuple[torch.Tensor, List[int]]` or `torch.Tensor`: + If `return_lengths` is True, returns a tuple of (encoded_tokens, lengths). + If `return_lengths` is False, returns only the encoded_tokens. + """ + if pad is not None: + assert max_length is not None, "max_length should be provided when pad is not None." + + if uncond_enabled is None: + uncond_enabled = [True] * len(texts) + elif isinstance(uncond_enabled, bool): + uncond_enabled = [uncond_enabled] * len(texts) + if len(uncond_enabled) != len(texts): + print(uncond_enabled, texts) + assert len(uncond_enabled) == len(texts), ( + f"Length of uncond_flags should be equal to the number of texts, " + f"but got {len(uncond_enabled)} and {len(texts)}." + ) + + # Prepare text/uncond tokens + # TODO: If len(texts) > 1, such as instruction + prompt in inpainting, we need to determine how to do uncond. + # Now all texts will be cond or uncond at the same time. + do_uncond_drop = (uncond_p is not None) and (random.random() < uncond_p) + text_tokens, lengths = [], [] + cum_length = 0 + for text, uncond_flag in zip(texts, uncond_enabled): + # If reach the max_length and there still have unencoded texts, give a warning message and break the loop. + if max_length is not None and cum_length >= max_length: + warnings.warn( + f"Text length exceeds the max_length({max_length}). The remaining texts will be ignored: " + f"{text[:80]}..." + ) + break + # Set add_special_tokens=False to avoid adding token in some LLMs. + if isinstance(text, str): + text_token = self.tokenizer.encode(text, add_special_tokens=False) + else: + text_token = text + if uncond_flag and do_uncond_drop: + text_token = [self.cfg_token_id] * len(text_token) + # Cutoff the text by max_length if necessary + if max_length is not None and (cum_length + len(text_token)) > max_length: + text_token = text_token[: max_length - cum_length] + text_tokens.extend(text_token) + lengths.append(len(text_token)) + cum_length += len(text_token) + + # Prepend/Append tokens if applicable + if pad is not None and (pad_length := max_length - len(text_tokens)) > 0: + if pad == "left": + text_tokens = [self.pad_token_id] * pad_length + text_tokens + elif pad == "right": + text_tokens = text_tokens + [self.pad_token_id] * pad_length + else: + raise ValueError(f"Unsupported padding method: {pad}.") + + if return_lengths: + return text_tokens, lengths + return text_tokens + + @staticmethod + def _check_key_number_matched(keys, data): + # Assert keys and token_source are matched + assert set(keys) == set(data.keys()), ( + f"Keys in the template and token source should be matched, but got {set(keys)} and {list(data.keys())}." + ) + key_counts = {k: 0 for k in keys} + for key in keys: + key_counts[key] += 1 + for key, count in key_counts.items(): + assert len(data[key]) == count, ( + f"Number of `{key}` in the token source should be matched with the template, but got " + f"{data[key]}({len(data[key])}) and {count}." + ) + + def _add_image_meta_info_token( + self, + token_seq, + token_count, + extra_token_pos, + add_timestep_token=False, + add_image_shape_token=False, + base_size=None, + ratio_idx=None, + image_type=None, + add_guidance_token=False, + ): + if add_image_shape_token: + token_seq.extend( + [self.special_token_map[f""], self.special_token_map[f""]] + ) + token_count += 2 + if add_timestep_token: + token_seq.extend([self.special_token_map[""]]) + extra_token_pos["timestep"].append(token_count) + if image_type is not None: + if image_type == "gen_image": + extra_token_pos["gen_timestep"].append(token_count) + elif image_type in ["joint_image"]: + extra_token_pos["cond_timestep"].append(token_count) + else: + raise ValueError(f"Unsupported image type: {image_type}.") + token_count += 1 + if add_guidance_token: + token_seq.extend([self.special_token_map[""]]) + extra_token_pos["guidance"].append(token_count) + token_count += 1 + return token_count + + def encode_sequence( + self, + template: str, + token_source: dict[str, list], + total_length=None, + add_timestep_token=False, + add_guidance_token=False, + last_key_only_prefix=False, + add_eos=True, + use_front_boi_token=True, + add_pad=True, + add_bos=True, + drop_last: str | bool = "auto", + add_image_shape_token=False, + ) -> tuple[list, dict[str, list]]: + r""" + Encode a sequence based on the template (e.g., `text-image` for t2i, `text-image-image` for instruction tuning) + and token source. + + Args: + template (`str`): + Template of the sequence. E.g., "text-gen_image" means the sequence is composed of text and an image. + "text-text-gen_image" means the sequence is composed of two sections of text and an image. + token_source (`Dict[str, List]`): + Token source for each key in the template, in order. + - text: List[Dict]. + - gen_image: List[Dict]. + - joint_image: List[Dict]. + total_length (`int`): + Total length of the encoded sequence, include padding tokens. + add_timestep_token (`bool`): + Whether to add timestep token before the image tokens. + (Right after the tokens) + add_guidance_token (`bool`): + Whether to add guidance token before the image tokens. + last_key_only_prefix (`bool`): + Whether to only use the modal prefix in the last key. + add_eos (`bool` or `'auto'`): + Whether to add eos token at the end of the sequence. If True, always add eos token. If 'auto', + add eos token only when the total_length is not reached and the last token is not . + use_front_boi_token (`bool`): + Whether to put the token at the front of iw, ih and timestep tokens. + add_pad (`bool` or `'auto'`): + Whether to add padding tokens to the sequence. + add_bos (`bool`): + Whether to add bos token at the beginning of the sequence. + drop_last (`bool` or `'auto'`): + - If auto, drop last tokens exceeding the total_length if the total_length is provided. If cut point is + in the middle of the image tokens, an error will raised. + - If True, drop last tokens exceeding the total_length. + - If False, keep the last tokens exceeding the total_length, even if the total_length is reached. + add_image_shape_token (`bool`): + Whether to add image shape token before the image tokens. (Right before the token) + + Returns: + `tuple[list, dict]`: A tuple containing: + - token_seq (`list`): Encoded token sequence. + - extra_token_pos (`dict`): Positions of extra tokens. + """ + if last_key_only_prefix: + assert add_eos is not True, "add_eos should not be True when last_key_only_prefix is True." + if drop_last is True and total_length is None: + raise ValueError("total_length should be provided when drop_last is True.") + + keys = template.split("-") + modal_length = len(keys) + index_indicator = {k: 0 for k in token_source} + for k, v in token_source.items(): + assert isinstance(v, (list, tuple)), ( + f"Value of `{k}` in the token source should be a list or tuple, but got {type(v)}." + ) + self._check_key_number_matched(keys, token_source) + + token_seq = [] + token_count = 0 + extra_token_pos = defaultdict(list) + if add_bos: + token_seq.append(self.bos_token_id) + token_count += 1 + # If drop_last is True, we check the token_count on the fly and exit the loop if the total_length is reached. + # This check is only applied to the block tokens. Block tokens mean the tokens that are unsplittable, like + # image tokens. Text tokens are splittable, so we don't need to check the token_count for text. + # If the loop is broken by drop_last, we don't add the eos token at the end because the sequence is not + # complete. + drop_last_break = False + for i, key in enumerate(keys): + source = token_source[key][index_indicator[key]] + if key == "text": + token_seq.extend(source) # text token sequence + extra_token_pos["_start"].append(token_count) + token_count += len(source) + extra_token_pos["_end"].append(token_count - 1) + + elif key == "gen_image": + if isinstance(source, int): + source = {"length": source} + extra_count = ( + 2 + + (1 if source.get("timestep", add_timestep_token) else 0) + + (1 if source.get("guidance", add_guidance_token) else 0) + + (2 if source.get("image_shape", add_image_shape_token) else 0) + ) + if drop_last is True and token_count + extra_count + source["length"] > total_length: + drop_last_break = True + break + if source.get("front_boi", use_front_boi_token): + token_seq.append(self.boi_token_id) + extra_token_pos["boi"].append(token_count) + token_count += 1 + token_count = self._add_image_meta_info_token( + token_seq=token_seq, + token_count=token_count, + extra_token_pos=extra_token_pos, + add_timestep_token=source.get("timestep", add_timestep_token), + add_guidance_token=source.get("guidance", add_guidance_token), + add_image_shape_token=source.get("image_shape", add_image_shape_token), + base_size=source.get("base_size"), + ratio_idx=source.get("ratio_idx"), + image_type=key, + ) + if not source.get("front_boi", use_front_boi_token): + token_seq.append(self.boi_token_id) + extra_token_pos["boi"].append(token_count) + token_count += 1 + if last_key_only_prefix and i == modal_length - 1: + pass # for AR inference + else: + token_seq.extend( + [self.img_token_id] * source["length"] # token number + + [self.eoi_token_id] + ) + extra_token_pos["_start"].append(token_count) + extra_token_pos["_start"].append(token_count) + token_count += source["length"] + extra_token_pos["_end"].append(token_count - 1) + extra_token_pos["_end"].append(token_count - 1) + extra_token_pos["eoi"].append(token_count) + token_count += 1 # + + elif key == "joint_image": + assert isinstance(source["length"], list) and len(source["length"]) == 2, ( + "joint_image length should be a list of two integers" + ) + extra_count = ( + 2 + + 1 + + ( # boi, eoi, joint_img_sep + 1 if source.get("timestep", add_timestep_token) else 0 + ) + + (2 if source.get("image_shape", add_image_shape_token) else 0) + ) + if drop_last is True and token_count + extra_count + sum(source["length"]) > total_length: + drop_last_break = True + break + if source.get("front_boi", use_front_boi_token): + token_seq.append(self.boi_token_id) # Use patched boi for Janus, otherwise useing default + extra_token_pos["boi"].append(token_count) + token_count += 1 + token_count = self._add_image_meta_info_token( + token_seq=token_seq, + token_count=token_count, + extra_token_pos=extra_token_pos, + add_timestep_token=source.get("timestep", add_timestep_token), + add_image_shape_token=source.get("image_shape", add_image_shape_token), + base_size=source.get("base_size"), + ratio_idx=source.get("ratio_idx"), + image_type=key, + ) + if not source.get("front_boi", use_front_boi_token): + token_seq.append(self.boi_token_id) + extra_token_pos["boi"].append(token_count) + token_count += 1 + if last_key_only_prefix and i == modal_length - 1: + pass # for AR inference + else: + token_seq.extend([self.img_token_id] * source["length"][0]) + extra_token_pos["_start"].append(token_count) + extra_token_pos["_start"].append(token_count) + extra_token_pos["_start"].append(token_count) + token_count += source["length"][0] + extra_token_pos["_end"].append(token_count - 1) + extra_token_pos["_end"].append(token_count - 1) + + token_seq.extend([self.special_token_map[""]]) + extra_token_pos["joint_img_sep"].append(token_count) + token_count += 1 + + token_seq.extend([self.img_token_id] * source["length"][1]) + extra_token_pos["_start"].append(token_count) + extra_token_pos["_start"].append(token_count) + token_count += source["length"][1] + extra_token_pos["_end"].append(token_count - 1) + extra_token_pos["_end"].append(token_count - 1) + extra_token_pos["_end"].append(token_count - 1) + + token_seq.extend([self.eoi_token_id]) + extra_token_pos["eoi"].append(token_count) + token_count += 1 # + + else: + raise ValueError(f"Not supported key: {key}") + index_indicator[key] += 1 + + if add_eos is True and not drop_last_break: + # Typically used for t2i task. + token_seq.append(self.eos_token_id) + extra_token_pos["eos"].append(token_count) + token_count += 1 + elif add_eos == "auto" and not drop_last_break: + # Typically used for lm and mmu task. + if token_seq[-1] != self.eos_token_id and (total_length is None or token_count < total_length): + token_seq.append(self.eos_token_id) + extra_token_pos["eos"].append(token_count) + token_count += 1 + + if total_length: + # Check token count and clip sequence if necessary + if token_count > total_length and drop_last: + # Assert clip position is not in the middle of the block-wise tokens (gen_image, joint_image) + for start_key, end_key in [ + ("_start", "_end"), + ("_start", "_end"), + ("_start", "_end"), + ("_start", "_end"), + ]: + if start_key in extra_token_pos and end_key in extra_token_pos: + assert all( + (start > total_length or end + 1 < total_length) + for start, end in zip(extra_token_pos[start_key], extra_token_pos[end_key]) + ), ( + "Clip position should not be in the middle of the image tokens.\n" + f"Below is the text:\n{self._shorten_text(self.tokenizer.decode(token_seq))}" + ) + token_seq = token_seq[:total_length] + + # Pad the sequence if necessary + pad_num = max(0, total_length - len(token_seq)) + if add_pad and pad_num: + token_seq.extend([self.pad_token_id] * pad_num) + extra_token_pos["first_pad"].append(token_count) + + return token_seq, extra_token_pos + + def batch_gen_infer( + self, + infer_fn, + prompt_list: list, + negative_prompt_list: list | None = None, + infer_fn_kwargs_list: list[dict[str, int]] | None = None, + do_classifier_free_guidance=False, + condition_repeat_times: int = 1, + uncondition_repeat_times: int = 1, + ): + r""" + Batch inference for the AR-like model training of the text-to-image/instruction tuning tasks. + + Args: + infer_fn (`callable`): + Inference function to encode the prompt. + prompt_list (`list`): + List of prompts. Each element can be a single prompt or a list of prompts passed to the infer_fn. + negative_prompt_list (`list`, *optional*): + List of negative prompts. Only used when do_classifier_free_guidance is True. If None, will use + token sequence as negative prompt. + infer_fn_kwargs_list (`List[Dict[str, int]]`, *optional*): + List of keyword arguments for the infer_fn. + do_classifier_free_guidance (`bool`): + Whether to do classifier-free guidance. + condition_repeat_times (`int`): + Support multi-condition. + uncondition_repeat_times (`int`): + Support multi-uncondition. + """ + if infer_fn_kwargs_list is None: + infer_fn_kwargs_list = [{} for _ in prompt_list] + + # [n_output, bsz] + cond_results_list = None + uncond_results_list = None + output_type_list = [] + + for prompt_idx, (prompt, infer_fn_kwargs) in enumerate(zip(prompt_list, infer_fn_kwargs_list)): + if not isinstance(prompt, (list, tuple)): + prompt = [prompt] + cond_kwargs = {"uncond_p": 0.0} if do_classifier_free_guidance else {} + results = infer_fn( + *prompt, + **infer_fn_kwargs, + **cond_kwargs, + ) + output_type_list.append((type(results), len(results) if isinstance(results, (list, tuple)) else 1)) + if isinstance(results, dict): + raise ValueError("Make batch on dict is not supported. Please return list or tuple for infer_fn.") + if not isinstance(results, (list, tuple)): + results = (results,) + if cond_results_list is None: + cond_results_list = [[] for _ in results] + uncond_results_list = [[] for _ in results] + for i, result in enumerate(results): + cond_results_list[i].append(result) + + if do_classifier_free_guidance: + if negative_prompt_list is None: + uncond_kwargs = {"uncond_p": 1.0} + uncond_results = infer_fn( + *prompt, + **infer_fn_kwargs, + **uncond_kwargs, + ) + else: + negative_prompt = negative_prompt_list[prompt_idx] + if not isinstance(negative_prompt, (list, tuple)): + negative_prompt = [negative_prompt] + uncond_results = infer_fn( + *negative_prompt, + **infer_fn_kwargs, + ) + if isinstance(uncond_results, TokenizerEncodeOutput): + uncond_results_list.append(uncond_results) + else: + for i, result in enumerate(uncond_results): + uncond_results_list[i].append(result) + + assert all(output_type_list[0] == n for n in output_type_list), ( + f"Number of outputs should be equal for all samples, but got {output_type_list}." + ) + output_type, output_num = output_type_list[0] + + def make_batch(batch_cond_item, batch_uncond_item): + # Process each output item to make batch + first = batch_cond_item[0] # The first element in the batch + if isinstance(first, torch.Tensor): + stacked_item = torch.stack( + self.pad( + batch_cond_item * condition_repeat_times + batch_uncond_item * uncondition_repeat_times, + ) + ) + + elif first is None: + assert all(item is None for item in batch_cond_item + batch_uncond_item), ( + f"The first cond item is None, but some items are not None:\n\n" + f"condition: {batch_cond_item}\n\n" + f"uncondition: {batch_uncond_item}" + ) + stacked_item = None + + elif isinstance(first, (list, tuple)): + # If the output item is a list or tuple, we treat it as a whole, and won't make nested batch any more. + stacked_item = batch_cond_item * condition_repeat_times + batch_uncond_item * uncondition_repeat_times + + elif isinstance(first, TokenizerEncodeOutput): + stacked_item = {} + # Traverse not-None attributes + for key in list(first.keys()): + merged_list = [cond_item[key] for cond_item in batch_cond_item] * condition_repeat_times + [ + uncond_item[key] for uncond_item in batch_uncond_item + ] * uncondition_repeat_times + if isinstance(first[key], torch.Tensor): + if "mask" in key: + pad_val = 0.0 + elif key == "tokens": + pad_val = self.special_token_map[""] + else: + pad_val = False # Should not pad for other tensors + stacked_item[key] = torch.stack(self.pad(merged_list, pad_val=pad_val), dim=0) + elif isinstance(first[key], list): + stacked_item[key] = merged_list + elif first[key] is None: + pass + else: + raise ValueError(f"Unsupported type of {key}: {type(first[key])}.") + stacked_item = TokenizerEncodeOutput(stacked_item) + + else: + raise TypeError(f"Making batch on type {type(first)} is not supported.") + + return stacked_item + + stacked_outputs = [] + for cond_results, uncond_results in zip(cond_results_list, uncond_results_list): + stacked_outputs.append(make_batch(cond_results, uncond_results)) + + if output_type is list: + return stacked_outputs + elif output_type is tuple: + return tuple(stacked_outputs) + elif output_num == 1: + return stacked_outputs[0] + else: + raise ValueError(f"Unsupported output type: {output_type}.") + + @staticmethod + def parse_extra_token_pos(extra_token_pos, prefix, tokens, rng=None): + if rng is None: + rng = slice(None) + image_slices = ( + [ + slice(start, end + 1) + for start, end in zip( + extra_token_pos[f"<{prefix}>_start"][rng], extra_token_pos[f"<{prefix}>_end"][rng] + ) + ] + if f"<{prefix}>_start" in extra_token_pos and f"<{prefix}>_end" in extra_token_pos + else [] + ) + if image_slices: + image_mask = torch.zeros_like(tokens, dtype=torch.bool) + for image_slice in image_slices: + image_mask[image_slice] = True + else: + image_mask = None + return image_slices, image_mask + + def encode_general( + self, + sections: list[dict[str, Any]] | None = None, + max_token_length: int | None = None, + add_eos="auto", + use_text_mask=True, + add_pad="auto", + add_bos=True, + drop_last="auto", + ) -> TokenizerEncodeOutput: + r""" + General encode function to encode a sequence with multiple sections of text and images. + Each section is a dict with a `type` key and other keys depending on the type. + + Supported section types: + + - text: dict with keys: + - text (`str` or `List[int]`): Text to be encoded. Either `text` or `tokens` should be provided. + - tokens (`List[int]`): Pre-encoded text tokens. Either `text` or `tokens` should be provided. + - uncond_enabled (`bool`): Whether to enable uncondition for this text section. + - uncond_p (`float`): Probability to drop the text section for uncondition. + - max_length (`int`): Maximum length of the text section. + - ignore (`bool`): Whether to ignore this text section in the text mask. + - start_offset (`int`): Start offset of the text mask. + - end_offset (`int`): End offset of the text mask. + + - gen_image: dict with keys: + - token_length (`int`): Number of image tokens. + - add_timestep_token (`bool`): Whether to add timestep token before the image tokens. + - add_guidance_token (`bool`): Whether to add guidance token before the image tokens. + - use_front_boi_token (`bool`): Whether to put the token. + - add_image_shape_token (`bool`): Whether to add image shape token before the image tokens. + - base_size (`int`): Base size of the image. + - ratio_idx (`int`): Ratio index of the image. + + - joint_image: dict with keys: + - token_length (`List[int]`): Number of image tokens for the two images. + - add_timestep_token (`bool`): Whether to add timestep token before the image tokens. + - use_front_boi_token (`bool`): Whether to put the token. + - add_image_shape_token (`bool`): Whether to add image shape token before the image tokens. + - base_size (`int`): Base size of the image. + - ratio_idx (`int`): Ratio index of the image. + + Args: + sections (`List[Dict[str, Any]]`): + List of sections to be encoded. + max_token_length (`int`): + Maximum length of the encoded token sequence. + add_eos (`bool` or `'auto'`): + Whether to add eos token at the end of the sequence. + use_text_mask (`bool`): + Whether to generate text mask. + add_pad (`bool` or `'auto'`): + Whether to add padding tokens to the sequence. If True and total_length is not reached, + add padding tokens. + add_bos (`bool`): + Whether to add bos token at the beginning of the sequence. + drop_last (`bool` or `'auto'`): + - If auto, drop last tokens exceeding the total_length if the total_length is provided. + If cut point is in the middle of the image tokens, an error will raised. + - If True, drop last tokens exceeding the total_length. If cut point is in the + middle of the image tokens, all the successive image tokens will be dropped. + - If False, keep the last tokens exceeding the total_length, even if the total_length + is reached. + + Returns: + `TokenizerEncodeOutput`: Encoded token sequence and extra information. + """ + if sections is None: + raise ValueError("sections must be provided.") + template = "-".join([section["type"] for section in sections]) + + sections = deepcopy(sections) + token_source = defaultdict(list) + text_mask_specs = [] + for section in sections: + if section["type"] == "text": + text = self.encode_text( + section["text"] if "text" in section else section["tokens"], + uncond_enabled=section.get("uncond_enabled"), + uncond_p=section.get("uncond_p"), + max_length=section.get("max_length"), + ) + token_source["text"].append(text) + text_mask_specs.append( + dict( + ignore=section.get("ignore", False), + start_offset=section.get("start_offset", 0), + end_offset=section.get("end_offset", 0), + ) + ) + elif section["type"] == "gen_image": + token_source["gen_image"].append( + dict( + length=section["token_length"], + timestep=section.get("add_timestep_token", False), + guidance=section.get("add_guidance_token", False), + front_boi=section.get("use_front_boi_token", False), + image_shape=section.get("add_image_shape_token", False), + base_size=section.get("base_size"), + ratio_idx=section.get("ratio_idx"), + ) + ) + elif section["type"] == "joint_image": + token_source["joint_image"].append( + dict( + length=section["token_length"], + timestep=section.get("add_timestep_token", False), + front_boi=section.get("use_front_boi_token", False), + image_shape=section.get("add_image_shape_token", False), + base_size=section.get("base_size"), + ratio_idx=section.get("ratio_idx"), + ) + ) + else: + raise ValueError(f"Invalid section type: {section['type']}") + + # Combine text and image tokens + full_token_seq, extra_token_pos = self.encode_sequence( + template=template, + token_source=dict(token_source), + total_length=max_token_length, + add_eos=add_eos, + add_pad=add_pad, + add_bos=add_bos, + drop_last=drop_last, + ) + full_seq_token_tensor = torch.tensor(full_token_seq, dtype=torch.long) + + timestep_scatter_index = ( + torch.tensor(extra_token_pos["timestep"], dtype=torch.long) if "timestep" in extra_token_pos else None + ) + guidance_scatter_index = ( + torch.tensor(extra_token_pos["guidance"], dtype=torch.long) if "guidance" in extra_token_pos else None + ) + cond_timestep_scatter_index = ( + torch.tensor(extra_token_pos["cond_timestep"], dtype=torch.long) + if "cond_timestep" in extra_token_pos + else None + ) + gen_timestep_scatter_index = ( + torch.tensor(extra_token_pos["gen_timestep"], dtype=torch.long) + if "gen_timestep" in extra_token_pos + else None + ) + + # Gen image mask + gen_image_slices, gen_image_mask = self.parse_extra_token_pos(extra_token_pos, "img", full_seq_token_tensor) + # Joint image + joint_image_slices, _ = self.parse_extra_token_pos(extra_token_pos, "joint_img", full_seq_token_tensor) + # Conditional vae image + cond_vae_image_slices, cond_vae_image_mask = self.parse_extra_token_pos( + extra_token_pos, "vae_img", full_seq_token_tensor + ) + # Conditional vit image + cond_vit_image_slices, cond_vit_image_mask = self.parse_extra_token_pos( + extra_token_pos, "vit_img", full_seq_token_tensor + ) + # All image slices (gen_image, joint_image) + all_image_slices = ( + [ + slice(start, end + 1) + for start, end in zip(extra_token_pos["_start"], extra_token_pos["_end"]) + ] + if "_start" in extra_token_pos and "_end" in extra_token_pos + else [] + ) + + # Text mask + text_slices = ( + [ + slice(start, end + 1) + for start, end in zip(extra_token_pos["_start"], extra_token_pos["_end"]) + ] + if "_start" in extra_token_pos and "_end" in extra_token_pos + else [] + ) + assert len(text_slices) <= len(text_mask_specs), ( + f"Number of text slices ({len(text_slices)}) should be less than or equal to " + f"number of text mask specs ({len(text_mask_specs)})" + ) + if use_text_mask: + text_mask = torch.zeros_like(full_seq_token_tensor, dtype=torch.float32) + for text_slice, mask_spec in zip(text_slices, text_mask_specs): + if not mask_spec["ignore"]: + real_slice = slice( + text_slice.start + mask_spec["start_offset"], text_slice.stop + mask_spec["end_offset"] + ) + text_mask[real_slice] = 1.0 + else: + text_mask = None + + # real_pos is the first position of the token + real_pos = torch.tensor(extra_token_pos.get("first_pad", [full_seq_token_tensor.shape[0]]), dtype=torch.long) + + return TokenizerEncodeOutput( + tokens=full_seq_token_tensor, + timestep_scatter_index=timestep_scatter_index, + guidance_scatter_index=guidance_scatter_index, + text_slices=text_slices, + gen_image_slices=gen_image_slices, + joint_image_slices=joint_image_slices, + cond_vae_image_slices=cond_vae_image_slices, + cond_vit_image_slices=cond_vit_image_slices, + text_mask=text_mask, + gen_image_mask=gen_image_mask, + cond_vae_image_mask=cond_vae_image_mask, + cond_vit_image_mask=cond_vit_image_mask, + real_pos=real_pos, + all_image_slices=all_image_slices, + cond_timestep_scatter_index=cond_timestep_scatter_index, + gen_timestep_scatter_index=gen_timestep_scatter_index, + ) + + def get_cot_sections(self, cot_text, uncond_kwargs, cot_max_length=None, drop_think=False): + if not cot_text: # None or empty + return [] + if "" in cot_text and "" in cot_text: + before_think_sec = cot_text.split("")[0] + after_think_sec = cot_text.split("")[1] + think_sec = cot_text.split("")[1].split("")[0] + return ( + self.get_cot_sections(before_think_sec, uncond_kwargs, drop_think=drop_think) + + ( + [ + dict(type="text", text=""), + dict(type="text", text=think_sec, max_length=cot_max_length, **uncond_kwargs), + dict(type="text", text=""), + ] + if not drop_think + else [] + ) + + self.get_cot_sections(after_think_sec, uncond_kwargs, drop_think=drop_think) + ) + + if "" in cot_text and "" in cot_text: + before_recaption_sec = cot_text.split("")[0] + after_recaption_sec = cot_text.split("")[1] + recaption_sec = cot_text.split("")[1].split("")[0] + return ( + self.get_cot_sections(before_recaption_sec, uncond_kwargs, drop_think=drop_think) + + [ + dict(type="text", text=""), + dict(type="text", text=recaption_sec, max_length=cot_max_length, **uncond_kwargs), + dict(type="text", text=""), + ] + + self.get_cot_sections(after_recaption_sec, uncond_kwargs, drop_think=drop_think) + ) + + return [ + dict(type="text", text=cot_text, **uncond_kwargs), + ] + + def apply_general_template( + self, + message_list, + max_length=None, + add_assistant_prefix=False, + answer="auto", + bot_task="auto", + sequence_template="instruct", + uncond_p=0.0, + cfg_factor=1, + batchify=False, + image_base_size=1024, + drop_think=False, + ): + # If cfg_factor > 1, we need to repeat the unconditioned part + if batchify: + assert isinstance(message_list[0], list), ( + f"When batchify is True, message_list should be a list of list, but got [{type(message_list[0])}, ...]." + ) + return self.batch_gen_infer( + infer_fn=self.apply_general_template, + prompt_list=[[]], + infer_fn_kwargs_list=[ + dict( + message_list=message_list_i, + max_length=max_length, + add_assistant_prefix=add_assistant_prefix, + answer=answer, + bot_task=bot_task, + sequence_template=sequence_template, + image_base_size=image_base_size, + drop_think=drop_think, + ) + for message_list_i in message_list + ], + do_classifier_free_guidance=cfg_factor > 1, + condition_repeat_times=1, + uncondition_repeat_times=cfg_factor - 1, + ) + + conv = Conversation() + uncond_kwargs = dict(uncond_enabled=uncond_p == 1.0, uncond_p=uncond_p) + + def process_successive_message( + _message_list, _cur_message_idx, role, prefix, suffix, answer_prefix="", answer_suffix="" + ): + _sub_sections = [] + while _cur_message_idx < len(message_list) and _message_list[_cur_message_idx]["role"] == role: + message = _message_list[_cur_message_idx] + if message["type"] == "text": + text = message["content"] + if role == "system": + _sub_sections.append(dict(type="text", text=text)) + elif role == "assistant": + if ("" in text and "" in text) or ( + "" in text and "" in text + ): + _sub_sections.extend(self.get_cot_sections(text, uncond_kwargs, drop_think=drop_think)) + else: + _sub_sections.append(dict(type="text", text=text, **uncond_kwargs)) + else: + _sub_sections.append( + dict(type="text", text=f"{answer_prefix}{text}{answer_suffix}", **uncond_kwargs) + ) + elif message["type"] == "gen_image": + info = message["content"] + assert isinstance(info, ImageInfo), f"Expected ImageInfo, but got {type(info)}" + if role == "assistant": + _sub_sections.append(dict(type="text", text=answer_prefix)) + _sub_sections.append(dict(type=message["type"], **info.meta_info)) + if role == "assistant": + _sub_sections.append(dict(type="text", text=answer_suffix)) + elif message["type"] == "joint_image": + info = message["content"] + assert isinstance(info, JointImageInfo), f"Expected JointImageInfo, but got {type(info)}" + _sub_sections.append(dict(type=message["type"], **info.meta_info)) + else: + raise ValueError(f"Unknown message type: {message['type']}") + _cur_message_idx += 1 + if len(_sub_sections) > 0: + # Add role prefix and suffix + _sub_sections.insert(0, dict(type="text", text=prefix)) + _sub_sections.append(dict(type="text", text=suffix)) + return _sub_sections, _cur_message_idx + + # Define assistant prefix and suffix + if (answer == "auto" and sequence_template == "instruct") or answer is True: + answer_prefix, answer_suffix = "", "" + else: + answer_prefix, answer_suffix = "", "" + if sequence_template == "pretrain": + system_suffix = "" + user_prefix = "" + user_suffix = "" + bot_prefix = "" + bot_suffix = "" + else: + system_suffix = f"{conv.sep}" + user_prefix = f"{conv.roles[0]}: " + user_suffix = f"{conv.sep}" + bot_prefix = f"{conv.roles[1]}: " + bot_suffix = f"{conv.sep}" + + # Process successive user and assistant messages + sections = [] + cur_message_idx = 0 + final_role = None + while cur_message_idx < len(message_list): + # Process successive system messages + sub_sections, cur_message_idx = process_successive_message( + message_list, cur_message_idx, role="system", prefix="", suffix=system_suffix + ) + # Add to the template and sections + sections.extend(sub_sections) + if len(sub_sections) > 0: + final_role = "system" + + # Process successive user messages + sub_sections, cur_message_idx = process_successive_message( + message_list, cur_message_idx, role="user", prefix=user_prefix, suffix=user_suffix + ) + # Add to the template and sections + sections.extend(sub_sections) + if len(sub_sections) > 0: + final_role = "user" + + # Process successive assistant messages + sub_sections, cur_message_idx = process_successive_message( + message_list, + cur_message_idx, + role="assistant", + prefix=bot_prefix, + suffix=bot_suffix, + answer_prefix=answer_prefix, + answer_suffix=answer_suffix, + ) + # Add to the template and sections + sections.extend(sub_sections) + if len(sub_sections) > 0: + final_role = "assistant" + + if add_assistant_prefix: + if final_role == "assistant": + # Avoid adding prefix twice + _bot_prefix = "" + # Remove the final bot_suffix + if len(sections) > 0 and sections[-1]["type"] == "text" and sections[-1]["text"] == bot_suffix: + sections = sections[:-1] + else: + _bot_prefix = bot_prefix + # We can add special tokens for the bot latest message according to different tasks + bot_response_prefix = dict( + auto=_bot_prefix, + image="", + think=f"{_bot_prefix}", + recaption=f"{_bot_prefix}", + img_ratio=f"{_bot_prefix}{answer_prefix}", + )[bot_task] + sections.append(dict(type="text", text=bot_response_prefix)) + + output = self.encode_general( + sections=sections, + use_text_mask=False, + add_eos=False, + add_pad=False, + ) + + if max_length is not None: + if output.tokens.shape[-1] > max_length: + raise ValueError( + f"Encoded token length {output.tokens.shape[-1]} exceeds max_length {max_length}.\n" + f"Please set a larger max_length or check the input messages:\n{message_list}" + ) + + return output, sections + + def apply_chat_template( + self, + batch_prompt: list[str] | None = None, + batch_message_list: list[list[dict[str, Any]]] | None = None, + mode: str = "gen_text", + batch_gen_image_info: list[ImageInfo] | None = None, + batch_cond_image_info: list[JointImageInfo] | list[list[JointImageInfo]] | None = None, + batch_system_prompt: list[str] | None = None, + batch_cot_text: list[str] | None = None, + max_length: int | None = None, + bot_task: str = "auto", # auto/image/think/recaption/img_ratio + image_base_size: int = 1024, + sequence_template: str = "pretrain", + cfg_factor: int = 1, + add_assistant_prefix: bool | None = None, + drop_think: bool = False, + ) -> dict[str, Any]: + assert bot_task in ["image", "auto", "think", "recaption", "img_ratio"], ( + f"bot_task should be one of ['image', 'auto', 'think', 'recaption', 'img_ratio'], but got {bot_task}." + ) + + if batch_message_list is None: + # Simple text-to-image or text-cot-to-image task + batch_size = len(batch_prompt) + + # Batchify inputs + if not isinstance(batch_system_prompt, list): + batch_system_prompt = [batch_system_prompt] * batch_size + if not isinstance(batch_gen_image_info, list): + batch_gen_image_info = [batch_gen_image_info] * batch_size + if batch_cot_text is not None: + assert len(batch_cot_text) == batch_size, ( + f"batch_cot_text should have the same length as batch_size ({batch_size}), " + f"but got {len(batch_cot_text)}." + ) + else: + batch_cot_text = [None] * batch_size + if batch_cond_image_info is not None: + assert len(batch_cond_image_info) == batch_size, ( + f"batch_cond_image_info should have the same length as batch_size ({batch_size}), " + f"but got {len(batch_cond_image_info)}." + ) + batch_cond_image_info = [ + cond_image_info if isinstance(cond_image_info, list) else [cond_image_info] + for cond_image_info in batch_cond_image_info + ] + else: + batch_cond_image_info = [[] for _ in range(batch_size)] + + # Convert single round materials into standard message list + batch_message_list = [] + for prompt, system_prompt, cot_text, gen_image_info, cond_image_info_list in zip( + batch_prompt, + batch_system_prompt, + batch_cot_text, + batch_gen_image_info, + batch_cond_image_info, + ): + message_list = [] + # 1. system prompt section + if system_prompt: + message_list.append(dict(role="system", type="text", content=system_prompt, context_type="str")) + # 2. user inputs sections + # 2.1 image inputs + if len(cond_image_info_list) > 0: + message_list.extend( + [ + dict(role="user", type="joint_image", content=cond_image_info, context_type="image_info") + for cond_image_info in cond_image_info_list + ] + ) + # 2.2 text inputs + message_list.append(dict(role="user", type="text", content=prompt, context_type="str")) + # 3. assistant answer sections + if cot_text is not None: + message_list.append(dict(role="assistant", type="text", content=cot_text, context_type="str")) + if mode == "gen_image": + message_list.append( + dict(role="assistant", type="gen_image", content=gen_image_info, context_type="image_info") + ) + batch_message_list.append(message_list) + + output, sections = self.apply_general_template( + message_list=batch_message_list, + max_length=max_length, + add_assistant_prefix=default(add_assistant_prefix, mode != "gen_image"), + bot_task=bot_task, + sequence_template=sequence_template, + cfg_factor=cfg_factor, + batchify=True, + image_base_size=image_base_size, + drop_think=drop_think, + ) + return dict(output=output, sections=sections) diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py new file mode 100644 index 000000000..adb582741 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py @@ -0,0 +1,2796 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import glob +import inspect +import logging +import math +import os +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from typing import Any, cast + +import numpy as np +import regex as re +import torch +import torch.nn.functional as F +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils.outputs import BaseOutput +from diffusers.utils.torch_utils import randn_tensor +from einops import rearrange +from PIL import Image +from safetensors.torch import load_file +from torch import nn +from torchvision import transforms +from transformers import PretrainedConfig, Siglip2ImageProcessorFast +from transformers.cache_utils import Cache +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, +) +from transformers.modeling_utils import PreTrainedModel +from vllm.config import CacheConfig +from vllm.distributed import ( + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.fused_moe import SharedFusedMoE +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.utils import ( + PPMissingLayer, + is_pp_missing_parameter, + make_layers, +) +from vllm.v1.attention.backend import AttentionType + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.distributed.parallel_state import get_pp_group +from vllm_omni.diffusion.layers.rope import RotaryEmbedding + +logger = logging.getLogger(__name__) + + +def _is_moe(config: PretrainedConfig) -> bool: + num_experts = getattr(config, "num_experts", None) + if isinstance(num_experts, int): + return num_experts > 1 + if isinstance(num_experts, list) and num_experts: + # Ensure all elements are integers before calling max. + if all(isinstance(e, int) for e in num_experts): + return max(num_experts) > 1 + else: + return False + return False + + +def _get_cla_factor(config: PretrainedConfig) -> int: + if not getattr(config, "use_cla", False): + return 1 + return getattr(config, "cla_share_factor", 1) + + +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +) -> tuple[list[int], int]: + """ + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +def real_batched_index_select(t, dim, idx): + """index_select for batched index and batched t""" + assert t.ndim >= 2 and idx.ndim >= 2, f"{t.ndim=} {idx.ndim=}" + assert len(t) == len(idx), f"{len(t)=} != {len(idx)=}" + return torch.stack([torch.index_select(t[i], dim - 1, idx[i]) for i in range(len(t))]) + + +def conv_nd(dims, *args, **kwargs): # noqa: N802 + """ + Create a 1D, 2D, or 3D convolution module. + """ + if dims == 1: + return nn.Conv1d(*args, **kwargs) + elif dims == 2: + return nn.Conv2d(*args, **kwargs) + elif dims == 3: + return nn.Conv3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") + + +def normalization(channels, **kwargs): + """ + Make a standard normalization layer. + :param channels: number of input channels. + :return: a nn.Module for normalization. + """ + return nn.GroupNorm(32, channels, **kwargs) + + +def linear(*args, **kwargs): + """ + Create a linear module. + """ + return nn.Linear(*args, **kwargs) + + +def zero_module(module): + """ + Zero out the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().zero_() + return module + + +def _to_tuple(x, dim=2): + if isinstance(x, int): + return (x,) * dim + elif len(x) == dim: + return x + else: + raise ValueError(f"Expected length {dim} or int, but got {x}") + + +def get_meshgrid_nd(start, *args, dim=2): # noqa: N802 + if len(args) == 0: + # start is grid_size + num = _to_tuple(start, dim=dim) + start = (0,) * dim + stop = num + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + start = _to_tuple(start, dim=dim) + stop = _to_tuple(args[0], dim=dim) + num = [stop[i] - start[i] for i in range(dim)] + # assert num are all integers + num_int = [int(x) for x in num] + assert (torch.tensor(num) == torch.tensor(num_int)).all(), f"num should be int, but got {num}" + num = num_int + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 + stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 + num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False) + axis_grid = [] + for i in range(dim): + a, b, n = start[i], stop[i], num[i] + g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n] + axis_grid.append(g) + grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [H, W] + grid = torch.stack(grid, dim=0) # [dim, H, W] + + return grid + + +def build_2d_rope( + seq_len: int, + n_elem: int, + image_infos: list[tuple[slice, tuple[int, int]]] | None = None, + device: torch.device | None = None, + base: int = 10000, + base_rescale_factor: float = 1.0, + return_all_pos: bool = False, +): + assert n_elem % 4 == 0, f"n_elem must be divisible by 4, but got {n_elem}." + + # theta + if base_rescale_factor != 1.0: + base *= base_rescale_factor ** (n_elem / (n_elem - 2)) + theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, device=device).float() / n_elem)) + theta = theta.reshape(1, n_elem // 4, 2) # [1, half_d, 2] + + # position indices + if image_infos is None: + image_infos = [] + + image_infos_list = [image_infos] + sample_seq_lens = [seq_len] + + # Prepare position indices for each sample + x_sections = [] + y_sections = [] + for sample_id, sample_image_infos in enumerate(image_infos_list): + last_pos = 0 + for sec_slice, (h, w) in sample_image_infos: + L = sec_slice.start # start from 0, so image_slice.start is just L + # previous text + if last_pos < L: + y_sections.append(torch.arange(last_pos, L)) + x_sections.append(torch.arange(last_pos, L)) + elif h is None: + # Interleave data has overlapped positions for tokens. + y_sections.append(torch.arange(sec_slice.start, sec_slice.stop)) + x_sections.append(torch.arange(sec_slice.start, sec_slice.stop)) + continue + else: + # Interleave data has overlapped positions for noised image and the successive clean image, + # leading to last_pos (= last text end L + noise w * h) > L (last text end L). + pass + # current image + beta_y = L + (w * h - h) / 2 + beta_x = L + (w * h - w) / 2 + grid = get_meshgrid_nd((beta_y, beta_x), (beta_y + h, beta_x + w)) # [2, h, w] # noqa: N802 + grid = grid.reshape(2, -1) # (y, x) + y_sections.append(grid[0]) + x_sections.append(grid[1]) + # step + last_pos = L + w * h + # final text + y_sections.append(torch.arange(last_pos, sample_seq_lens[sample_id])) + x_sections.append(torch.arange(last_pos, sample_seq_lens[sample_id])) + + x_pos = torch.cat(x_sections).long() + y_pos = torch.cat(y_sections).long() + # If there are overlap positions, we need to remove them. + x_pos = x_pos[:seq_len] + y_pos = y_pos[:seq_len] + all_pos = torch.stack((y_pos, x_pos), dim=1).unsqueeze(1).to(device) # [seq_len, 1, 2] + + # calc rope + idx_theta = (all_pos * theta).reshape(all_pos.shape[0], n_elem // 2).repeat(1, 1) + + cos = torch.cos(idx_theta) + sin = torch.sin(idx_theta) + + if return_all_pos: + return cos, sin, all_pos + + return cos, sin + + +def build_batch_2d_rope( + seq_len: int, + n_elem: int, + image_infos: list[list[tuple[slice, tuple[int, int]]]] | None = None, + device: torch.device | None = None, + base: int = 10000, + base_rescale_factor: float = 1.0, + return_all_pos: bool = False, +): + cos_list, sin_list, all_pos_list = [], [], [] + if image_infos is None: + image_infos = [None] + for i, image_info in enumerate(image_infos): + res = build_2d_rope( + seq_len, + n_elem, + image_infos=image_info, + device=device, + base=base, + base_rescale_factor=base_rescale_factor, + return_all_pos=return_all_pos, + ) + if isinstance(res, tuple) and len(res) == 3: + cos, sin, all_pos = res + elif isinstance(res, tuple) and len(res) == 2: + cos, sin = res + all_pos = None + else: + raise ValueError( + "build_2d_rope must return a tuple of length 2 or 3 " + f"when return_all_pos={return_all_pos}, got: {type(res)} with length " + f"{len(res) if isinstance(res, tuple) else 'N/A'}" + ) + cos_list.append(cos) + sin_list.append(sin) + all_pos_list.append(all_pos) + stacked_cos = torch.stack(cos_list, dim=0) + stacked_sin = torch.stack(sin_list, dim=0) + if return_all_pos: + return stacked_cos, stacked_sin, all_pos_list + + return stacked_cos, stacked_sin + + +def get_full_state_dict(model_path): + files = glob.glob(os.path.join(model_path, "*.safetensors")) + full_sd = {} + for f in files: + full_sd.update(load_file(f)) + return full_sd + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb( + q, k, cos, sin, position_ids=None, unsqueeze_dim=1, mla=False +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): + The query tensor. + k (`torch.Tensor`): + The key tensor. + cos (`torch.Tensor`): + The cosine part of the rotary embedding. + sin (`torch.Tensor`): + The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + The position indices of the tokens corresponding to the query and key tensors. For example, this can be + used to pass offsetted position ids when working with a KV-cache. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + + Returns: + `Tuple[torch.Tensor, torch.Tensor]`: + A tuple comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + if position_ids is not None: + cos = cos[position_ids] + sin = sin[position_ids] + + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + + if mla: + b, h, s, d = q.shape + q = q.reshape(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + b, h, s, d = k.shape + k = k.reshape(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def default(value, default_value): + return value if value is not None else default_value + + +class Resolution: + def __init__(self, size, *args): + if isinstance(size, str): + if "x" in size: + size = size.split("x") + size = (int(size[0]), int(size[1])) + else: + size = int(size) + if len(args) > 0: + size = (size, args[0]) + if isinstance(size, int): + size = (size, size) + + self.h = self.height = size[0] + self.w = self.width = size[1] + self.r = self.ratio = self.height / self.width + + def __getitem__(self, idx): + if idx == 0: + return self.h + elif idx == 1: + return self.w + else: + raise IndexError(f"Index {idx} out of range") + + def __str__(self): + return f"{self.h}x{self.w}" + + +class ResolutionGroup: + def __init__(self, base_size=None, step=None, align=1): + self.align = align + self.base_size = base_size + assert base_size % align == 0, f"base_size {base_size} is not divisible by align {align}" + if base_size is not None and not isinstance(base_size, int): + raise ValueError(f"base_size must be None or int, but got {type(base_size)}") + if step is None: + step = base_size // 16 + if step is not None and step > base_size // 2: + raise ValueError(f"step must be smaller than base_size // 2, but got {step} > {base_size // 2}") + + self.step = step + self.data = self._calc_by_step() + + self.ratio = np.array([x.ratio for x in self.data]) + self.attr = ["" for _ in range(len(self.data))] + self.prefix_space = 0 + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx] + + def __repr__(self): + prefix = self.prefix_space * " " + prefix_close = (self.prefix_space - 4) * " " + res_str = f"ResolutionGroup(base_size={self.base_size}, step={self.step}, data=" + attr_maxlen = max([len(x) for x in self.attr] + [5]) + res_str += ( + f"\n{prefix}ID: height width ratio {' ' * max(0, attr_maxlen - 4)}count h/16 w/16 tokens\n{prefix}" + ) + res_str += ("\n" + prefix).join( + [ + f"{i:2d}: ({x.h:4d}, {x.w:4d}) {self.ratio[i]:.4f} {self.attr[i]:>{attr_maxlen}s} " + f"({x.h // 16:3d}, {x.w // 16:3d}) {x.h // 16 * x.w // 16:6d}" + for i, x in enumerate(self.data) + ] + ) + res_str += f"\n{prefix_close})" + return res_str + + def _calc_by_step(self): + assert self.align <= self.step, f"align {self.align} must be smaller than step {self.step}" + + min_height = self.base_size // 2 + min_width = self.base_size // 2 + max_height = self.base_size * 2 + max_width = self.base_size * 2 + + resolutions = [Resolution(self.base_size, self.base_size)] + + cur_height, cur_width = self.base_size, self.base_size + while True: + if cur_height >= max_height and cur_width <= min_width: + break + + cur_height = min(cur_height + self.step, max_height) + cur_width = max(cur_width - self.step, min_width) + resolutions.append(Resolution(cur_height // self.align * self.align, cur_width // self.align * self.align)) + + cur_height, cur_width = self.base_size, self.base_size + while True: + if cur_height <= min_height and cur_width >= max_width: + break + + cur_height = max(cur_height - self.step, min_height) + cur_width = min(cur_width + self.step, max_width) + resolutions.append(Resolution(cur_height // self.align * self.align, cur_width // self.align * self.align)) + + resolutions = sorted(resolutions, key=lambda x: x.ratio) + + return resolutions + + def get_target_size(self, width, height): + ratio = height / width + idx = np.argmin(np.abs(self.ratio - ratio)) + reso = self.data[idx] + return reso.w, reso.h + + def get_base_size_and_ratio_index(self, width, height): + ratio = height / width + idx = np.argmin(np.abs(self.ratio - ratio)) + return self.base_size, idx + + +class ImageInfo: + """Class to store image information for processing and generation.""" + + def __init__( + self, + image_type: str = None, + image_tensor: torch.Tensor = None, + image_width: int = None, + image_height: int = None, + token_width: int = None, + token_height: int = None, + image_token_length: int = None, + base_size: int = None, + ratio_index: int = None, + **kwargs, + ): + self.image_type = image_type + self.image_tensor = image_tensor + self.image_width = image_width + self.w = image_width + self.image_height = image_height + self.h = image_height + self.token_width = token_width + self.tk_w = token_width + self.token_height = token_height + self.tk_h = token_height + self.image_token_length = default( + image_token_length, + token_width * token_height if token_width is not None and token_height is not None else None, + ) + self.base_size = base_size + self.ratio_index = ratio_index + + self.add_timestep_token = kwargs.get("add_timestep_token", True) + self.add_guidance_token = kwargs.get("add_guidance_token", False) + self.use_front_boi_token = kwargs.get("use_front_boi_token", True) + self.add_image_shape_token = kwargs.get("add_image_shape_token", True) + + def __getitem__(self, key: str) -> Any: + """Allow dictionary-like access to attributes.""" + if hasattr(self, key): + return getattr(self, key) + raise KeyError(f"Key '{key}' not found in ImageInfo") + + def __setitem__(self, key: str, value: Any) -> None: + """Allow dictionary-like assignment to attributes.""" + if hasattr(self, key): + setattr(self, key, value) + else: + raise KeyError(f"Key '{key}' not found in ImageInfo") + + def __contains__(self, key: str) -> bool: + """Check if the key exists in the ImageInfo object.""" + return hasattr(self, key) + + def __repr__(self): + return ( + f"ImageInfo(image_type={self.image_type}, image_tensor={self.image_tensor}, " + f"image_width={self.image_width}, image_height={self.image_height}, " + f"token_width={self.token_width}, token_height={self.token_height}, " + f"image_token_length={self.image_token_length}, " + f"base_size={self.base_size}, ratio_index={self.ratio_index}" + ) + + @property + def meta_info(self): + # Used for image sections of tkwrapper.encode_general() + if self.image_type in ["vae", "gen_image"]: + return dict( + token_length=self.image_token_length, + add_timestep_token=self.add_timestep_token, + add_guidance_token=self.add_guidance_token, + use_front_boi_token=self.use_front_boi_token, + add_image_shape_token=self.add_image_shape_token, + base_size=self.base_size, + ratio_idx=self.ratio_index, + # for rope 2d + token_height=self.token_height, + token_width=self.token_width, + # for bc + image_height=self.image_height, + image_width=self.image_width, + ) + else: + raise ValueError(f"Unknown image type '{self.image_type}'") + + +class ImageTensor(torch.Tensor): + # This class is just for type hinting purposes. Attribute `i` should be defined + # as an instance attribute of the torch.Tensor instance, like: tensor.i = ImageInfo(...) + i: ImageInfo + vision_encoder_kwargs: dict + + +class JointImageInfo: + def __init__(self, vae_image_info: ImageInfo, vision_image_info: ImageInfo, vision_encoder_kwargs: dict = None): + self.vae_image_info = vae_image_info + self.vision_image_info = vision_image_info + self.vision_encoder_kwargs = vision_encoder_kwargs + + # Define key attributes to align with ImageInfo for uniformity + self.image_type = "joint_image" + self.image_token_length = vae_image_info.image_token_length + vision_image_info.image_token_length + + self.add_timestep_token = vae_image_info.add_timestep_token + self.use_front_boi_token = vae_image_info.use_front_boi_token + self.add_image_shape_token = vae_image_info.add_image_shape_token + + def __repr__(self): + return f"JointImageInfo(vae_image={self.vae_image_info}, vision_image={self.vision_image_info})" + + @property + def meta_info(self): + # Used for image sections of tkwrapper.encode_general() + return dict( + token_length=[self.vae_image_info.image_token_length, self.vision_image_info.image_token_length], + add_timestep_token=self.add_timestep_token, + use_front_boi_token=self.use_front_boi_token, + add_image_shape_token=self.add_image_shape_token, + base_size=self.vae_image_info.base_size, + ratio_idx=self.vae_image_info.ratio_index, + # for rope 2d + token_height=[self.vae_image_info.token_height, self.vision_image_info.token_height], + token_width=[self.vae_image_info.token_width, self.vision_image_info.token_width], + # for bc + image_height=[self.vae_image_info.image_height, self.vision_image_info.image_height], + image_width=[self.vae_image_info.image_width, self.vision_image_info.image_width], + ) + + @property + def num_special_tokens(self): + return ( + 2 # + + + (1 if self.add_timestep_token else 0) + + (2 if self.add_image_shape_token else 0) + + 1 # + ) + + def copy(self, copy_image_tensor=True): + if copy_image_tensor and ( + self.vae_image_info.image_tensor is None or self.vision_image_info.image_tensor is None + ): + raise ValueError("image_tensor is None, cannot copy") + return JointImageInfo( + self.vae_image_info.copy(copy_image_tensor), + self.vision_image_info.copy(copy_image_tensor), + self.vision_encoder_kwargs, + ) + + def zeros_(self): + self.vae_image_info.zeros_() + self.vision_image_info.zeros_() + + +class JointImage: + def __init__(self, vae_image: ImageTensor, vision_image: ImageTensor): + self.vae_image = vae_image + self.vision_image = vision_image + self.i = JointImageInfo(vae_image.i, vision_image.i) + + +class Config: + def __init__(self, config): + if config is not None: + for key, value in config.items(): + setattr(self, key, value) + + +class LightProjector(nn.Module): + def __init__(self, config): + config = Config(config) + super().__init__() + + if config.projector_type == "linear": + modules = nn.Linear(config.input_dim, config.n_embed) + + elif config.projector_type == "mlp_gelu": + modules = [nn.Linear(config.input_dim, config.n_embed)] + for _ in range(1, config.depth): + modules.append(nn.GELU()) + modules.append(nn.Linear(config.n_embed, config.n_embed)) + modules = nn.Sequential(*modules) + + else: + raise ValueError(f"Unknown projector type: {config.projector_type}") + + self.layers = modules + + +class HunYuanRotary2DEmbedder: + r""" + A RoPE wrapper specifically designed for HunYuan-Image attention. + + This class implements Rotary Position Embedding (RoPE) specifically optimized for + the HunYuan-Image model's attention mechanism. It handles the application of rotary + embeddings to query and key tensors with support for custom position embeddings + and attention metadata. + + Example: + ```python + embedder = HunYuanRotaryEmbedder(num_heads=num_h, num_kv_heads=num_kv, head_dim=h_d) + q, k = embedder(q, k, hidden_states, custom_pos_emb, first_step, attn_meta) + ``` + + Args: + num_heads (`int`): + Number of attention heads in the model. + num_kv_heads (`int`): + Number of key-value heads (for grouped-query attention). + head_dim (`int`): + Dimension of each attention head. + + Methods: + __call__: Applies rotary position embedding to query and key tensors. + """ + + def __init__(self, num_heads: int, num_kv_heads: int, head_dim: int): + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = head_dim + self.custom_pos_emb: tuple[torch.Tensor, torch.Tensor] | None = None + self.rope = RotaryEmbedding(is_neox_style=True) + + def _prepare_cos_sin( + self, + custom_pos_emb: tuple[torch.Tensor, torch.Tensor], + first_step: bool, + device: torch.device, + ): + """Returns cos/sin on the target device based on first_step and caching strategy.""" + if first_step: + cos_input, sin_input = custom_pos_emb + cos = cos_input.to(device) + sin = sin_input.to(device) + self.custom_pos_emb = None + else: + if self.custom_pos_emb is None: + cos_input, sin_input = custom_pos_emb + cos = cos_input.to(device) + sin = sin_input.to(device) + self.custom_pos_emb = (cos, sin) + else: + cos, sin = self.custom_pos_emb + return cos, sin + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + hidden_states: torch.Tensor, + custom_pos_emb: tuple[torch.Tensor, torch.Tensor], + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + hidden_states_shape = hidden_states.shape + hidden_states = hidden_states.reshape(-1, hidden_states_shape[-1]) + if kwargs.get("mode", "gen_text") != "gen_image": + return q, k + + first_step = kwargs.get("first_step", False) + device = q.device + # 1. Prepare cos/sin + cos, sin = self._prepare_cos_sin(custom_pos_emb, first_step, device) + + # 2. Shape validation + query_lens: list[int] = kwargs.get("query_lens") + bs = len(query_lens) + q_len = query_lens[0] + assert hidden_states.shape[0] == bs * q_len, f"{hidden_states.shape[0]} != {bs * q_len}" + + # 3. Reshape + transpose for apply_rotary_pos_emb + # Assume q shape [B*L, H*D] -> [2, L, H, D] -> [2, H, L, D] + q = q.reshape(bs, q_len, self.num_heads, self.head_dim) + k = k.reshape(bs, q_len, self.num_kv_heads, self.head_dim) + + q = self.rope(q.to(torch.float32), cos, sin) + k = self.rope(k.to(torch.float32), cos, sin) + + # 5. Restore original shape + convert to bfloat16 + q = q.reshape(hidden_states.shape[0], self.num_heads * self.head_dim).to(torch.bfloat16) + k = k.reshape(hidden_states.shape[0], self.num_kv_heads * self.head_dim).to(torch.bfloat16) + hidden_states = hidden_states.reshape(hidden_states_shape) + return q, k + + +class ImageKVCacheManager: + """ + Manages specialized caching and updating of KV-Cache for image tokens in multimodal models. + """ + + def __init__(self, image_token_len: int = 4097): + """ + Args: + image_token_len: Number of tokens per image (including special placeholders), + default 4097 (timestamp + 4096 image tokens). + """ + self.image_token_len: int = image_token_len + self.image_kv_cache: tuple[torch.Tensor, torch.Tensor] = None + + def _save_image_kv_caches( + self, + key: torch.Tensor, + value: torch.Tensor, + seq_len: int, + ) -> None: + bs, q_len, num_kv_heads, head_dim = key.shape + assert q_len == seq_len, f"for first-step, {q_len} != {seq_len}" + + key = key.reshape(-1, num_kv_heads, head_dim) + value = value.reshape(-1, num_kv_heads, head_dim) + + cached_prompt_len = seq_len - self.image_token_len - 1 + cached_key = [key[:cached_prompt_len], key[seq_len - 1 : seq_len]] + cached_value = [value[:cached_prompt_len], value[seq_len - 1 : seq_len]] + + if bs > 1: + assert bs == 2, "for cfg case, bs must be 2" + cached_key.append(key[seq_len : seq_len + cached_prompt_len]) + cached_key.append(key[-1:]) + + cached_value.append(value[seq_len : seq_len + cached_prompt_len]) + cached_value.append(value[-1:]) + + cached_key = torch.cat(cached_key, dim=0) + cached_value = torch.cat(cached_value, dim=0) + self.image_kv_cache_map = (cached_key, cached_value) + + def _update_image_kv_caches( + self, + key: torch.Tensor, + value: torch.Tensor, + seq_len: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + cached_key, cached_value = self.image_kv_cache_map + bs, q_len, num_kv_heads, head_dim = key.shape + + cached_prompt_len = cached_key.shape[0] // bs - 1 + assert (cached_prompt_len + 1) == (seq_len - q_len), f"{cached_prompt_len + 1} != {seq_len - q_len}" + + key = key.reshape(-1, num_kv_heads, head_dim) + value = value.reshape(-1, num_kv_heads, head_dim) + + new_key = [ + cached_key[:cached_prompt_len], + key[:q_len], + cached_key[cached_prompt_len : cached_prompt_len + 1], + ] + new_value = [ + cached_value[:cached_prompt_len], + value[:q_len], + cached_value[cached_prompt_len : cached_prompt_len + 1], + ] + + if bs > 1: + assert bs == 2, "for cfg case, bs must be 2" + new_key.append(cached_key[cached_prompt_len + 1 : cached_prompt_len + 1 + cached_prompt_len]) + new_key.append(key[q_len:]) + new_key.append(cached_key[-1:]) + + new_value.append(cached_value[cached_prompt_len + 1 : cached_prompt_len + 1 + cached_prompt_len]) + new_value.append(value[q_len:]) + new_value.append(cached_value[-1:]) + + new_key = torch.cat(new_key, dim=0) + new_value = torch.cat(new_value, dim=0) + new_key = new_key.reshape(bs, seq_len, num_kv_heads, head_dim) + new_value = new_value.reshape(bs, seq_len, num_kv_heads, head_dim) + + return new_key.contiguous(), new_value.contiguous() + + def __call__( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + self.image_token_len = kwargs.get("num_image_tokens") + first_step = kwargs.get("first_step") + + query_lens = kwargs.get("query_lens") + seq_lens = kwargs.get("seq_lens") + bs = len(query_lens) + q_len = query_lens[0] + seq_len = seq_lens[0] + assert query.shape[0] == bs * q_len, f"{query.shape[0]} != {bs * q_len}" + + head_num_per_rank = query.shape[1] + kv_head_num_per_rank = key.shape[1] + repeat_num = head_num_per_rank // kv_head_num_per_rank + head_dim = query.shape[2] + + query = query.reshape(bs, q_len, head_num_per_rank, head_dim) + key = key.reshape(bs, q_len, kv_head_num_per_rank, head_dim) + value = value.reshape(bs, q_len, kv_head_num_per_rank, head_dim) + + if first_step: + self.image_kv_cache_map = None + self._save_image_kv_caches(key, value, seq_len) + else: + key, value = self._update_image_kv_caches(key, value, seq_len) + + query = query.transpose(1, 2).contiguous() + key = key.transpose(1, 2).contiguous() + value = value.transpose(1, 2).contiguous() + + key = repeat_kv(key, repeat_num) + value = repeat_kv(value, repeat_num) + + attention_mask = attention_mask.contiguous() + + attn_output = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, dropout_p=0.0) + + attn_output = attn_output.transpose(1, 2).contiguous() # [bs, q_len, heads, head_dim] + attn_output = attn_output.reshape(bs * q_len, head_num_per_rank, head_dim) + return attn_output + + +@dataclass +class CausalMMOutputWithPast(CausalLMOutputWithPast): + diffusion_prediction: torch.Tensor | None = None + + +class HunyuanImage3Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`HunyuanImage3Model`]. It is used to instantiate + an Hunyuan model according to the specified arguments, defining the model architecture. Instantiating a + configuration with the defaults will yield a similar configuration to that of the Hunyuan-7B. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 32000): + Vocabulary size of the Hunyuan Image 3 model. Defines the number of different tokens that can be + represented by the `inputs_ids` passed when calling [`HunyuanImage3Model`] + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 11008): + Dimension of the MLP representations or shared MLP representations. + moe_intermediate_size (`int` or `List`, *optional*, defaults to 11008): + Dimension of the MLP representations in MoE. Use a list if you want a different size per layer. + num_hidden_layers (`int`, *optional*, defaults to 32): + Number of hidden layers in the Transformer decoder. + num_attention_heads (`int`, *optional*, defaults to 32): + Number of attention heads for each attention layer in the Transformer decoder. + num_key_value_heads (`int`, *optional*): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to + `num_attention_heads`. + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 2048): + The maximum sequence length that this model might ever be used with. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + rms_norm_eps (`float`, *optional*, defaults to 1e-06): + The epsilon used by the rms normalization layers. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if `config.is_decoder=True`. + pad_token_id (`int`, *optional*): + Padding token id. + bos_token_id (`int`, *optional*, defaults to 1): + Beginning of stream token id. + eos_token_id (`int`, *optional*, defaults to 2): + End of stream token id. + pretraining_tp (`int`, *optional*, defaults to 1): + Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this + document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is + necessary to ensure exact reproducibility of the pretraining results. Please refer to [this + issue](https://github.com/pytorch/pytorch/issues/76232). + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether to tie weight embeddings + rope_theta (`float`, *optional*, defaults to 10000.0): + The base period of the RoPE embeddings. + rope_scaling (`Dict`, *optional*): + Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling + strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is + `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update + `max_position_embeddings` to the expected new maximum. See the following thread for more information on how + these scaling strategies behave: + https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an + experimental feature, subject to breaking API changes in future versions. + attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): + Whether to use a bias in the query, key, value and output projection layers during self-attention. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + use_qk_norm (`bool`, *optional*, defaults to `False`): + Whether query and key in attention use norm + use_cla (`bool`, *optional*, defaults to `False`): + Whether to use CLA in attention + cla_share_factor (`int`, *optional*, defaults to 1): + The share factor of CLA + num_experts (`int` or `List`, *optional*, defaults to 1): + The number of experts for moe. If it is a list, it will be used as the number of experts for each layer. + num_shared_expert (`int` or `List`, *optional*, defaults to 1): + The number of shared experts for moe. If it is a list, it will be used as the number of shared experts + for each layer. + moe_topk (`int` or `List`, *optional*, defaults to 1): + The topk value for moe. If it is a list, it will be used as the topk value for each layer. + capacity_factor (Not used) (`float` or `List`, *optional*, defaults to 1.0): + The capacity factor for moe. If it is a list, it will be used as the capacity factor for each layer. + moe_layer_num_skipped (`int`, *optional*, defaults to 0): + First moe_layer_num_skipped layers do not use MoE. + """ + + model_type = "Hunyuan" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=290943, + hidden_size=4096, + intermediate_size: int = 11008, + moe_intermediate_size: int | list = None, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + attention_head_dim=None, + hidden_act="silu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-5, + use_cache=True, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + eod_token_id=3, + im_start_id=4, + im_end_id=5, + text_start_id=6, + text_end_id=7, + image_token_id=8, + video_start_id=9, + video_end_id=10, + im_newline_id=11, + mask_init_id=12, + pretraining_tp=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + mlp_bias=False, + attention_dropout=0.0, + use_qk_norm=False, + use_rotary_pos_emb=True, + use_cla=False, + cla_share_factor=1, + norm_type="hf_rms", + num_experts: int | list = 1, + use_mixed_mlp_moe=False, + num_shared_expert: int | list = 1, + moe_topk: int | list = 1, + capacity_factor: float | list = 1.0, + moe_drop_tokens=False, + moe_random_routing_dropped_token=False, + use_mla=False, + kv_lora_rank=512, + q_lora_rank=1536, + qk_rope_head_dim=64, + v_head_dim=128, + qk_nope_head_dim=128, + moe_layer_num_skipped=0, + norm_topk_prob=True, + routed_scaling_factor=1.0, + group_limited_greedy=False, + n_group=None, + topk_group=None, + add_classification_head=False, + class_num=0, + pool_type="last", + pad_id=-1, + # Added + moe_impl="eager", + vae_downsample_factor=(16, 16), # (h, w) + img_proj_type="unet", + patch_size=1, + patch_embed_hidden_dim=1024, + image_base_size=1024, + vae=None, + vit=None, + vit_processor=None, + vit_aligner=None, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.moe_impl = moe_impl + self.num_experts = num_experts + self.use_mixed_mlp_moe = use_mixed_mlp_moe + self.num_shared_expert = num_shared_expert + self.moe_topk = moe_topk + self.capacity_factor = capacity_factor + self.moe_drop_tokens = moe_drop_tokens + self.moe_random_routing_dropped_token = moe_random_routing_dropped_token + + if attention_head_dim is not None: + self.attention_head_dim = attention_head_dim + else: + self.attention_head_dim = self.hidden_size // num_attention_heads + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.pretraining_tp = pretraining_tp + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.mlp_bias = mlp_bias + self.attention_dropout = attention_dropout + self.use_qk_norm = use_qk_norm + self.use_rotary_pos_emb = use_rotary_pos_emb + self.use_cla = use_cla + self.cla_share_factor = cla_share_factor + self.norm_type = norm_type + # MLA args + self.use_mla = use_mla + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.v_head_dim = v_head_dim + + # DeepSeek related args + self.moe_layer_num_skipped = moe_layer_num_skipped + self.norm_topk_prob = norm_topk_prob + self.routed_scaling_factor = routed_scaling_factor + self.group_limited_greedy = group_limited_greedy + self.n_group = n_group + self.topk_group = topk_group + self.add_classification_head = add_classification_head + self.class_num = class_num + self.pool_type = pool_type + self.pad_id = pad_id + + if self.class_num is not None: + self.dense_list = [self.hidden_size, self.class_num] + + # ViT args + self.vit = vit + self.vit_processor = vit_processor + self.vit_aligner = vit_aligner + + # Image Gen args + self.vae = vae + self.vae_downsample_factor = vae_downsample_factor + self.img_proj_type = img_proj_type + self.patch_size = patch_size + self.patch_embed_hidden_dim = patch_embed_hidden_dim + self.image_base_size = image_base_size + + # token id + self.eod_token_id = eod_token_id + self.im_start_id = im_start_id + self.im_end_id = im_end_id + self.text_start_id = text_start_id + self.text_end_id = text_end_id + self.image_token_id = image_token_id + self.video_start_id = video_start_id + self.video_end_id = video_end_id + self.im_newline_id = im_newline_id + self.mask_init_id = mask_init_id + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +class HunyuanImage3ImageProcessor: + def __init__(self, config): + self.config = config + + self.reso_group = ResolutionGroup(base_size=config.image_base_size) + self.vae_processor = transforms.Compose( + [ + transforms.ToTensor(), + transforms.Normalize([0.5], [0.5]), # transform to [-1, 1] + ] + ) + self.vision_encoder_processor = Siglip2ImageProcessorFast.from_dict(config.vit_processor) + + def build_image_info(self, image_size): + # parse image size (HxW, H:W, or ) + if isinstance(image_size, str): + if image_size.startswith("")) + reso = self.reso_group[ratio_index] + image_size = reso.height, reso.width + elif "x" in image_size: + image_size = [int(s) for s in image_size.split("x")] + elif ":" in image_size: + image_size = [int(s) for s in image_size.split(":")] + else: + raise ValueError( + f"`image_size` should be in the format of 'HxW', 'H:W' or , got {image_size}." + ) + assert len(image_size) == 2, f"`image_size` should be in the format of 'HxW', got {image_size}." + elif isinstance(image_size, (list, tuple)): + assert len(image_size) == 2 and all(isinstance(s, int) for s in image_size), ( + f"`image_size` should be a tuple of two integers or a string in the format of 'HxW', got {image_size}." + ) + else: + raise ValueError( + f"`image_size` should be a tuple of two integers or a string in the format of 'WxH', got {image_size}." + ) + image_width, image_height = self.reso_group.get_target_size(image_size[1], image_size[0]) + token_height = image_height // (self.config.vae_downsample_factor[0] * self.config.patch_size) + token_width = image_width // (self.config.vae_downsample_factor[1] * self.config.patch_size) + base_size, ratio_idx = self.reso_group.get_base_size_and_ratio_index(image_size[1], image_size[0]) + image_info = ImageInfo( + image_type="gen_image", + image_width=image_width, + image_height=image_height, + token_width=token_width, + token_height=token_height, + base_size=base_size, + ratio_index=ratio_idx, + ) + return image_info + + +class HunYuanMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + prefix: str = "", + reduce_results: bool = True, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=[intermediate_size] * 2, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=intermediate_size, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + reduce_results=reduce_results, + ) + if hidden_act != "silu": + raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported for now.") + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class HunYuanSparseMoeBlock(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: QuantizationConfig | None = None, + layer_id: int = -1, + prefix: str = "", + enable_eplb: bool = False, + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + self.n_routed_experts = config.num_experts + + if self.tp_size > config.num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than the number of experts {config.num_experts}." + ) + + # Get layer_id topk if config.moe_topk is a list + if isinstance(config.moe_topk, list): + assert layer_id >= 0 + assert len(config.moe_topk) > layer_id + top_k = config.moe_topk[layer_id] + else: + top_k = config.moe_topk + + # If it is moe, moe_intermediate_size is preferred + intermediate_size = config.intermediate_size + if config.moe_intermediate_size is not None: + intermediate_size = ( + config.moe_intermediate_size + if isinstance(config.moe_intermediate_size, int) + else config.moe_intermediate_size[layer_id] + ) + + self.enable_eplb = False + self.n_logical_experts = self.n_routed_experts + self.n_redundant_experts = 0 + + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + quant_config=None, + prefix=f"{prefix}.gate", + ) + if config.use_mixed_mlp_moe > 0: + # Get layer_id num_shared_expert if config.num_shared_expert is + # a list. + if isinstance(config.num_shared_expert, list): + assert layer_id >= 0 + assert len(config.num_shared_expert) > layer_id + num_shared_expert = config.num_shared_expert[layer_id] + else: + num_shared_expert = config.num_shared_expert + + self.shared_mlp = HunYuanMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size * num_shared_expert, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_mlp", + ) + else: + self.shared_mlp = None + + self.experts = HunyuanFusedMoE( + shared_experts=self.shared_mlp, + num_experts=self.n_routed_experts, + top_k=top_k, + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + reduce_results=False, + renormalize=top_k > 1, + quant_config=quant_config, + prefix=f"{prefix}.experts", + enable_eplb=self.enable_eplb, + num_redundant_experts=self.n_redundant_experts, + pcp_size=1, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # NOTE: hidden_states can have either 1D or 2D shape. + orig_shape = hidden_states.shape + hidden_dim = hidden_states.shape[-1] + hidden_states = hidden_states.view(-1, hidden_dim) + + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts(hidden_states=hidden_states, router_logits=router_logits) + if self.shared_mlp is not None: + final_hidden_states = final_hidden_states[0] + final_hidden_states[1] + + if self.tp_size > 1: + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states.view(orig_shape) + + +class HunYuanAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + rope_theta: float = 10000, + rope_scaling: dict[str, Any] | None = None, + max_position_embeddings: int = 8192, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + cache_config: CacheConfig | None = None, + prefix: str = "", + layer_id: int = -1, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.num_key_value_heads = ( + config.num_key_value_heads // tp_size if config.num_key_value_heads else self.num_heads + ) + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + + if hasattr(config, "head_dim") and config.head_dim: + self.head_dim = config.head_dim + elif hasattr(config, "attention_head_dim"): + self.head_dim = config.attention_head_dim + else: + self.head_dim = self.hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.use_qk_norm = getattr(config, "use_qk_norm", False) + self.layer_id = layer_id + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + if rope_scaling is not None: + # for t2t, fallback rotary_emb type from 'custom' to default. + rope_scaling["rope_type"] = "default" + self.rotary_emb = get_rope( + head_size=self.head_dim, + max_position=max_position_embeddings, + rope_parameters=rope_scaling, + is_neox_style=True, + ) + # which attention, from vllm or omni? + # try use Atten from diffusion + self.attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + causal=False, + softmax_scale=self.scaling, + num_kv_heads=self.num_kv_heads, + # cache_config=cache_config, + # quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + # default image_token_len = timestamp + 4096*image_tokes + self.image_attn = ImageKVCacheManager(image_token_len=4097) + self.image_rope2d_emb = HunYuanRotary2DEmbedder( + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + ) + + if self.use_qk_norm: + self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_states: tuple[torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + custom_pos_emb: tuple[torch.FloatTensor] | None = None, + **kwargs, + ) -> torch.Tensor: + bsz, q_len, _ = hidden_states.size() + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + past_key_value: Cache | None = kwargs.get("past_key_value", None) + if past_key_value is not None: + position_ids = kwargs.get("position_ids") + key_states = k.reshape(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = v.reshape(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + cache_kwargs = {"cache_position": position_ids} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_id, cache_kwargs) + # for image_generation + if kwargs.get("mode", "gen_text") == "gen_image": + # assert positions is None, "positions should be None for image attention" + q, k = self.image_rope2d_emb(q, k, hidden_states, custom_pos_emb, **kwargs) + else: + q, k = self.rotary_emb(positions, q, k) + if self.use_qk_norm: + q = self.query_layernorm(q.view(-1, self.num_heads, self.head_dim).contiguous()) + k = self.key_layernorm(k.view(-1, self.num_kv_heads, self.head_dim).contiguous()) + # for image_generation + if kwargs.get("mode", "gen_text") == "gen_image": + attn_output = self.image_attn(q, k, v, attention_mask=attention_mask, **kwargs) + else: + attn_output = self.attn(q, k, v) + # For o_proj + attn_output = attn_output.view(q.shape[0], -1) + output, _ = self.o_proj(attn_output) + output = output.reshape(bsz, q_len, -1) + return output, None, past_key_value + + +class HunyuanFusedMoE(SharedFusedMoE): + def __init__(self, *, prefix: str = "", **kwargs): + super().__init__(prefix=prefix, **kwargs) + self._prefix = prefix + + self._init_hook_handle = self.register_forward_pre_hook(self._initialize_kernel_hook, with_kwargs=True) + + def _initialize_kernel_hook(self, module, args, kwargs): + if self.quant_method: + self.quant_method.process_weights_after_loading(self) + self._init_hook_handle.remove() + + def forward(self, hidden_states, router_logits): + from vllm.model_executor.layers.fused_moe.layer import get_forward_context + + ctx = get_forward_context() + if not ctx.remaining_moe_layers: + import re + + moe_names = [name for name in ctx.no_compile_layers.keys() if ".mlp.experts" in name] + + def get_layer_num(name): + match = re.search(r"layers\.(\d+)\.mlp", name) + return int(match.group(1)) if match else -1 + + moe_names.sort(key=get_layer_num, reverse=True) + ctx.remaining_moe_layers.extend(moe_names) + return super().forward(hidden_states, router_logits) + + +class HunyuanImage3DecoderLayer(nn.Module): + def __init__(self, config: HunyuanImage3Config, layer_idx: int, prefix: str = ""): + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + self.intermediate_size = ( + config.intermediate_size + if isinstance(config.intermediate_size, int) + else config.intermediate_size[layer_idx] + ) + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is not None and getattr(config, "original_max_position_embeddings", None): + rope_scaling["original_max_position_embeddings"] = config.original_max_position_embeddings + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + attention_bias = getattr(config, "attention_bias", False) or getattr(config, "bias", False) + + cla_factor = _get_cla_factor(config) + attention_type = ( + AttentionType.ENCODER_DECODER if layer_idx >= 0 and layer_idx % cla_factor != 0 else AttentionType.DECODER + ) + if attention_type == AttentionType.DECODER: + self.self_attn = HunYuanAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads), + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=None, + bias=attention_bias, + cache_config=None, + prefix=f"{prefix}.self_attn", + layer_id=layer_idx, + ) + else: + raise RuntimeError(f"Unsupported attention type: {attention_type}") + + if ( + (isinstance(config.num_experts, int) and config.num_experts > 1) + or (isinstance(config.num_experts, list) and max(config.num_experts) > 1) + ) and layer_idx >= config.moe_layer_num_skipped: + self.mlp = HunYuanSparseMoeBlock(config, layer_id=layer_idx, prefix=f"{prefix}.mlp") + else: + self.mlp = HunYuanMLP(self.hidden_size, self.intermediate_size, config.hidden_act) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_value: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + custom_pos_emb: tuple[torch.FloatTensor] | None = None, + **kwargs, + ) -> tuple[torch.FloatTensor | Any]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): + attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1, + query_sequence_length, key_sequence_length)` if default attention is used. + position_ids (`torch.LongTensor`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + custom_pos_emb (`Tuple[torch.FloatTensor]`, *optional*): custom position embedding for rotary + position embedding + """ + if "padding_mask" in kwargs: + logger.warning( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use " + "`attention_mask` instead.`" + ) + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + custom_pos_emb=custom_pos_emb, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + hidden_states = residual + hidden_states + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +class HunyuanImage3PreTrainedModel(PreTrainedModel): + config_class = HunyuanImage3Config + base_model_prefix = "" + supports_gradient_checkpointing = True + _no_split_modules = ["HunyuanImage3DecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class HunyuanImage3Model(nn.Module): + def __init__(self, config: HunyuanImage3Config, prefix: str = ""): + super().__init__() + quant_config = None + lora_config = None + self.num_redundant_experts = 0 + self.config = config + self.quant_config = quant_config + self.padding_idx = config.pad_token_id + lora_vocab = (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1)) if lora_config else 0 + self.vocab_size = config.vocab_size + lora_vocab + self.org_vocab_size = config.vocab_size + self.wte = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank): + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + quant_config=quant_config, + ) + else: + self.embed_tokens = PPMissingLayer() + # maybe should not use make_layers() + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: HunyuanImage3DecoderLayer( + config=config, + layer_idx=int(prefix.split(".")[-1]), + prefix=prefix, + ), + prefix=f"{prefix}.layers", + ) + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + def _split_qkv_weight(self, qkv: torch.Tensor): + num_attention_heads = self.config.num_attention_heads + num_kv_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads) + num_key_value_groups = num_attention_heads // num_kv_heads + hidden_size = self.config.hidden_size + + if hasattr(self.config, "head_dim"): + attention_head_dim = self.config.head_dim + elif hasattr(self.config, "attention_head_dim"): + attention_head_dim = self.config.attention_head_dim + else: + attention_head_dim = self.config.hidden_size // num_attention_heads + + qkv = qkv.reshape(num_kv_heads, num_key_value_groups + 2, attention_head_dim, hidden_size) + q, k, v = torch.split(qkv, (num_key_value_groups, 1, 1), dim=1) + q = q.reshape(-1, hidden_size) + k = k.reshape(-1, hidden_size) + v = v.reshape(-1, hidden_size) + return torch.concat((q, k, v)) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + if _is_moe(self.config): + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + fused_moe_expert_mapping = HunyuanFusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + num_redundant_experts=self.num_redundant_experts, + ) + expert_weights_remapping = { + "gate_proj": ("gate_and_up_proj", 1, 2), + "up_proj": ("gate_and_up_proj", 0, 2), + } + return fused_moe_expert_mapping, expert_weights_remapping + else: + return [], {} + + # rename for delay load + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + cla_factor = _get_cla_factor(self.config) + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + + num_attention_heads = self.config.num_attention_heads + num_kv_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads) + split_params_mapping = [ + (".gate_up_proj", ".gate_and_up_proj", 2, [(1, 1), (0, 1)], None), + ( + ".qkv_proj", + ".qkv_proj", + num_attention_heads + num_kv_heads * 2, + [("q", num_attention_heads), ("k", num_kv_heads), ("v", num_kv_heads)], + self._split_qkv_weight, + ), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + expert_params_mapping, expert_weights_remapping = self.get_expert_mapping() + + # List of unexpected keywords in weight names + unexpected_keywords = [ + "vae", + "vision_aligner", + "vision_model", + "final_layer", + "patch_embed", + "timestep_emb", + "time_embed", + "time_embed_2", + "guidance_emb", + "timestep_r_emb", + ] + + def contains_unexpected_keyword(name, keywords): + """Check if the name contains any unexpected keywords""" + for keyword in keywords: + if keyword in name: + return True + return False + + for name, loaded_weight in weights: + # print(f"Loading weight name: {name}, tp_rank: {tp_rank}", flush=True) + if contains_unexpected_keyword(name, unexpected_keywords): + print(f"Skipping unexpected weight name: {name}") + continue + if "rotary_emb.inv_freq" in name: + continue + if "gate_proj_bias" in name: + name = name.replace("gate_proj_bias", "gate_proj.bias") + if "up_proj_bias" in name: + name = name.replace("up_proj_bias", "up_proj.bias") + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + # With tie_word_embeddings, we can skip lm_head.weight + # The weight might appear unnecessarily in the files if the model is + # processed with quantization, LoRA, fine-tuning, etc. + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): + # Loading kv cache scales for compressed-tensors quantization + param = params_dict[scale_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + loaded_weight = loaded_weight[0] + weight_loader(param, loaded_weight) + continue + + is_found = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts" in name: + continue + # cross layer only have q_proj, skip qkv pack + if weight_name == ".q_proj": + match = re.search(r"layers\.\d+", name) + if match: + layer_id = int(match.group(0).split(".")[-1]) + if cla_factor > 1 and layer_id % cla_factor != 0: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name) + is_found = True + break + if is_found: + continue + + for ( + param_name, + weight_name, + den, + split_param, + func, + ) in split_params_mapping: + if weight_name not in name: + continue + if "mlp.experts" in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if is_pp_missing_parameter(name, self): + continue + + assert loaded_weight.shape[0] % den == 0 + units = loaded_weight.shape[0] // den + param = params_dict[name] + weight_loader = param.weight_loader + offset = 0 + for shard_id, num in split_param: + new_offset = offset + num * units + if func: + weight_loader(param, func(loaded_weight)[offset:new_offset], shard_id) + else: + weight_loader(param, loaded_weight[offset:new_offset], shard_id) + offset = new_offset + + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + is_expert_weight = False + is_found = False + found_num = 0 + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + offset = 0 + den = 1 + for ( + mapped_weight_substr, + origin_weight_info, + ) in expert_weights_remapping.items(): + if mapped_weight_substr in weight_name: + origin_weight_name, offset, den = origin_weight_info + weight_name = weight_name.replace(mapped_weight_substr, origin_weight_name) + break + if weight_name not in name: + continue + is_expert_weight = True + name_mapped = name.replace(weight_name, param_name) + found_num += 1 + if is_pp_missing_parameter(name_mapped, self): + continue + param = params_dict[name_mapped] + weight_loader = cast(Callable[..., bool], param.weight_loader) + assert loaded_weight.shape[0] % den == 0 + units = loaded_weight.shape[0] // den + + success = weight_loader( + param, + loaded_weight[offset * units : offset * units + units], + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + loaded_params.add(name_mapped) + is_found = True + if found_num == den: + break + if is_found: + continue + if is_expert_weight: + # We've checked that this is an expert weight + # However it's not mapped locally to this rank + # So we simply skip it + continue + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + if is_pp_missing_parameter(name, self): + continue + + if "mlp.gate.wg." in name: + name = name.replace("wg.", "") + if name == "ln_f.weight": + name = "norm.weight" + if name == "wte.weight": + name = "embed_tokens.weight" + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + custom_pos_emb: tuple[torch.FloatTensor] | None = None, + mode: str = "gen_text", + first_step: bool | None = None, + query_lens: list[int] | None = None, + seq_lens: list[int] | None = None, + num_image_tokens: int | None = None, + gen_timestep_scatter_index: torch.Tensor | None = None, + ) -> tuple | BaseModelOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if inputs_embeds is None: + inputs_embeds = self.wte(input_ids) + + # embed positions + hidden_states = inputs_embeds + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for layer_idx, decoder_layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = decoder_layer( + positions=None, + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + custom_pos_emb=custom_pos_emb, + mode=mode, + first_step=first_step, + query_lens=query_lens, + seq_lens=seq_lens, + num_image_tokens=num_image_tokens, + gen_timestep_scatter_index=gen_timestep_scatter_index, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = next_decoder_cache + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +class ClassifierFreeGuidance: + def __init__( + self, + use_original_formulation: bool = False, + start: float = 0.0, + stop: float = 1.0, + ): + super().__init__() + self.use_original_formulation = use_original_formulation + + def __call__( + self, + pred_cond: torch.Tensor, + pred_uncond: torch.Tensor | None, + guidance_scale: float, + step: int, + ) -> torch.Tensor: + shift = pred_cond - pred_uncond + pred = pred_cond if self.use_original_formulation else pred_uncond + pred = pred + guidance_scale * shift + + return pred + + +@dataclass +class HunyuanImage3Text2ImagePipelineOutput(BaseOutput): + samples: list[Any] | np.ndarray + + +class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): + r""" + Pipeline for condition-to-sample generation using Stable Diffusion. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + model ([`ModelMixin`]): + A model to denoise the diffused latents. + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `diffusion_model` to denoise the diffused latents. Can be one of + [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. + """ + + model_cpu_offload_seq = "" + _optional_components = [] + _exclude_from_cpu_offload = [] + _callback_tensor_inputs = ["latents"] + + def __init__( + self, + model, + scheduler: SchedulerMixin, + vae, + progress_bar_config: dict[str, Any] | None = None, + ): + super().__init__() + + if progress_bar_config is None: + progress_bar_config = {} + if not hasattr(self, "_progress_bar_config"): + self._progress_bar_config = {} + self._progress_bar_config.update(progress_bar_config) + self.register_modules( + model=model, + scheduler=scheduler, + vae=vae, + ) + + # should be a tuple or a list corresponding to the size of latents (batch_size, channel, *size) + # if None, will be treated as a tuple of 1 + self.latent_scale_factor = self.model.config.vae_downsample_factor + self.image_processor = VaeImageProcessor(vae_scale_factor=self.latent_scale_factor) + + # Must start with APG_mode_ + self.cfg_operator = ClassifierFreeGuidance() + + @staticmethod + def denormalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: + """ + Denormalize an image array to [0,1]. + """ + return (images / 2 + 0.5).clamp(0, 1) + + @staticmethod + def pt_to_numpy(images: torch.Tensor) -> np.ndarray: + """ + Convert a PyTorch tensor to a NumPy image. + """ + images = images.cpu().permute(0, 2, 3, 1).float().numpy() + return images + + @staticmethod + def numpy_to_pil(images: np.ndarray): + """ + Convert a numpy image or a batch of images to a PIL image. + """ + if images.ndim == 3: + images = images[None, ...] + images = (images * 255).round().astype("uint8") + if images.shape[-1] == 1: + # special case for grayscale (single channel) images + pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images] + else: + pil_images = [Image.fromarray(image) for image in images] + + return pil_images + + def prepare_extra_func_kwargs(self, func, kwargs): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + extra_kwargs = {} + + for k, v in kwargs.items(): + accepts = k in set(inspect.signature(func).parameters.keys()) + if accepts: + extra_kwargs[k] = v + return extra_kwargs + + def prepare_latents(self, batch_size, latent_channel, image_size, dtype, device, generator, latents=None): + if self.latent_scale_factor is None: + latent_scale_factor = (1,) * len(image_size) + elif isinstance(self.latent_scale_factor, int): + latent_scale_factor = (self.latent_scale_factor,) * len(image_size) + elif isinstance(self.latent_scale_factor, tuple) or isinstance(self.latent_scale_factor, list): + assert len(self.latent_scale_factor) == len(image_size), ( + "len(latent_scale_factor) should be the same as len(image_size)" + ) + latent_scale_factor = self.latent_scale_factor + else: + raise ValueError( + f"latent_scale_factor should be either None, int, tuple of int, or list of int, " + f"but got {self.latent_scale_factor}" + ) + + latents_shape = ( + batch_size, + latent_channel, + *[int(s) // f for s, f in zip(image_size, latent_scale_factor)], + ) + if isinstance(generator, list) and 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." + ) + + if latents is None: + latents = randn_tensor(latents_shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + + # Check existence to make it compatible with FlowMatchEulerDiscreteScheduler + if hasattr(self.scheduler, "init_noise_sigma"): + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * self.scheduler.init_noise_sigma + + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def guidance_rescale(self): + return self._guidance_rescale + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + def set_scheduler(self, new_scheduler): + self.register_modules(scheduler=new_scheduler) + + @torch.no_grad() + def __call__( + self, + batch_size: int, + image_size: list[int], + num_inference_steps: int = 50, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + guidance_scale: float = 5.0, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + return_dict: bool = True, + guidance_rescale: float = 0.0, + callback_on_step_end: Callable[[int, int, dict], None] + | PipelineCallback + | MultiPipelineCallbacks + | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + model_kwargs: dict[str, Any] | None = None, + **kwargs, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `List[str]`): + The text to guide image generation. + image_size (`Tuple[int]` or `List[int]`): + The size (height, width) of the generated image. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 7.5): + A higher guidance scale value encourages the model to generate samples closely linked to the + `condition` at the expense of lower sample quality. Guidance scale is enabled when `guidance_scale > 1`. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for sample + generation. Can be used to tweak the same generation with different conditions. If not provided, + a latents tensor is generated by sampling using the supplied random `generator`. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated sample. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~DiffusionPipelineOutput`] instead of a + plain tuple. + guidance_rescale (`float`, *optional*, defaults to 0.0): + Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when + using zero terminal SNR. + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of + each denoising step during the inference. with the following arguments: `callback_on_step_end(self: + DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a + list of all tensors as specified by `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + + Examples: + + Returns: + [`~DiffusionPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~DiffusionPipelineOutput`] is returned, + otherwise a `tuple` is returned where the first element is a list with the generated samples. + """ + callback_steps = kwargs.pop("callback_steps", None) + pbar_steps = kwargs.pop("pbar_steps", None) + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + self._guidance_scale = guidance_scale + self._guidance_rescale = guidance_rescale + + cfg_factor = 1 + self.do_classifier_free_guidance + + # Define call parameters + device = self._execution_device + + # Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + timesteps, + sigmas, + ) + + # Prepare latent variables + latents = self.prepare_latents( + batch_size=batch_size, + latent_channel=self.model.config.vae["latent_channels"], + image_size=image_size, + dtype=torch.bfloat16, + device=device, + generator=generator, + latents=latents, + ) + + # Prepare extra step kwargs. + _scheduler_step_extra_kwargs = self.prepare_extra_func_kwargs(self.scheduler.step, {"generator": generator}) + + # Prepare model kwargs + input_ids = model_kwargs.pop("input_ids") + attention_mask = self.model._prepare_attention_mask_for_generation( # noqa + input_ids, + self.model.generation_config, + model_kwargs=model_kwargs, + ) + b, _, q_len1, seq_len = attention_mask.shape + query_lens = [q_len1] * b + seq_lens = [seq_len] * b + model_kwargs["query_lens"] = query_lens + model_kwargs["seq_lens"] = seq_lens + model_kwargs["attention_mask"] = attention_mask.to(latents.device) + + # Sampling loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * cfg_factor) + # latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + t_expand = t.repeat(latent_model_input.shape[0]) + + model_inputs = self.model.prepare_inputs_for_generation( + input_ids, + images=latent_model_input, + timestep=t_expand, + **model_kwargs, + ) + + with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): + model_output = self.model.forward_call(**model_inputs, first_step=(i == 0)) + pred = model_output["diffusion_prediction"] + pred = pred.to(dtype=torch.float32) + + # perform guidance + if self.do_classifier_free_guidance: + pred_cond, pred_uncond = pred.chunk(2) + pred = self.cfg_operator(pred_cond, pred_uncond, self.guidance_scale, step=i) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(pred, t, latents, **_scheduler_step_extra_kwargs, return_dict=False)[0] + + if i != len(timesteps) - 1: + model_kwargs = self.model._update_model_kwargs_for_generation( # noqa + model_output, + model_kwargs, + ) + if input_ids.shape[1] != model_kwargs["position_ids"].shape[1]: + input_ids = torch.gather(input_ids, 1, index=model_kwargs["position_ids"]) + attention_mask = model_kwargs.get("attention_mask") + b, _, q_len1, seq_len = attention_mask.shape + query_lens = [q_len1] * b + seq_lens = [seq_len] * b + model_kwargs["query_lens"] = query_lens + model_kwargs["seq_lens"] = seq_lens + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if hasattr(self.vae.config, "scaling_factor") and self.vae.config.scaling_factor: + latents = latents / self.vae.config.scaling_factor + if hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor: + latents = latents + self.vae.config.shift_factor + + if hasattr(self.vae, "ffactor_temporal"): + latents = latents.unsqueeze(2) + + with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=True): + image = self.vae.decode(latents, return_dict=False, generator=generator)[0] + + if hasattr(self.vae, "ffactor_temporal"): + assert image.shape[2] == 1, "image should have shape [B, C, T, H, W] and T should be 1" + image = image.squeeze(2) + + do_denormalize = [True] * image.shape[0] + image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) + + if not return_dict: + return (image,) + + return HunyuanImage3Text2ImagePipelineOutput(samples=image) + + +def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. + dim (int): the dimension of the output. + max_period (int): controls the minimum frequency of the embeddings. + + Returns: + embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. + + .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + """ + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__( + self, + hidden_size, + act_layer=nn.GELU, + frequency_embedding_size=256, + max_period=10000, + out_size=None, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + if out_size is None: + out_size = hidden_size + + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True, **factory_kwargs), + act_layer(), + nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs), + ) + nn.init.normal_(self.mlp[0].weight, std=0.02) + nn.init.normal_(self.mlp[2].weight, std=0.02) + + def forward(self, t): + t_freq = timestep_embedding(t, self.frequency_embedding_size, self.max_period).type(self.mlp[0].weight.dtype) + t_emb = self.mlp(t_freq) + return t_emb + + +class ResBlock(nn.Module): + r""" + A residual block that can optionally change the number of channels. + + Args: + in_channels (`int`): + The number of input channels. + emb_channels (`int`): + The number of timestep embedding channels. + dropout (`float`): + The rate of dropout. + out_channels (`int`, *optional*): + If specified, the number of output channels. + use_conv (`bool`, *optional*): + If True and out_channels is specified, use a spatial convolution instead of a + smaller 1x1 convolution to change the channels in the skip connection. + dims (`int`, *optional*): + Determines if the signal is 1D, 2D, or 3D. + up (`bool`, *optional*): + If True, use this block for upsampling. + down (`bool`, *optional*): + If True, use this block for downsampling. + + """ + + def __init__( + self, + in_channels, + emb_channels, + out_channels=None, + dropout=0.0, + use_conv=False, + dims=2, + up=False, + down=False, + device=None, + dtype=None, + ) -> None: + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.in_channels = in_channels + self.dropout = dropout + self.out_channels = out_channels or self.in_channels + self.use_conv = use_conv + + self.in_layers = nn.Sequential( + normalization(self.in_channels, **factory_kwargs), + nn.SiLU(), + conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs), # noqa: N802 + ) + + self.updown = up or down + self.h_upd = self.x_upd = nn.Identity() + + self.emb_layers = nn.Sequential(nn.SiLU(), linear(emb_channels, 2 * self.out_channels, **factory_kwargs)) + + self.out_layers = nn.Sequential( + normalization(self.out_channels, **factory_kwargs), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1, **factory_kwargs)), # noqa: N802 + ) + + if self.out_channels == self.in_channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs) # noqa: N802 + else: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 1, **factory_kwargs) # noqa: N802 + + def forward(self, x, emb) -> torch.Tensor: + if self.updown: + in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] + h = in_rest(x) + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + else: + h = self.in_layers(x) + + emb_out = self.emb_layers(emb) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + + # Adaptive Group Normalization + out_norm, out_rest = self.out_layers[0], self.out_layers[1:] + scale, shift = torch.chunk(emb_out, 2, dim=1) + h = out_norm(h) * (1.0 + scale) + shift + h = out_rest(h) + + return self.skip_connection(x) + h + + +class UNetDown(nn.Module): + def __init__( + self, patch_size, in_channels, emb_channels, hidden_channels, out_channels, dropout=0.0, device=None, dtype=None + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + + self.patch_size = patch_size + assert self.patch_size in [1, 2, 4, 8] + + self.model = nn.ModuleList( + [ + conv_nd( + 2, in_channels=in_channels, out_channels=hidden_channels, kernel_size=3, padding=1, **factory_kwargs + ) # noqa: N802 + ] + ) + + if self.patch_size == 1: + self.model.append( + ResBlock( + in_channels=hidden_channels, + emb_channels=emb_channels, + out_channels=out_channels, + dropout=dropout, + **factory_kwargs, + ) + ) + else: + for i in range(self.patch_size // 2): + self.model.append( + ResBlock( + in_channels=hidden_channels, + emb_channels=emb_channels, + out_channels=hidden_channels if (i + 1) * 2 != self.patch_size else out_channels, + dropout=dropout, + down=True, + **factory_kwargs, + ) + ) + + def forward(self, x, t): + assert x.shape[2] % self.patch_size == 0 and x.shape[3] % self.patch_size == 0 + for module in self.model: + if isinstance(module, ResBlock): + x = module(x, t) + else: + x = module(x) + _, _, token_h, token_w = x.shape + x = rearrange(x, "b c h w -> b (h w) c") + return x, token_h, token_w + + +class UNetUp(nn.Module): + def __init__( + self, + patch_size, + in_channels, + emb_channels, + hidden_channels, + out_channels, + dropout=0.0, + device=None, + dtype=None, + out_norm=False, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + + self.patch_size = patch_size + assert self.patch_size in [1, 2, 4, 8] + + self.model = nn.ModuleList() + + if self.patch_size == 1: + self.model.append( + ResBlock( + in_channels=in_channels, + emb_channels=emb_channels, + out_channels=hidden_channels, + dropout=dropout, + **factory_kwargs, + ) + ) + else: + for i in range(self.patch_size // 2): + self.model.append( + ResBlock( + in_channels=in_channels if i == 0 else hidden_channels, + emb_channels=emb_channels, + out_channels=hidden_channels, + dropout=dropout, + up=True, + **factory_kwargs, + ) + ) + + if out_norm: + self.model.append( + nn.Sequential( + normalization(hidden_channels, **factory_kwargs), + nn.SiLU(), + conv_nd( + 2, + in_channels=hidden_channels, + out_channels=out_channels, + kernel_size=3, + padding=1, + **factory_kwargs, + ), # noqa: N802 + ) + ) + else: + self.model.append( + conv_nd( + 2, + in_channels=hidden_channels, + out_channels=out_channels, + kernel_size=3, + padding=1, + **factory_kwargs, + ) # noqa: N802 + ) + + # batch_size, seq_len, model_dim + def forward(self, x, t, token_h, token_w): + x = rearrange(x, "b (h w) c -> b c h w", h=token_h, w=token_w) + for module in self.model: + if isinstance(module, ResBlock): + x = module(x, t) + else: + x = module(x) + return x diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py new file mode 100644 index 000000000..9a8f26d38 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py @@ -0,0 +1,1017 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import logging +from collections.abc import Iterable +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from transformers.generation.configuration_utils import GenerationConfig +from transformers.generation.utils import ALL_CACHE_NAMES, GenerationMixin +from transformers.models.siglip2 import Siglip2VisionConfig, Siglip2VisionModel +from transformers.utils.generic import ModelOutput +from vllm.config.vllm import get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_rank +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.transformers_utils.config import get_config + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.request import OmniDiffusionRequest + +from .autoencoder import AutoencoderKLConv3D +from .hunyuan_image_3_tokenizer import TokenizerWrapper +from .hunyuan_image_3_transformer import ( + CausalMMOutputWithPast, + HunyuanImage3ImageProcessor, + HunyuanImage3Model, + HunyuanImage3PreTrainedModel, + HunyuanImage3Text2ImagePipeline, + ImageInfo, + JointImageInfo, + LightProjector, + TimestepEmbedder, + UNetDown, + UNetUp, + build_batch_2d_rope, + get_full_state_dict, + real_batched_index_select, +) + +logger = logging.getLogger(__name__) + +BatchRaggedImages = torch.Tensor | list[torch.Tensor | list[torch.Tensor]] +BatchRaggedTensor = torch.Tensor | list[torch.Tensor] + + +def default(val, d): + return val if val is not None else d + + +def to_device(data, device): + if device is None: + return data + if isinstance(data, torch.Tensor): + return data.to(device) + elif isinstance(data, list): + return [to_device(x, device) for x in data] + else: + return data + + +class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin): + def __init__(self, od_config: OmniDiffusionConfig) -> None: + self.hf_config = get_config(od_config.model, trust_remote_code=True) + super().__init__(self.hf_config) + # update diffusion config + self.generation_config = GenerationConfig.from_pretrained(od_config.model) + self.od_config = od_config + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder=None, + revision=od_config.revision, + prefix="", + fall_back_to_pt=True, + ) + ] + self.model = HunyuanImage3Model(self.hf_config) + self.vae = AutoencoderKLConv3D.from_config(self.hf_config.vae) + self._pipeline = None + self._tkwrapper = TokenizerWrapper(od_config.model) + self.image_processor = HunyuanImage3ImageProcessor(self.hf_config) + self.hf_config.vit.pop("use_return_dict", None) + vision_config = Siglip2VisionConfig(**self.hf_config.vit) + self.vision_model = Siglip2VisionModel(vision_config).vision_model + # self.vision_model = vision_model.vision_model + self.vision_aligner = LightProjector(self.hf_config.vit_aligner) + self.timestep_emb = TimestepEmbedder(hidden_size=self.hf_config.hidden_size) + if self.hf_config.img_proj_type != "unet": + raise ValueError(f"Unknown img_proj_type: {self.hf_config.img_proj_type}") + + self.patch_embed = UNetDown( + patch_size=self.hf_config.patch_size, + emb_channels=self.hf_config.hidden_size, + in_channels=self.hf_config.vae["latent_channels"], + hidden_channels=self.hf_config.patch_embed_hidden_dim, + out_channels=self.hf_config.hidden_size, + ) + self.time_embed = TimestepEmbedder(hidden_size=self.hf_config.hidden_size) + self.final_layer = UNetUp( + patch_size=self.hf_config.patch_size, + emb_channels=self.hf_config.hidden_size, + in_channels=self.hf_config.hidden_size, + hidden_channels=self.hf_config.patch_embed_hidden_dim, + out_channels=self.hf_config.vae["latent_channels"], + out_norm=True, + ) + self.time_embed_2 = TimestepEmbedder(hidden_size=self.hf_config.hidden_size) + self.lm_head = nn.Linear(self.hf_config.hidden_size, self.hf_config.vocab_size, bias=False) + self.vllm_config = get_current_vllm_config() + self.post_init() + + def pre_load(self): + tp_rank = get_tensor_model_parallel_rank() + state_dict = get_full_state_dict(self.od_config.model) + non_layer_prefixes = [ + "vae", + "vision_model", + "vision_aligner", + "lm_head", + "patch_embed", + "timestep_emb", + "model.wte", + "model.ln_f", + "time_embed", + "time_embed_2", + "final_layer.model", + ] + filtered_sd = {} + for k, v in state_dict.items(): + if any(k.startswith(prefix) for prefix in non_layer_prefixes): + filtered_sd[k] = v + + missing, unexpected = self.load_state_dict(filtered_sd, strict=False) + + for prefix in non_layer_prefixes: + if hasattr(self, prefix.split(".")[0]): + module = dict(self.named_modules()).get(prefix) + if module: + module.to(f"cuda:{tp_rank}") + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + self.pre_load() + skip_prefixes = ["lm_head."] if self.hf_config.tie_word_embeddings else [] + # List of unexpected keywords in weight names + unexpected_keywords = [ + "vae", + "vision_aligner", + "vision_model", + "final_layer", + "patch_embed", + "timestep_emb", + "time_embed", + "time_embed_2", + "guidance_emb", + "timestep_r_emb", + ] + skip_prefixes.extend(unexpected_keywords) + loader = AutoWeightsLoader( + self, + skip_prefixes=skip_prefixes, + ) + return loader.load_weights(weights) + + def prepare_seed(self, seed=None, batch_size=1): + # random seed + if seed is not None: + return [seed + i for i in range(batch_size)] + else: + import random + + return [random.randint(0, 2**32 - 1) for _ in range(batch_size)] + + @property + def pipeline(self): + if self._pipeline is None: + # shift hard code + self.scheduler = FlowMatchEulerDiscreteScheduler( + num_train_timesteps=1000, + shift=self.generation_config.flow_shift, + use_dynamic_shifting=False, + base_shift=0.5, + max_shift=1.15, + time_shift_type="exponential", + stochastic_sampling=False, + ) + self._pipeline = HunyuanImage3Text2ImagePipeline(model=self, scheduler=self.scheduler, vae=self.vae) + return self._pipeline + + @staticmethod + def get_pos_emb(custom_pos_emb, position_ids): + cos, sin = custom_pos_emb + cos = real_batched_index_select(cos, dim=1, idx=position_ids) + sin = real_batched_index_select(sin, dim=1, idx=position_ids) + return cos, sin + + def instantiate_vae_image_tokens( + self, + x: torch.Tensor, + images: BatchRaggedImages, + ts: BatchRaggedTensor, + image_mask: torch.Tensor, + ): + r""" + Instantiate the VAE image embeddings into the input embedding sequence. + Args: + x (`torch.Tensor`): + Input sequence tensor with shape `(batch_size, seq_len, n_embd)`. + images (`BatchRaggedImages`): + Batch of images to embed. Can be: + - A 4-D tensor (batch, channels, height, width) + - A list of 4-D tensors (variable number of images per batch) + - A list of lists of 3-D tensors (ragged batch structure) + ts (`BatchRaggedTensor`, *optional*): + Timestep tensor(s) for conditioning. Can be: + - A 1-D tensor (single timestep per batch) + - A list of 1-D tensors (variable timesteps per batch) + image_mask (`torch.Tensor`, *optional*): + Boolean mask tensor with shape `(batch_size, seq_len)` indicating which positions + should be replaced with image embeddings. + """ + batch_size, seq_len, n_embd = x.shape + + if isinstance(images, list): + index = torch.arange(seq_len, device=x.device).unsqueeze(0).repeat(batch_size, 1) + t_emb = [] + for i, (image_i, t_i) in enumerate(zip(images, ts)): + if isinstance(image_i, torch.Tensor): + # time_embed needs a 1-D tensor as input + t_i_emb = self.time_embed(t_i) + # n_{i} x one_image_seq_len x n_embd + image_i_seq, _, _ = self.patch_embed(image_i, t_i_emb) + # 1 x (n_{i} * one_image_seq_len) + image_i_scatter_index = index[i : i + 1].masked_select(image_mask[i : i + 1].bool()).reshape(1, -1) + x[i : i + 1].scatter_( + dim=1, + index=image_i_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd), + # 1 x (n_{i} * one_image_seq_len) x n_embd + src=image_i_seq.reshape(1, -1, n_embd), # 1 x (n_{i} * one_image_seq_len) x n_embd + ) + t_emb.append(t_i_emb) + elif isinstance(image_i, list): + # time_embed needs a 1-D tensor as input + t_i_emb = self.time_embed(t_i) # n_{i} x d + image_i_seq_list = [] + for j in range(len(image_i)): + image_ij = image_i[j] + if image_ij.dim() == 4: + assert image_i[j].shape[0] == 1, "image_i[j] should have a batch dimension of 1" + elif image_ij.dim() == 3: + image_ij = image_ij.unsqueeze(0) + else: + raise ValueError(f"image_i[j] should have 3 or 4 dimensions, got {image_ij.dim()}") + # 1 x one_image_seq_len_{j} x n_embd + image_i_seq_j, _, _ = self.patch_embed(image_ij, t_i_emb[j : j + 1]) + image_i_seq_list.append(image_i_seq_j) + # 1 x sum_{j}(one_image_seq_len_{j}) x n_embd + image_i_seq = torch.cat(image_i_seq_list, dim=1) + # 1 x sum_{j}(one_image_seq_len_{j}) + image_i_scatter_index = index[i : i + 1].masked_select(image_mask[i : i + 1].bool()).reshape(1, -1) + x[i : i + 1].scatter_( + dim=1, + index=image_i_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd), + # 1 x sum_{j}(one_image_seq_len_{j}) x n_embd + src=image_i_seq.reshape(1, -1, n_embd), # 1 x sum_{j}(one_image_seq_len_{j}) x n_embd + ) + t_emb.append(t_i_emb) + else: + raise TypeError(f"image_i should be a torch.Tensor or a list, got {type(image_i)}") + token_h, token_w = None, None + else: + # images is a 4-D tensor + batch_size, seq_len, n_embd = x.shape + index = torch.arange(seq_len, device=x.device).unsqueeze(0).repeat(batch_size, 1) + t_emb = self.time_embed(ts) + image_seq, token_h, token_w = self.patch_embed(images, t_emb) + image_scatter_index = index.masked_select(image_mask.bool()).reshape(batch_size, -1) + x.scatter_( + dim=1, + index=image_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd), + src=image_seq, + ) + + return x, token_h, token_w + + def instantiate_timestep_tokens( + self, + x: torch.Tensor, + t: BatchRaggedTensor, + timestep_scatter_index: BatchRaggedTensor, + ): + batch_size, seq_len, n_embd = x.shape + # batch_size x n x n_embd + timestep_scatter_src = self.timestep_emb(t.reshape(-1)).reshape(batch_size, -1, n_embd) + x.scatter_( + dim=1, + index=timestep_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd), + src=timestep_scatter_src, + ) + + return x + + def instantiate_vit_image_tokens( + self, + x: torch.Tensor, + cond_vit_images: torch.Tensor | list[torch.Tensor], + cond_vit_image_mask: torch.Tensor, + vit_kwargs: dict[str, Any], + ): + # 1. Forward the vit encoder and vit aligner to get the vit image embeddings and align them to the + # transformer hidden size + cond_vit_image_embeds = [] + for batch_idx, image in enumerate(cond_vit_images): + cur_kwargs = {k: v[batch_idx] for k, v in vit_kwargs.items()} + image_embed = self.vision_model(image, **cur_kwargs).last_hidden_state + image_embed = self.vision_aligner(image_embed) + n, seq_len, dim = image_embed.shape + image_embed = image_embed.reshape(n * seq_len, dim) + cond_vit_image_embeds.append(image_embed) + + # 2. Instantiate the vit image embeddings into the input sequence + batch_size, seq_len, n_embd = x.shape + index = torch.arange(seq_len, device=x.device).unsqueeze(0).repeat(batch_size, 1) + + for i, (image_embed, mask) in enumerate(zip(cond_vit_image_embeds, cond_vit_image_mask)): + image_scatter_index = index[i : i + 1].masked_select(mask.bool()).reshape(1, -1) + x[i : i + 1].scatter_( + dim=1, + index=image_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd), + src=image_embed.reshape(1, -1, n_embd), + ) + + return x + + def ragged_final_layer(self, x, image_mask, timestep, token_h, token_w, first_step): + bsz, seq_len, n_embd = x.shape + if first_step: + image_output = x.masked_select(image_mask.unsqueeze(-1).bool()).reshape(bsz, -1, n_embd) + else: + image_output = x[:, 1:, :] + timestep_emb = self.time_embed_2(timestep) + pred = self.final_layer(image_output, timestep_emb, token_h, token_w) + return pred + + @staticmethod + def build_batch_rope_image_info(output, sections): + rope_image_info = [] + for image_slices, sections_i in zip(output.all_image_slices, sections): + image_shapes = [] + for section in sections_i: + if "image" in section["type"]: + if isinstance(section["token_height"], list): + assert len(section["token_height"]) == len(section["token_width"]), ( + f"token_height and token_width should have the same length, " + f"but got {len(section['token_height'])} and {len(section['token_width'])}" + ) + image_shapes.extend(list(zip(section["token_height"], section["token_width"]))) + else: + image_shapes.append((section["token_height"], section["token_width"])) + assert len(image_slices) == len(image_shapes), ( + f"Size miss matching: Image slices({len(image_slices)}) != image shapes({len(image_shapes)})" + ) + rope_image_info.append(list(zip(image_slices, image_shapes))) + return rope_image_info + + def vae_encode(self, image, cfg_factor=1): + config = self.vae.config + + with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=True): + vae_encode_result = self.vae.encode(image) + if isinstance(vae_encode_result, torch.Tensor): + latents = vae_encode_result + else: + latents = vae_encode_result.latent_dist.sample() + if hasattr(config, "shift_factor") and config.shift_factor: + latents.sub_(config.shift_factor) + if hasattr(config, "scaling_factor") and config.scaling_factor: + latents.mul_(config.scaling_factor) + + if hasattr(self.vae, "ffactor_temporal"): + assert latents.shape[2] == 1, "latents should have shape [B, C, T, H, W] and T should be 1" + latents = latents.squeeze(2) + + # Here we always use t=0 to declare it is a clean conditional image + t = torch.zeros((latents.shape[0],)) + + if cfg_factor > 1: + t = t.repeat(cfg_factor) + latents = latents.repeat(cfg_factor, 1, 1, 1) + + return t, latents + + def _encode_cond_image( + self, + batch_cond_image_info_list: list[list[JointImageInfo]], + cfg_factor: int = 1, + ): + # VAE encode one by one, as we assume cond images have different sizes + batch_cond_vae_images, batch_cond_t, batch_cond_vit_images = [], [], [] + for cond_image_info_list in batch_cond_image_info_list: + cond_vae_image_list, cond_t_list, cond_vit_image_list = [], [], [] + for image_info in cond_image_info_list: + cond_t_, cond_vae_image_ = self.vae_encode( + image_info.vae_image_info.image_tensor.to(self.device), + ) + cond_vit_image_list.append(image_info.vision_image_info.image_tensor) + cond_vae_image_list.append(cond_vae_image_.squeeze(0)) + cond_t_list.append(cond_t_) + batch_cond_vae_images.append(cond_vae_image_list) + batch_cond_t.append(cond_t_list) + batch_cond_vit_images.append(torch.cat(cond_vit_image_list, dim=0)) + + # If only one cond image for each sample and all have the same size, we can batch them together + # In this case, cond_vae_images is a 4-D tensor. + if all([len(items) == 1 for items in batch_cond_vae_images]) and all( + items[0].shape == batch_cond_vae_images[0][0].shape for items in batch_cond_vae_images + ): + cond_vae_images = torch.stack([items[0] for items in batch_cond_vae_images], dim=0) + cond_t = torch.cat([items[0] for items in batch_cond_t], dim=0) + if cfg_factor > 1: + cond_t = cond_t.repeat(cfg_factor) + cond_vae_images = cond_vae_images.repeat(cfg_factor, 1, 1, 1) + else: + # In this case, cond_vae_images is a list of 4-D tensors or a list of lists of 3-D tensors. + cond_t = [torch.cat(item, dim=0) for item in batch_cond_t] + cond_vae_images = [] + for items in batch_cond_vae_images: + if all(items[0].shape == item.shape for item in items): + cond_vae_images.append(torch.stack(items, dim=0)) + else: + cond_vae_images.append(items) + if cfg_factor > 1: + cond_t = cond_t * cfg_factor + cond_vae_images = cond_vae_images * cfg_factor + + if cfg_factor > 1: + batch_cond_vit_images = batch_cond_vit_images * cfg_factor + + return cond_vae_images, cond_t, batch_cond_vit_images + + @staticmethod + def check_inputs(prompt=None, message_list=None): + if prompt is None and message_list is None: + raise ValueError("Either `prompt` or `message_list` should be provided.") + if prompt is not None and message_list is not None: + raise ValueError("Only one of `prompt` or `message_list` should be provided.") + if prompt is not None: + assert isinstance(prompt, str) or isinstance(prompt, list), ( + f"`prompt` should be a string or a list of strings, but got {type(prompt)}." + ) + if isinstance(prompt, list): + assert len(prompt) > 0 and all(isinstance(p, str) for p in prompt), ( + "`prompt` should be a non-empty list of strings." + ) + if message_list is not None: + if not isinstance(message_list, list): + raise ValueError(f"`message_list` should be a list of messages, but got {type(message_list)}.") + assert len(message_list) > 0, "`message_list` should be a non-empty list." + for message in message_list: + assert isinstance(message, list) or isinstance(message, dict), ( + f"Each message should be a list of dicts or a dict, but got {type(message)}." + ) + + def prepare_model_inputs( + self, + prompt=None, + mode="gen_image", + system_prompt=None, + cot_text=None, + num_inference_steps=50, + guidance_scale=5.0, + image_size="auto", + message_list=None, + device=None, + max_new_tokens=None, + **kwargs, + ): + # 1. Sanity check + self.check_inputs(prompt, message_list) + device = default(device, self.device) + + # 2. Format inputs + batch_message_list = message_list + batch_prompt = prompt + batch_cot_text = cot_text + batch_system_prompt = system_prompt + batch_gen_image_info = None + # TODO: construct with user input images + batch_cond_image_info = None + + # -- 2.1 message_list + if batch_message_list is not None: + if isinstance(batch_message_list[0], dict): + batch_message_list = [batch_message_list] + batch_size = len(batch_message_list) + + batch_gen_image_info = [ + [message["content"] for message in message_list_ if message["type"] == "gen_image"] + for message_list_ in batch_message_list + ] + # At most one gen_image is allowed for each message_list + batch_gen_image_info = [info[-1] if len(info) > 0 else None for info in batch_gen_image_info] + # Multiple cond images are allowed. + batch_cond_image_info = [ + [message["content"] for message in message_list_ if message["type"] == "joint_image"] + for message_list_ in batch_message_list + ] + + # -- 2.2 Prompt, cot text, system prompt + else: + if isinstance(batch_prompt, str): + batch_prompt = [batch_prompt] + batch_size = len(batch_prompt) + + if batch_cot_text is not None: + if isinstance(batch_cot_text, str): + batch_cot_text = [batch_cot_text] + else: + assert isinstance(batch_cot_text, list) and len(batch_cot_text) == batch_size, ( + "`cot_text` should be a string or a list of strings with the same length as `prompt`." + ) + + if batch_system_prompt is not None: + if isinstance(batch_system_prompt, str): + batch_system_prompt = [batch_system_prompt] + else: + assert isinstance(batch_system_prompt, list) and len(batch_system_prompt) == batch_size, ( + "`system_prompts` should be a string or a list of strings with the same length as `prompt`." + ) + + if mode == "gen_image": + batch_gen_image_info = [self.image_processor.build_image_info(image_size) for _ in range(batch_size)] + + # -- 2.3 seed + generator = kwargs.get("generator", None) + if generator is None: + seeds = self.prepare_seed(seed=kwargs.get("seed"), batch_size=batch_size) + generator = [torch.Generator(self.device).manual_seed(seed) for seed in seeds] + + # 3. apply chat template + cfg_factor = {"gen_text": 1, "gen_image": 2} + bot_task = kwargs.pop("bot_task", "auto") + # If `drop_think` enabled, always drop parts in the context. + drop_think = kwargs.get("drop_think", self.generation_config.drop_think) + # Apply batched prompt or batched message_list to build input sequence with associated info. + out = self._tkwrapper.apply_chat_template( + batch_prompt=batch_prompt, + batch_message_list=batch_message_list, + mode=mode, + batch_gen_image_info=batch_gen_image_info, + batch_cond_image_info=batch_cond_image_info, + batch_system_prompt=batch_system_prompt, + batch_cot_text=batch_cot_text, + max_length=kwargs.get("max_length"), + bot_task=bot_task, + image_base_size=self.config.image_base_size, + sequence_template="pretrain", + cfg_factor=cfg_factor[mode], + drop_think=drop_think, + ) + output, sections = out["output"], out["sections"] + + # 4. Encode conditional images + if batch_cond_image_info is not None and len(batch_cond_image_info[0]) > 0: + cond_vae_images, cond_timestep, cond_vit_images = self._encode_cond_image( + batch_cond_image_info, cfg_factor[mode] + ) + # Pack vit kwargs. Siglip2-so requires spatial_shapes and attention_mask for inference. + vit_kwargs = {"spatial_shapes": [], "attention_mask": []} + for cond_image_info in batch_cond_image_info: + vit_kwargs["spatial_shapes"].append( + torch.stack([item.vision_encoder_kwargs["spatial_shapes"] for item in cond_image_info]) + ) + vit_kwargs["attention_mask"].append( + torch.stack([item.vision_encoder_kwargs["pixel_attention_mask"] for item in cond_image_info]) + ) + if cfg_factor[mode] > 1: + vit_kwargs["spatial_shapes"] = vit_kwargs["spatial_shapes"] * cfg_factor[mode] + vit_kwargs["attention_mask"] = vit_kwargs["attention_mask"] * cfg_factor[mode] + else: + cond_vae_images, cond_timestep, cond_vit_images = None, None, None + vit_kwargs = None + + # 5. Build position embeddings + rope_image_info = self.build_batch_rope_image_info(output, sections) + if mode == "gen_text": + seq_len = self.generation_config.max_length + else: + seq_len = output.tokens.shape[1] + cos, sin = build_batch_2d_rope( + image_infos=rope_image_info, + seq_len=seq_len, + n_elem=self.config.attention_head_dim, + device=device, + base=self.config.rope_theta, + ) + + # 6. Build kv cache + if bot_task == "img_ratio": + max_new_tokens = 1 + + # 7. Build position ids + batch_input_pos = torch.arange(0, output.tokens.shape[1], dtype=torch.long, device=device)[None].expand( + batch_size * cfg_factor[mode], -1 + ) # use expand to share indices to save memory + + # 8. Build model input kwargs + tkw = self._tkwrapper + if image_size == "auto": + extra_auto_stops = [tkw.special_token_map[f""] for i in range(33)] + else: + extra_auto_stops = [tkw.boi_token_id] + stop_token_id = dict( + auto=[tkw.eos_token_id] + extra_auto_stops, + image=[tkw.eos_token_id], + recaption=[tkw.end_recaption_token_id, tkw.end_answer_token_id, tkw.eos_token_id], + think=[tkw.end_recaption_token_id, tkw.end_answer_token_id, tkw.eos_token_id], + img_ratio=extra_auto_stops, + ) + model_input_kwargs = dict( + input_ids=output.tokens.to(device), + position_ids=batch_input_pos, + past_key_values=None, + custom_pos_emb=(cos, sin), + mode=mode, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + image_mask=to_device(output.gen_image_mask, device), + gen_timestep_scatter_index=to_device(output.gen_timestep_scatter_index, device), + cond_vae_images=to_device(cond_vae_images, device), + cond_timestep=to_device(cond_timestep, device), + cond_vae_image_mask=to_device(output.cond_vae_image_mask, device), + cond_vit_images=to_device(cond_vit_images, device), + cond_vit_image_mask=to_device(output.cond_vit_image_mask, device), + vit_kwargs={k: to_device(v, self.device) for k, v in vit_kwargs.items()} + if vit_kwargs is not None + else None, + cond_timestep_scatter_index=to_device(output.cond_timestep_scatter_index, device), + # for inner usage + tokenizer_output=output, + batch_gen_image_info=batch_gen_image_info, + generator=generator, + # generation config + eos_token_id=stop_token_id[bot_task], + max_new_tokens=max_new_tokens, + ) + + return model_input_kwargs + + def _prepare_attention_mask_for_generation( + self, + inputs_tensor: torch.Tensor, + generation_config: GenerationConfig, + model_kwargs: dict[str, Any], + ) -> torch.Tensor: + # create `4d` bool attention mask (b, 1, seqlen, seqlen) using this implementation to bypass the 2d requirement + # in the `transformers.generation_utils.GenerationMixin.generate`. + # This implementation can handle sequences with text and image modalities, where text tokens use causal + # attention and image tokens use full attention. + bsz, seq_len = inputs_tensor.shape + tokenizer_output = model_kwargs["tokenizer_output"] + batch_image_slices = [ + tokenizer_output.joint_image_slices[i] + tokenizer_output.gen_image_slices[i] for i in range(bsz) + ] + attention_mask = torch.ones(seq_len, seq_len, dtype=torch.bool).tril(diagonal=0).repeat(bsz, 1, 1) + for i in range(bsz): + for j, image_slice in enumerate(batch_image_slices[i]): + attention_mask[i, image_slice, image_slice] = True + attention_mask = attention_mask.unsqueeze(1) + return attention_mask + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + tokenizer_output=None, + batch_gen_image_info=None, + generator=None, + **kwargs, + ): + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + if input_ids.shape[1] != kwargs["position_ids"].shape[1]: # in decode steps + input_ids = torch.gather(input_ids, dim=1, index=kwargs["position_ids"]) + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "attention_mask": attention_mask, + "position_ids": kwargs["position_ids"], + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "custom_pos_emb": kwargs["custom_pos_emb"], + "mode": kwargs["mode"], + "images": kwargs.get("images"), + "image_mask": kwargs.get("image_mask"), + "timestep": kwargs.get("timestep"), + "gen_timestep_scatter_index": kwargs.get("gen_timestep_scatter_index"), + "cond_vae_images": kwargs.get("cond_vae_images"), + "cond_timestep": kwargs.get("cond_timestep"), + "cond_vae_image_mask": kwargs.get("cond_vae_image_mask"), + "cond_vit_images": kwargs.get("cond_vit_images"), + "cond_vit_image_mask": kwargs.get("cond_vit_image_mask"), + "vit_kwargs": kwargs.get("vit_kwargs"), + "cond_timestep_scatter_index": kwargs.get("cond_timestep_scatter_index"), + "query_lens": kwargs.get("query_lens"), + "seq_lens": kwargs.get("seq_lens"), + "num_image_tokens": kwargs.get("num_image_tokens"), + } + ) + return model_inputs + + def _update_model_kwargs_for_generation( + self, + outputs: ModelOutput, + model_kwargs: dict[str, Any], + is_encoder_decoder: bool = False, + num_new_tokens: int = 1, + ) -> dict[str, Any]: + mode = model_kwargs["mode"] + + updated_model_kwargs = { + "mode": mode, + "custom_pos_emb": model_kwargs["custom_pos_emb"], + } + + # update past_key_values keeping its naming used in model code + for possible_cache_name in ALL_CACHE_NAMES: + if possible_cache_name in outputs: + # TODO (joao): remove output/input mismatch when these old models (xlnet, reformer) are deprecated + if possible_cache_name in ("past_buckets_states", "mems"): + cache_name = "past_key_values" + else: + cache_name = possible_cache_name + updated_model_kwargs[cache_name] = getattr(outputs, possible_cache_name) + break + + if "tokenizer_output" in model_kwargs: + if mode == "gen_text": + # When enable batching, we use right padding, which requires a real_pos to index the valid + # end position of the sequence. If tokenizer_output in model_kwargs, it means we are in the + # prefill step of generation. + real_pos = to_device(model_kwargs["tokenizer_output"].real_pos, self.device) + updated_model_kwargs["position_ids"] = real_pos + else: + # position ids + image_mask = model_kwargs["image_mask"] + bsz, seq_len = image_mask.shape + index = torch.arange(seq_len, device=image_mask.device).unsqueeze(0).repeat(bsz, 1) + position_ids = index.masked_select(image_mask.bool()).reshape(bsz, -1) + timestep_position_ids = index[ + torch.arange(bsz), model_kwargs["gen_timestep_scatter_index"][:, -1] + ].unsqueeze(-1) + updated_model_kwargs["position_ids"] = torch.cat([timestep_position_ids, position_ids], dim=1) + + # attention mask + mask_list = [] + for attention_mask_i, position_ids_i in zip( + model_kwargs["attention_mask"], updated_model_kwargs["position_ids"] + ): + mask_list.append(torch.index_select(attention_mask_i, dim=1, index=position_ids_i.reshape(-1))) + attention_mask = torch.stack(mask_list, dim=0) + updated_model_kwargs["attention_mask"] = attention_mask + updated_model_kwargs["gen_timestep_scatter_index"] = model_kwargs["gen_timestep_scatter_index"] + + else: + if mode == "gen_text": + # Now we are in the decode steps. + updated_model_kwargs["position_ids"] = model_kwargs["position_ids"] + 1 + else: + updated_model_kwargs["position_ids"] = model_kwargs["position_ids"] + updated_model_kwargs["attention_mask"] = model_kwargs["attention_mask"] + updated_model_kwargs["gen_timestep_scatter_index"] = model_kwargs["gen_timestep_scatter_index"] + + return updated_model_kwargs + + def _generate( + self, + generator: list[torch.Generator] | None = None, + **kwargs, + ): + mode = kwargs.get("mode", "gen_text") + # verbose > 1 not support + if mode == "gen_text": + raise NotImplementedError("Not support gen text for hunyuan image") + + elif mode == "gen_image": + batch_gen_image_info: list[ImageInfo] = kwargs.get("batch_gen_image_info") + if batch_gen_image_info is None: + raise ValueError("`batch_gen_image_info` should be provided when `mode` is `gen_image`.") + + image_info: ImageInfo = batch_gen_image_info[0] + num_image_tokens = ( + image_info.image_token_length + + (1 if image_info.add_timestep_token else 0) + + (1 if image_info.add_guidance_token else 0) + ) + kwargs["num_image_tokens"] = num_image_tokens + # 50 and 5.0 hard code + from vllm.forward_context import set_forward_context + + with set_forward_context(None, self.vllm_config): + results = self.pipeline( + batch_size=len(batch_gen_image_info), + image_size=[batch_gen_image_info[0].image_height, batch_gen_image_info[0].image_width], + num_inference_steps=kwargs.get("num_inference_steps", 50), + guidance_scale=kwargs.get("guidance_scale", 5.0), + generator=generator, + model_kwargs=kwargs, + ) + samples = results[0] + return samples + + else: + raise ValueError(f"Unknown mode {mode}, only `gen_text` and `gen_image` are supported.") + + @staticmethod + def _check_inputs(cond, target, check_list): + if cond: + for name, item in check_list: + assert item is not None, f"`{name}` should be provided when `{target}`." + + def forward_call( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + custom_pos_emb: tuple[torch.FloatTensor] | None = None, + mode: str = "gen_text", + first_step: bool | None = None, + # for gen image + images: BatchRaggedImages | None = None, + image_mask: torch.Tensor | None = None, + timestep: BatchRaggedTensor | None = None, + gen_timestep_scatter_index: torch.Tensor | None = None, + # for cond image + cond_vae_images: BatchRaggedImages | None = None, + cond_timestep: BatchRaggedTensor | None = None, + cond_vae_image_mask: torch.Tensor | None = None, + cond_vit_images: BatchRaggedImages | None = None, + cond_vit_image_mask: torch.Tensor | None = None, + vit_kwargs: dict[str, Any] | None = None, + cond_timestep_scatter_index: torch.Tensor | None = None, + query_lens: list[int] | None = None, + seq_lens: list[int] | None = None, + num_image_tokens: int | None = None, + ) -> tuple | CausalMMOutputWithPast: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + # Sanity Check of Inputs + self._check_inputs( + mode == "gen_image", + "in `gen_image` mode", + [ + ("images", images), + ("timestep", timestep), + ("gen_timestep_scatter_index", gen_timestep_scatter_index), + ], + ) + self._check_inputs( + mode == "gen_image" and first_step, + "in `gen_image` mode at the first step", + [ + ("image_mask", image_mask), + ], + ) + self._check_inputs( + cond_vae_images is not None, + "`cond_vae_images` is provided", + [ + ("cond_timestep", cond_timestep), + ("cond_vae_image_mask", cond_vae_image_mask), + ("cond_timestep_scatter_index", cond_timestep_scatter_index), + ], + ) + self._check_inputs( + cond_vit_images is not None, + "`cond_vit_images` is provided", + [ + ("cond_vit_image_mask", cond_vit_image_mask), + ("vit_kwargs", vit_kwargs), + ], + ) + + custom_pos_emb = self.get_pos_emb(custom_pos_emb, position_ids) + + inputs_embeds = self.model.wte(input_ids) + + bsz, seq_len, n_embd = inputs_embeds.shape + + # Instantiate placeholder tokens: , for the gen image + if mode == "gen_text": + # For gen_text, make sure gen_timestep_scatter_index is None + gen_timestep_scatter_index = None + token_h, token_w = None, None + else: + if first_step: + inputs_embeds, token_h, token_w = self.instantiate_vae_image_tokens( + inputs_embeds, images, timestep, image_mask + ) + inputs_embeds = self.instantiate_timestep_tokens(inputs_embeds, timestep, gen_timestep_scatter_index) + else: + t_emb = self.time_embed(timestep) + image_emb, token_h, token_w = self.patch_embed(images, t_emb) + timestep_emb = self.timestep_emb(timestep).reshape(bsz, -1, n_embd) + inputs_embeds = torch.cat([timestep_emb, image_emb], dim=1) + + # Instantiate placeholder tokens: , for cond images + # Should only run once with kv-cache enabled. + if cond_vae_images is not None: + inputs_embeds, _, _ = self.instantiate_vae_image_tokens( + inputs_embeds, cond_vae_images, cond_timestep, cond_vae_image_mask + ) + inputs_embeds = self.instantiate_timestep_tokens(inputs_embeds, cond_timestep, cond_timestep_scatter_index) + if cond_vit_images is not None: + inputs_embeds = self.instantiate_vit_image_tokens( + inputs_embeds, cond_vit_images, cond_vit_image_mask, vit_kwargs + ) + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + custom_pos_emb=custom_pos_emb, + mode=mode, + first_step=first_step, + query_lens=query_lens, + seq_lens=seq_lens, + num_image_tokens=num_image_tokens, + gen_timestep_scatter_index=gen_timestep_scatter_index, + ) + hidden_states = outputs[0] + + if mode == "gen_text": + hidden_states = self.model.ln_f(hidden_states) + logits = self.lm_head(hidden_states) + logits = logits.float() + diffusion_prediction = None + else: + logits = None + hidden_states = hidden_states.to(input_ids.device) + assert hidden_states.numel() == bsz * seq_len * n_embd, ( + f"Shape mismatch: {hidden_states.shape} cannot reshape to ({bsz}, {seq_len}, {n_embd})" + ) + hidden_states = hidden_states.reshape(bsz, seq_len, n_embd) + diffusion_prediction = self.ragged_final_layer( + hidden_states, image_mask, timestep, token_h, token_w, first_step + ) + + if not return_dict: + output = (logits,) + outputs[1:] + (diffusion_prediction,) + return output + + output = CausalMMOutputWithPast( + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + diffusion_prediction=diffusion_prediction, + ) + + return output + + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | list[str] = "", + image_size="auto", + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 50, + guidance_scale: float = 5.0, + system_prompt: str | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + **kwargs, + ) -> DiffusionOutput: + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + generator = req.sampling_params.generator or generator + height = req.sampling_params.height or height + width = req.sampling_params.width or width + num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps + if req.sampling_params.guidance_scale_provided: + guidance_scale = req.sampling_params.guidance_scale + if guidance_scale <= 1.0: + logger.warning("HunyuanImage3.0 does not support guidance_scale <= 1.0, will set it to 1.0 + epsilon.") + guidance_scale = 1.0 + np.finfo(float).eps + image_size = (height, width) + model_inputs = self.prepare_model_inputs( + prompt=prompt, + cot_text=None, + system_prompt=system_prompt, + mode="gen_image", + generator=generator, + image_size=image_size, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + ) + outputs = self._generate(**model_inputs, **kwargs) + return DiffusionOutput(output=outputs[0]) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 024352479..3d236ac31 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -85,6 +85,11 @@ _DIFFUSION_MODELS = { "pipeline_sd3", "StableDiffusion3Pipeline", ), + "HunyuanImage3ForCausalMM": ( + "hunyuan_image_3", + "pipeline_hunyuan_image_3", + "HunyuanImage3Pipeline", + ), "Flux2KleinPipeline": ( "flux2_klein", "pipeline_flux2_klein", diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 38f438062..0e12460e8 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -87,18 +87,17 @@ class DiffusionModelRunner: return nullcontext() # Load model within forward context - with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config): - load_config = LoadConfig() - model_loader = DiffusersPipelineLoader(load_config) - time_before_load = time.perf_counter() + load_config = LoadConfig() + model_loader = DiffusersPipelineLoader(load_config) + time_before_load = time.perf_counter() - with get_memory_context(): - with DeviceMemoryProfiler() as m: - self.pipeline = model_loader.load_model( - od_config=self.od_config, - load_device=load_device, - ) - time_after_load = time.perf_counter() + with get_memory_context(): + with DeviceMemoryProfiler() as m: + self.pipeline = model_loader.load_model( + od_config=self.od_config, + load_device=load_device, + ) + time_after_load = time.perf_counter() logger.info( "Model loading took %.4f GiB and %.6f seconds", diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 0fad90a5b..6f5bd00e6 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -13,10 +13,11 @@ from contextlib import AbstractContextManager, nullcontext import torch import zmq -from vllm.config import VllmConfig +from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config from vllm.distributed.device_communicators.shm_broadcast import MessageQueue from vllm.logger import init_logger from vllm.utils.mem_utils import GiB_bytes +from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.diffusion.data import ( DiffusionOutput, @@ -84,13 +85,16 @@ class DiffusionWorker: current_omni_platform.set_device(self.device) # Create vllm_config for parallel configuration - vllm_config = VllmConfig() + vllm_config = VllmConfig(compilation_config=CompilationConfig()) vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size self.vllm_config = vllm_config # Initialize distributed environment - with set_forward_context(vllm_config=vllm_config, omni_diffusion_config=self.od_config): + with ( + set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config), + set_current_vllm_config(self.vllm_config), + ): init_distributed_environment(world_size=world_size, rank=rank) logger.info(f"Worker {self.rank}: Initialized device and distributed environment.") @@ -104,16 +108,17 @@ class DiffusionWorker: tensor_parallel_size=parallel_config.tensor_parallel_size, pipeline_parallel_size=parallel_config.pipeline_parallel_size, ) + init_workspace_manager(self.device) - # Create model runner and load model - self.model_runner = DiffusionModelRunner( - vllm_config=self.vllm_config, - od_config=self.od_config, - device=self.device, - ) - self.model_runner.load_model( - memory_pool_context_fn=self._maybe_get_memory_pool_context, - ) + # Create model runner and load model + self.model_runner = DiffusionModelRunner( + vllm_config=self.vllm_config, + od_config=self.od_config, + device=self.device, + ) + self.model_runner.load_model( + memory_pool_context_fn=self._maybe_get_memory_pool_context, + ) assert self.model_runner.pipeline is not None self.lora_manager = DiffusionLoRAManager( pipeline=self.model_runner.pipeline, diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py index 535f04f7d..11a3c07e1 100644 --- a/vllm_omni/entrypoints/async_omni_diffusion.py +++ b/vllm_omni/entrypoints/async_omni_diffusion.py @@ -98,6 +98,10 @@ class AsyncOmniDiffusion: od_config.model_class_name = "BagelPipeline" od_config.tf_model_config = TransformerConfig() od_config.update_multimodal_support() + elif architectures and len(architectures) == 1: + od_config.model_class_name = architectures[0] + else: + raise # Initialize engine self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config) diff --git a/vllm_omni/entrypoints/omni_diffusion.py b/vllm_omni/entrypoints/omni_diffusion.py index 5ad9a91c8..9f6dde15b 100644 --- a/vllm_omni/entrypoints/omni_diffusion.py +++ b/vllm_omni/entrypoints/omni_diffusion.py @@ -80,6 +80,8 @@ class OmniDiffusion: od_config.model_class_name = "BagelPipeline" od_config.tf_model_config = TransformerConfig() od_config.update_multimodal_support() + elif architectures and len(architectures) == 1: + od_config.model_class_name = architectures[0] else: raise