feat: Add initial model files

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
+utils/__pycache__/import_utils.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text
+utils/__pycache__/import_utils.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text

__init__.py

@@ -0,0 +1,18 @@
+from typing import TYPE_CHECKING
+
+from utils import _LazyModule
+from utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_hunyuan_image_3 import *
+    from .modeling_hunyuan_image_3 import *
+    from .autoencoder_kl_3d import *
+    from .image_processor import *
+    from .siglip2 import *
+    from .tokenization_hunyuan_image_3 import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

autoencoder_kl_3d.py

@@ -0,0 +1,1081 @@
+"""
+Reference code
+[FLUX] https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/autoencoder.py
+[DCAE] https://github.com/mit-han-lab/efficientvit/blob/master/efficientvit/models/efficientvit/dc_ae.py
+"""
+import os
+from dataclasses import dataclass
+from typing import Tuple, Optional
+import math
+import random
+import numpy as np
+from einops import rearrange
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from safetensors import safe_open
+import os
+from collections import OrderedDict
+from collections.abc import Iterable
+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.torch_utils import randn_tensor
+from diffusers.utils import BaseOutput
+
+
+
+class DiagonalGaussianDistribution(object):
+    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: Optional[torch.Generator] = 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
+
+    def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        else:
+            reduce_dim = list(range(1, self.mean.ndim))
+            if other is None:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+                    dim=reduce_dim,
+                )
+            else:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean - other.mean, 2) / other.var +
+                    self.var / other.var -
+                    1.0 -
+                    self.logvar +
+                    other.logvar,
+                    dim=reduce_dim,
+                )
+
+    def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims,
+        )
+
+    def mode(self) -> torch.Tensor:
+        return self.mean
+
+@dataclass
+class DecoderOutput(BaseOutput):
+    sample: torch.FloatTensor
+    posterior: Optional[DiagonalGaussianDistribution] = 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):
+    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 Downsample(nn.Module):
+    def __init__(self, in_channels: int, add_temporal_downsample: bool = True):
+        super().__init__()
+        self.add_temporal_downsample = add_temporal_downsample
+        stride = (2, 2, 2) if add_temporal_downsample else (1, 2, 2)  # THW
+        # no asymmetric padding in torch conv, must do it ourselves
+        self.conv = Conv3d(in_channels, in_channels, kernel_size=3, stride=stride, padding=0)
+
+    def forward(self, x: Tensor):
+        spatial_pad = (0, 1, 0, 1, 0, 0)  # WHT
+        x = nn.functional.pad(x, spatial_pad, mode="constant", value=0)
+
+        temporal_pad = (0, 0, 0, 0, 0, 1) if self.add_temporal_downsample else (0, 0, 0, 0, 1, 1)
+        x = nn.functional.pad(x, temporal_pad, mode="replicate")
+
+        x = self.conv(x)
+        return x
+
+
+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 Upsample(nn.Module):
+    def __init__(self, in_channels: int, add_temporal_upsample: bool = True):
+        super().__init__()
+        self.add_temporal_upsample = add_temporal_upsample
+        self.scale_factor = (2, 2, 2) if add_temporal_upsample else (1, 2, 2)  # THW
+        self.conv = Conv3d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor):
+        x = nn.functional.interpolate(x, scale_factor=self.scale_factor, mode="nearest")
+        x = self.conv(x)
+        return x
+
+
+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):
+    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:
+        with torch.no_grad():
+            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):
+    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:
+        with torch.no_grad():
+            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):
+    _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,
+        shift_factor: Optional[float] = None,
+        downsample_match_channel: bool = True,
+        upsample_match_channel: bool = True,
+        only_encoder: bool = False,
+        only_decoder: bool = False,
+    ):
+        super().__init__()
+        self.ffactor_spatial = ffactor_spatial
+        self.ffactor_temporal = ffactor_temporal
+        self.scaling_factor = scaling_factor
+        self.shift_factor = shift_factor
+
+        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,
+            )
+
+        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.125
+
+        self.use_compile = False
+
+        self.empty_cache = torch.empty(0, device="cuda")
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (Encoder, Decoder)):
+            module.gradient_checkpointing = value
+
+    def enable_tiling_during_training(self, use_tiling: bool = True):
+        self.use_tiling_during_training = use_tiling
+
+    def disable_tiling_during_training(self):
+        self.enable_tiling_during_training(False)
+
+    def enable_temporal_tiling(self, use_tiling: bool = True):
+        self.use_temporal_tiling = use_tiling
+
+    def disable_temporal_tiling(self):
+        self.enable_temporal_tiling(False)
+
+    def enable_spatial_tiling(self, use_tiling: bool = True):
+        self.use_spatial_tiling = use_tiling
+
+    def disable_spatial_tiling(self):
+        self.enable_spatial_tiling(False)
+
+    def enable_tiling(self, use_tiling: bool = True):
+        self.enable_spatial_tiling(use_tiling)
+
+    def disable_tiling(self):
+        self.disable_spatial_tiling()
+
+    def enable_slicing(self):
+        self.use_slicing = True
+
+    def disable_slicing(self):
+        self.use_slicing = False
+
+    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):
+        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):
+        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):
+        B, C, T, H, W = z.shape
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))  # 24 * (1 - 0.125) = 21
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)  # 384 * 0.125 = 48
+        row_limit = self.tile_sample_min_size - blend_extent  # 384 - 48 = 336
+
+        # 分布式/多卡：输入不做 padding -> 每 rank 对解码输出做右/下 padding -> GPU all_gather -> rank0重组/融合/裁剪
+        if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+            rank = dist.get_rank()
+            world_size = dist.get_world_size()
+
+            # 统计tile
+            num_rows = math.ceil(H / overlap_size)
+            num_cols = math.ceil(W / overlap_size)
+            total_tiles = num_rows * num_cols
+            tiles_per_rank = math.ceil(total_tiles / world_size)
+
+            print(f"==={torch.distributed.get_rank()},  {total_tiles=}, {tiles_per_rank=}, {world_size=}")
+
+            # 本 rank 的 tile 索引（循环分配）：rank, rank+world_size,
+            my_linear_indices = list(range(rank, total_tiles, world_size))
+            if my_linear_indices == []:
+                my_linear_indices = [0]
+            print(f"==={torch.distributed.get_rank()},  {my_linear_indices=}")
+            decoded_tiles = [] # tiles
+            decoded_metas = [] # (ri, rj, pad_w, pad_h)
+            H_out_std = self.tile_sample_min_size
+            W_out_std = self.tile_sample_min_size
+            for lin_idx in my_linear_indices:
+                ri = lin_idx // num_cols
+                rj = lin_idx % num_cols
+                i = ri * overlap_size
+                j = rj * overlap_size
+                tile = z[
+                    :,
+                    :,
+                    :,
+                    i : i + self.tile_latent_min_size,
+                    j : j + self.tile_latent_min_size,
+                ]
+                dec = self.decoder(tile)
+                # 对边界 tile 的输出做右/下方向 padding 到标准尺寸
+                pad_h = max(0, H_out_std - dec.shape[-2])
+                pad_w = max(0, W_out_std - dec.shape[-1])
+                if pad_h > 0 or pad_w > 0:
+                    dec = F.pad(dec, (0, pad_w, 0, pad_h, 0, 0), "constant", 0)
+                decoded_tiles.append(dec)
+                decoded_metas.append(torch.tensor([ri, rj, pad_w, pad_h], device=z.device, dtype=torch.int64))
+            
+            # 各rank数量不一定相同，进行padding到相同长度
+            T_out = decoded_tiles[0].shape[2] if len(decoded_tiles) > 0 else (T-1)*self.ffactor_temporal+1
+            while len(decoded_tiles) < tiles_per_rank:
+                decoded_tiles.append(torch.zeros([1, 3, T_out, self.tile_sample_min_size, self.tile_sample_min_size], device=z.device, dtype=dec.dtype))
+                decoded_metas.append(torch.tensor([-1, -1, self.tile_sample_min_size, self.tile_sample_min_size], device=z.device, dtype=torch.int64)) 
+                
+            # 进行gpu的all_gather
+            decoded_tiles = torch.stack(decoded_tiles, dim=0)
+            decoded_metas = torch.stack(decoded_metas, dim=0)
+            
+            tiles_gather_list = [torch.empty_like(decoded_tiles) for _ in range(world_size)]
+            metas_gather_list = [torch.empty_like(decoded_metas) for _ in range(world_size)]
+
+            dist.all_gather(tiles_gather_list, decoded_tiles)
+            dist.all_gather(metas_gather_list, decoded_metas)
+
+            if rank != 0:
+                # 非0号rank返回空占位，结果只在rank0上有效
+                return torch.empty(0, device=z.device)
+
+            # rank0：根据 (ri, rj) 元信息重建 tile 网格；跳过占位项 (ri, rj) == (-1, -1)
+            rows = [[None for _ in range(num_cols)] for _ in range(num_rows)]
+            for r in range(world_size):
+                gathered_tiles_r = tiles_gather_list[r]  # [tiles_per_rank, B, C, T, H, W]
+                gathered_metas_r = metas_gather_list[r]  # [tiles_per_rank, 4]，元素: (ri, rj, pad_w, pad_h)
+                for k in range(gathered_tiles_r.shape[0]):
+                    ri = int(gathered_metas_r[k][0])
+                    rj = int(gathered_metas_r[k][1])
+                    if ri < 0 or rj < 0:
+                        continue
+                    if ri < num_rows and rj < num_cols:
+                        # 去除padding
+                        pad_w = int(gathered_metas_r[k][2])
+                        pad_h = int(gathered_metas_r[k][3])
+                        h_end = None if pad_h == 0 else -pad_h
+                        w_end = None if pad_w == 0 else -pad_w
+                        rows[ri][rj] = gathered_tiles_r[k][:, :, :, :h_end, :w_end]
+
+            result_rows = []
+            for i, row in enumerate(rows):
+                result_row = []
+                for j, tile in enumerate(row):
+                    if tile is None:
+                        continue
+                    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
+
+        # 单卡：原有串行逻辑
+        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):
+        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):
+
+        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):
+
+        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 torch.distributed.is_initialized():
+            if torch.distributed.get_rank() != 0:
+                return self.empty_cache
+
+        if z.shape[-3] == 1:
+            decoded = decoded[:, :, -1:]
+        if not return_dict:
+            return (decoded,)
+
+        return DecoderOutput(sample=decoded)
+
+    def decode_dist(self, z: Tensor, return_dict: bool = True, generator=None):
+        z = z.cuda()
+        self.use_spatial_tiling = True
+        decoded = self.decode(z)
+        self.use_spatial_tiling = False
+        return decoded
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        return_posterior: bool = True,
+        return_dict: bool = True
+    ):
+        posterior = self.encode(sample).latent_dist
+        z = posterior.sample() if sample_posterior else posterior.mode()
+        dec = self.decode(z).sample
+        return DecoderOutput(sample=dec, posterior=posterior) if return_dict else (dec, posterior)
+
+    def random_reset_tiling(self, x: torch.Tensor):
+        if x.shape[-3] == 1:
+            self.disable_spatial_tiling()
+            self.disable_temporal_tiling()
+            return
+
+        # tiling在input_shape和sample_size上限制很多，任意的input_shape和sample_size很可能不满足条件，因此这里使用固定值
+        min_sample_size = int(1 / self.tile_overlap_factor) * self.ffactor_spatial
+        min_sample_tsize = int(1 / self.tile_overlap_factor) * self.ffactor_temporal
+        sample_size = random.choice([None, 1 * min_sample_size, 2 * min_sample_size, 3 * min_sample_size])
+        if sample_size is None:
+            self.disable_spatial_tiling()
+        else:
+            self.tile_sample_min_size = sample_size
+            self.tile_latent_min_size = sample_size // self.ffactor_spatial
+            self.enable_spatial_tiling()
+
+        sample_tsize = random.choice([None, 1 * min_sample_tsize, 2 * min_sample_tsize, 3 * min_sample_tsize])
+        if sample_tsize is None:
+            self.disable_temporal_tiling()
+        else:
+            self.tile_sample_min_tsize = sample_tsize
+            self.tile_latent_min_tsize = sample_tsize // self.ffactor_temporal
+            self.enable_temporal_tiling()
+
+def load_sharded_safetensors(model_dir):
+    """
+    手动加载分片的 safetensors 文件
+
+    Args:
+        model_dir: 包含分片文件的目录路径
+
+    Returns:
+        合并后的完整权重字典
+    """
+    # 获取所有分片文件并按编号排序
+    shard_files = []
+    for file in os.listdir(model_dir):
+        if file.endswith(".safetensors"):
+            shard_files.append(file)
+
+    # 按分片编号排序
+    shard_files.sort(key=lambda x: int(x.split("-")[1]))
+
+    print(f"找到 {len(shard_files)} 个分片文件")
+
+    # 合并所有权重
+    merged_state_dict = dict()
+
+    for shard_file in shard_files:
+        shard_path = os.path.join(model_dir, shard_file)
+        print(f"加载分片: {shard_file}")
+
+        # 使用 safetensors 加载当前分片
+        with safe_open(shard_path, framework="pt", device="cpu") as f:
+            for key in f.keys():
+                tensor = f.get_tensor(key)
+                merged_state_dict[key] = tensor
+
+    print(f"合并完成，总键数量: {len(merged_state_dict)}")
+    return merged_state_dict
+
+def load_weights(model, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    def update_state_dict(state_dict: dict[str, torch.Tensor], name, weight):
+        if name not in state_dict:
+            raise ValueError(f"Unexpected weight {name}")
+
+        model_tensor = state_dict[name]
+        if model_tensor.shape != weight.shape:
+            raise ValueError(
+                f"Shape mismatch for weight {name}: "
+                f"model tensor shape {model_tensor.shape} vs. "
+                f"loaded tensor shape {weight.shape}"
+            )
+        if isinstance(weight, torch.Tensor):
+            model_tensor.data.copy_(weight.data)
+        else:
+            raise ValueError(
+                f"Unsupported tensor type in load_weights "
+                f"for {name}: {type(weight)}"
+            )
+
+    loaded_params = set()
+    for name, load_tensor in weights.items():
+        updated = True
+        name = name.replace('vae.', '')
+        if name in model.state_dict():
+            update_state_dict(model.state_dict(), name, load_tensor)
+        else:
+            updated = False
+
+        if updated:
+            loaded_params.add(name)
+
+    return loaded_params
+
+def _worker(path, config, 
+    rank=None, world_size=None, port=None, req_queue=None, rsp_queue=None):
+    """
+    each rank's worker:
+      - idle: block on req_queue.get() (CPU blocking, no GPU)
+      - receive request: run runner.predict(), all ranks forward
+      - only rank0 put result to rsp_queue
+    """
+    # _tame_cpu_threads_and_comm()
+    # basic env
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+    os.environ["WORLD_SIZE"] = str(world_size)
+    os.environ["RANK"] = str(rank)
+    os.environ["LOCAL_RANK"] = str(rank)
+
+    # device binding should be early than all CUDA operations
+    visible = torch.cuda.device_count()
+    assert visible >= world_size, f"可见卡数 {visible} < world_size {world_size}"
+    local_rank = int(os.environ["LOCAL_RANK"])
+    
+    print(f"[worker {rank}] bind to cuda:{local_rank} (visible={visible})", flush=True)
+    if not torch.distributed.is_initialized():
+        dist.init_process_group("nccl")
+    torch.cuda.set_device(local_rank)
+    #from .. import load_vae
+
+    #vae = load_vae(vae_type, vae_precision, device, logger, args, weights_only, only_encoder, only_decoder, sample_size, skip_create_dist=True)
+    #vae = vae.cuda()
+    vae = AutoencoderKLConv3D.from_config(config)
+    merged_state_dict = load_sharded_safetensors(path)
+    loaded_params = load_weights(vae, merged_state_dict) 
+    vae = vae.cuda()
+    vae.eval()  # 关闭 Dropout、BatchNorm 训练行为
+    for param in vae.parameters():
+        param.requires_grad = False  #
+    
+    while True:
+        req = req_queue.get()  # blocking
+        if req == "__STOP__":
+            break
+        out = vae.decode_dist(req, return_dict=False)
+        if rank == 0:
+            rsp_queue.put(out)
+
+    #try:
+    #    while True:
+    #        # blocking on CPU queue
+    #        req = req_queue.get()  # blocking
+    #        if req == "__STOP__":
+    #            break
+    #        out = vae.decode_dist(req, return_dict=False)
+    #        if rank == 0:
+    #            rsp_queue.put(out)
+    #finally:
+    #    # destroy process group before exit
+    #    try:
+    #        dist.destroy_process_group()
+    #    except Exception:
+    #        pass
+
+#def _find_free_port():
+#    import socket
+#    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+#        s.bind(("127.0.0.1", 0))
+#        return s.getsockname()[1]
+
+# 避免端口冲突的常见做法
+def _find_free_port(start_port=8100, max_attempts=900):
+    import socket
+    """获取一个可用的端口"""
+    for port in range(start_port, start_port + max_attempts):
+        try:
+            with socket.socket() as s:
+                s.bind(('localhost', port))
+                return s.getsockname()[1]  # 返回实际绑定的端口
+        except OSError:
+            continue
+    raise RuntimeError("找不到可用端口")
+
+class AutoencoderKLConv3D_Dist(AutoencoderKLConv3D):
+    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,
+        shift_factor: Optional[float] = None,
+        downsample_match_channel: bool = True,
+        upsample_match_channel: bool = True,
+        only_encoder: bool = False,
+        only_decoder: bool = False,
+    ):
+        super().__init__(in_channels, out_channels, latent_channels, block_out_channels, layers_per_block, ffactor_spatial, ffactor_temporal, sample_size, sample_tsize, scaling_factor, shift_factor, downsample_match_channel, upsample_match_channel, only_encoder, only_decoder)
+
+    def create_dist(self, path, config, 
+    ):
+        self.world_size = 8
+        self.port = _find_free_port()
+        ctx = mp.get_context("spawn")
+        # 每个 rank 一个请求队列（纯 CPU），再加一个公共响应队列
+        self.req_queues = [ctx.Queue() for _ in range(self.world_size)]
+        self.rsp_queue = ctx.Queue()
+
+        self.procs = []
+        for rank in range(self.world_size):
+            p = ctx.Process(
+                target=_worker,
+                args=(
+                    path, config, 
+                    rank, self.world_size, self.port,
+                    self.req_queues[rank], self.rsp_queue,
+                ),
+                daemon=True,
+            )
+            p.start()
+            self.procs.append(p)
+    
+    def decode(self, z: Tensor, return_dict: bool = True, generator=None):
+        """
+        synchronous inference: put the same request to all ranks' queues.
+        return rank0's result.
+        """
+        # check alive
+        for p in self.procs:
+            if not p.is_alive():
+                raise RuntimeError("One of the processes is not alive")
+
+        # put to each rank's queue
+        for q in self.req_queues:
+            q.put(z)
+
+        # wait for rank0's result
+        return self.rsp_queue.get(timeout=None)

cache_utils.py

@@ -0,0 +1,226 @@
+import torch
+import torch.nn as nn
+import math
+from typing import Tuple
+
+def cache_init(cache_interval, max_order, num_steps=None,
+               enable_first_enhance=False, first_enhance_steps=3, 
+               enable_tailing_enhance=False, tailing_enhance_steps=1, 
+               low_freqs_order=0, high_freqs_order=2):
+    cache_dic = {}
+    cache_dic['counter']= 0
+    cache_dic['current_step'] = 0
+    cache_dic['cache_interval']= cache_interval
+    cache_dic['max_order'] = max_order
+    cache_dic['num_steps'] = num_steps
+
+    # enhance related utils
+    
+    # first enhance: fully compute first some steps, enhancing contour infos
+    cache_dic['enable_first_enhance'] = enable_first_enhance
+    cache_dic['first_enhance_steps'] = first_enhance_steps
+
+    # tailing enhance: fully compute the last 1 steps, enhancing details
+    cache_dic['enable_tailing_enhance'] = enable_tailing_enhance
+    cache_dic['tailing_enhance_steps'] = tailing_enhance_steps
+
+    # freqs related utils
+    cache_dic['low_freqs_order'] = low_freqs_order
+    cache_dic['high_freqs_order'] = high_freqs_order
+
+    # features for training-aware cache, here we don't use these
+    cache_dic['enable_force_control']= False 
+    cache_dic['force_compute']=False
+    return cache_dic
+
+class TaylorCacheContainer(nn.Module):
+    def __init__(self, max_order):
+        super().__init__()
+        self.max_order = max_order
+        # 逐个注册buffer
+        for i in range(max_order + 1):
+            self.register_buffer(f"derivative_{i}", None, persistent=False)
+            self.register_buffer(f"temp_derivative_{i}", None, persistent=False)
+    
+    def get_derivative(self, order):
+        return getattr(self, f"derivative_{order}")
+    
+    def set_derivative(self, order, tensor):
+        setattr(self, f"derivative_{order}", tensor)
+
+    def set_temp_derivative(self, order, tensor):
+        setattr(self, f"temp_derivative_{order}", tensor)
+
+    def get_temp_derivative(self, order):
+        return getattr(self, f"temp_derivative_{order}")
+    
+    def clear_temp_derivative(self):
+        for i in range(self.max_order + 1):
+            setattr(self, f"temp_derivative_{i}", None)
+
+    def move_temp_to_derivative(self):
+        for i in range(self.max_order + 1):
+            if self.get_temp_derivative(i) is not None:
+                setattr(self, f"derivative_{i}", self.get_temp_derivative(i))
+            else:
+                break
+        self.clear_temp_derivative()
+
+    def get_all_derivatives(self):
+        return [getattr(self, f"derivative_{i}") for i in range(self.max_order + 1)]
+
+    def get_all_filled_derivatives(self):
+        return [self.get_derivative(i) for i in range(self.max_order + 1) if self.get_derivative(i) is not None]
+
+    def taylor_formula(self, distance):
+        output = 0
+        for i in range(len(self.get_all_filled_derivatives())):
+            output += (1 / math.factorial(i)) * self.get_derivative(i) * (distance ** i)
+        return output
+    
+    def derivatives_computation(self, x, distance):
+        '''
+        x: tensor, the new x_0
+        distance: int, the distance between the current step and the last full computation step
+        '''
+        self.set_temp_derivative(0, x)
+        for i in range(self.max_order):
+            if self.get_derivative(i) is not None:
+                self.set_temp_derivative(i+1, (self.get_temp_derivative(i) - self.get_derivative(i)) / distance)
+            else:
+                break
+        self.move_temp_to_derivative()
+
+    def clear_derivatives(self):
+        for i in range(self.max_order + 1):
+            setattr(self, f"derivative_{i}", None)
+            setattr(self, f"temp_derivative_{i}", None)
+
+
+@torch.compile
+def decomposition_FFT(x: torch.Tensor, cutoff_ratio: float = 0.1) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Fast Fourier Transform frequency domain decomposition
+    
+    Args:
+        x: Input tensor [B, H*W, D]
+        cutoff_ratio: Cutoff frequency ratio (0~0.5)
+        
+    Returns:
+        Tuple of (low_freq, high_freq) tensors with same dtype as input
+    """
+    orig_dtype = x.dtype
+    device = x.device
+
+    x_fp32 = x.to(torch.float32)  # Convert to fp32 for FFT compatibility
+
+    B, HW, D = x_fp32.shape
+    freq = torch.fft.fft(x_fp32, dim=1)  # FFT on spatial dimension
+
+    freqs = torch.fft.fftfreq(HW, d=1.0, device=device)
+    cutoff = cutoff_ratio * freqs.abs().max()
+
+    # Create frequency masks
+    low_mask = freqs.abs() <= cutoff
+    high_mask = ~low_mask
+
+    low_mask = low_mask[None, :, None]  # Broadcast to (B, HW, D)
+    high_mask = high_mask[None, :, None]
+
+    low_freq_complex  = freq * low_mask
+    high_freq_complex = freq * high_mask
+
+    # IFFT and take real part
+    low_fp32  = torch.fft.ifft(low_freq_complex,  dim=1).real
+    high_fp32 = torch.fft.ifft(high_freq_complex, dim=1).real
+
+    low  = low_fp32.to(device=device, dtype=orig_dtype)
+    high = high_fp32.to(device=device, dtype=orig_dtype)
+
+    return low, high
+
+@torch.compile
+def reconstruction(low_freq: torch.Tensor, high_freq: torch.Tensor) -> torch.Tensor:
+    return low_freq + high_freq
+
+class CacheWithFreqsContainer(nn.Module):
+    def __init__(self, max_order):
+        super().__init__()
+        self.max_order = max_order
+        # 逐个注册buffer
+        for i in range(max_order + 1):
+            self.register_buffer(f"derivative_{i}_low_freqs", None, persistent=False)
+            self.register_buffer(f"derivative_{i}_high_freqs", None, persistent=False)
+            self.register_buffer(f"temp_derivative_{i}_low_freqs", None, persistent=False)
+            self.register_buffer(f"temp_derivative_{i}_high_freqs", None, persistent=False)
+    
+    def get_derivative(self, order, freqs):
+        return getattr(self, f"derivative_{order}_{freqs}")
+    
+    def set_derivative(self, order, freqs, tensor):
+        setattr(self, f"derivative_{order}_{freqs}", tensor)
+
+    def set_temp_derivative(self, order, freqs, tensor):
+        setattr(self, f"temp_derivative_{order}_{freqs}", tensor)
+
+    def get_temp_derivative(self, order, freqs):
+        return getattr(self, f"temp_derivative_{order}_{freqs}")
+    
+    def move_temp_to_derivative(self):
+        for i in range(self.max_order + 1):
+            if self.get_temp_derivative(i, "low_freqs") is not None:
+                setattr(self, f"derivative_{i}_low_freqs", self.get_temp_derivative(i, "low_freqs"))
+            if self.get_temp_derivative(i, "high_freqs") is not None:
+                setattr(self, f"derivative_{i}_high_freqs", self.get_temp_derivative(i, "high_freqs"))
+            else:
+                break
+        self.clear_temp_derivative()
+
+    def get_all_filled_derivatives(self, freqs):
+        return [self.get_derivative(i, freqs) for i in range(self.max_order + 1) if self.get_derivative(i, freqs) is not None]
+
+    def taylor_formula(self, distance):
+        low_freqs_output = 0
+        high_freqs_output = 0
+        for i in range(len(self.get_all_filled_derivatives("low_freqs"))):
+            low_freqs_output += (1 / math.factorial(i)) * self.get_derivative(i, "low_freqs") * (distance ** i)
+        for i in range(len(self.get_all_filled_derivatives("high_freqs"))):
+            high_freqs_output += (1 / math.factorial(i)) * self.get_derivative(i, "high_freqs") * (distance ** i)
+        return reconstruction(low_freqs_output, high_freqs_output)
+    
+    def hermite_formula(self, distance):
+        low_freqs_output = 0
+        high_freqs_output = 0
+        for i in range(len(self.get_all_filled_derivatives("low_freqs"))):
+            low_freqs_output += (1 / math.factorial(i)) * self.get_derivative(i, "low_freqs") * (distance ** i)
+        for i in range(len(self.get_all_filled_derivatives("high_freqs"))):
+            high_freqs_output += (1 / math.factorial(i)) * self.get_derivative(i, "high_freqs") * (distance ** i)
+        return reconstruction(low_freqs_output, high_freqs_output)
+
+    def derivatives_computation(self, x, distance, low_freqs_order, high_freqs_order):
+        '''
+        x: tensor, the new x_0
+        distance: int, the distance between the current step and the last full computation step
+        '''
+        x_low, x_high = decomposition_FFT(x, cutoff_ratio=0.1)
+        self.set_temp_derivative(0, "low_freqs", x_low)
+        self.set_temp_derivative(0, "high_freqs", x_high)
+        for i in range(low_freqs_order):
+            if self.get_derivative(i, "low_freqs") is not None:
+                self.set_temp_derivative(i+1, "low_freqs", (self.get_temp_derivative(i, "low_freqs") - self.get_derivative(i, "low_freqs")) / distance)
+        for i in range(high_freqs_order):
+            if self.get_derivative(i, "high_freqs") is not None:
+                self.set_temp_derivative(i+1, "high_freqs", (self.get_temp_derivative(i, "high_freqs") - self.get_derivative(i, "high_freqs")) / distance)
+        self.move_temp_to_derivative()
+        
+    def clear_temp_derivative(self):
+        for i in range(self.max_order + 1):
+            setattr(self, f"temp_derivative_{i}_low_freqs", None)
+            setattr(self, f"temp_derivative_{i}_high_freqs", None)
+
+    def clear_derivatives(self):
+        for i in range(self.max_order + 1):
+            setattr(self, f"derivative_{i}_low_freqs", None)
+            setattr(self, f"derivative_{i}_high_freqs", None)
+            setattr(self, f"temp_derivative_{i}_low_freqs", None)
+            setattr(self, f"temp_derivative_{i}_high_freqs", None)

config.json

@@ -0,0 +1,285 @@
+{
+    "cfg_distilled": true,
+    "use_meanflow": true,
+    "add_classification_head": false,
+    "anyres_pooling_size": 2,
+    "anyres_vit_max_image_size": null,
+    "anyres_vit_two_views": false,
+    "architectures": [
+        "HunyuanImage3ForCausalMM"
+    ],
+    "auto_map": {
+        "AutoConfig": "configuration_hunyuan_image_3.HunyuanImage3Config",
+        "AutoModel": "modeling_hunyuan_image_3.HunyuanImage3Model",
+        "AutoModelForCausalLM": "modeling_hunyuan_image_3.HunyuanImage3ForCausalMM"
+    },
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attention_head_dim": 128,
+    "bos_token_id": 127958,
+    "cla_share_factor": 2,
+    "class_num": 0,
+    "dense_list": [
+        4096,
+        0
+    ],
+    "eod_token_id": 3,
+    "eos_token_id": 127957,
+    "group_limited_greedy": false,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "im_end_id": 128001,
+    "im_newline_id": 11,
+    "im_start_id": 128000,
+    "image_token_id": 128006,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "kv_lora_rank": null,
+    "mask_init_id": 12,
+    "max_position_embeddings": 22800,
+    "mlp_bias": false,
+    "model_type": "hunyuan_image_3_moe",
+    "moe_drop_tokens": false,
+    "moe_intermediate_size": [
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072,
+        3072
+    ],
+    "moe_layer_num_skipped": 0,
+    "moe_random_routing_dropped_token": false,
+    "moe_topk": [
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8,
+        8
+    ],
+    "n_group": false,
+    "norm_topk_prob": true,
+    "norm_type": "rms",
+    "num_attention_heads": 32,
+    "num_experts": 64,
+    "num_hidden_layers": 32,
+    "num_key_value_heads": 8,
+    "num_media_embeds": 257,
+    "num_shared_expert": [
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1
+    ],
+    "pad_id": 128009,
+    "pad_token_id": 128009,
+    "pool_type": "last",
+    "position_embedding_xdrope": false,
+    "pretraining_tp": 1,
+    "q_lora_rank": null,
+    "qk_nope_head_dim": null,
+    "qk_rope_head_dim": null,
+    "rms_norm_eps": 1e-05,
+    "rope_scaling": {
+        "alpha": 1.0,
+        "beta_fast": 32,
+        "beta_slow": 1,
+        "factor": 1.0,
+        "mscale": 1.0,
+        "mscale_all_dim": 1.0,
+        "type": "custom"
+    },
+    "rope_theta": 10000.0,
+    "routed_scaling_factor": false,
+    "skip_cls_token": false,
+    "text_end_id": 7,
+    "text_start_id": 6,
+    "tie_word_embeddings": false,
+    "topk_group": false,
+    "torch_dtype": "bfloat16",
+    "transformers_version": "4.50.0",
+    "use_cache": true,
+    "use_cla": false,
+    "use_mixed_mlp_moe": true,
+    "use_mla": false,
+    "use_qk_norm": true,
+    "use_rotary_pos_emb": true,
+    "v_head_dim": null,
+    "video_end_id": 10,
+    "video_start_id": 9,
+    "vit_add_patchemb_bias": false,
+    "vit_input_resolution": 224,
+    "vit_mapping_type": "resampler",
+    "vit_norm_type": "fused",
+    "vit_patch": 1,
+    "vit_path": null,
+    "vit_remove_prenorm": false,
+    "vit_token": 64,
+    "vit_type": "siglip2-so400m-patch16-naflex",
+    "vit_used_rms_norm": false,
+    "vocab_size": 133120,
+    "xdrope_section": null,
+    "head_dim": 128,
+    "rope_type": "2d",
+    "vae_downsample_factor": [
+        16,
+        16
+    ],
+    "vit_downsample_factor": [
+        16,
+        16
+    ],
+    "cond_token_attn_type": "joint_full",
+    "cond_image_type": "vae_vit",
+    "vae_type": "hunyuan-image-vae-v1",
+    "vae_dtype": "float32",
+    "vae_autocast_dtype": "float16",
+    "vae": {
+        "_class_name": "AutoencoderKLConv3D",
+        "block_out_channels": [
+            128,
+            256,
+            512,
+            1024,
+            1024
+        ],
+        "in_channels": 3,
+        "out_channels": 3,
+        "latent_channels": 32,
+        "layers_per_block": 2,
+        "ffactor_spatial": 16,
+        "ffactor_temporal": 4,
+        "sample_size": 384,
+        "sample_tsize": 96,
+        "downsample_match_channel": true,
+        "upsample_match_channel": true,
+        "scaling_factor": 0.562679178327931
+    },
+    "vit": {
+        "_attn_implementation": "sdpa",
+        "attention_dropout": 0.0,
+        "hidden_act": "gelu_pytorch_tanh",
+        "hidden_size": 1152,
+        "intermediate_size": 4304,
+        "layer_norm_eps": 1e-06,
+        "num_attention_heads": 16,
+        "num_channels": 3,
+        "num_hidden_layers": 27,
+        "num_patches": 256,
+        "patch_size": 16,
+        "torch_dtype": "float32",
+        "output_attentions": false,
+        "output_hidden_states": false,
+        "use_return_dict": true
+    },
+    "vit_processor": {
+        "do_convert_rgb": null,
+        "do_normalize": true,
+        "do_rescale": true,
+        "do_resize": true,
+        "image_mean": [
+            0.5,
+            0.5,
+            0.5
+        ],
+        "image_processor_type": "Siglip2ImageProcessorFast",
+        "image_std": [
+            0.5,
+            0.5,
+            0.5
+        ],
+        "max_num_patches": 1024,
+        "patch_size": 16,
+        "processor_class": "Siglip2Processor",
+        "resample": 2,
+        "rescale_factor": 0.00392156862745098
+    },
+    "vit_aligner": {
+        "projector_type": "mlp_gelu",
+        "input_dim": 1152,
+        "n_embed": 4096,
+        "depth": 2,
+        "torch_dtype": "float32"
+    }
+}

configuration_hunyuan_image_3.py

@@ -0,0 +1,310 @@
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/LICENSE
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+from typing import List, Union, Optional
+
+
+logger = logging.get_logger(__name__)
+
+
+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: int = 290943,
+            hidden_size: int = 4096,
+            intermediate_size: int = 11008,
+            moe_intermediate_size: Union[int, List] = None,
+            num_hidden_layers: int = 32,
+            num_attention_heads: int = 32,
+            num_key_value_heads: Optional[int] = None,
+            attention_head_dim: Optional[int] = 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: Union[int, List] = 1,
+            use_mixed_mlp_moe=False,
+            num_shared_expert: Union[int, List] = 1,
+            moe_topk: Union[int, List] = 1,
+            capacity_factor: int = 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,
+            rope_type="2d",
+            cond_token_attn_type="full",
+            cond_image_type="vae_vit",
+            vae_type=None,
+            vae_dtype="float32",
+            vae_autocast_dtype="float16",
+            vae=None,
+            vit_type=None,
+            vit=None,
+            vit_processor=None,
+            vit_aligner=None,
+            cfg_distilled=False,
+            use_meanflow=False,
+            **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]
+
+        # Conditioning image configs
+        self.cond_token_attn_type = cond_token_attn_type
+        self.cond_image_type = cond_image_type
+
+        # ViT args
+        self.vit_type = vit_type
+        self.vit = vit
+        self.vit_processor = vit_processor
+        self.vit_aligner = vit_aligner
+
+        # Image Gen args
+        self.vae_type = vae_type
+        self.vae_dtype = vae_dtype
+        self.vae_autocast_dtype = vae_autocast_dtype
+        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
+        self.rope_type = rope_type
+
+        # 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
+
+        # flag of cfg distilled model
+        self.cfg_distilled = cfg_distilled
+        # flag of meanflow distilled model
+        self.use_meanflow = use_meanflow
+        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,
+        )
+
+
+__all__ = ["HunyuanImage3Config"]

generation_config.json

@@ -0,0 +1,21 @@
+{
+    "disable_compile": true,
+    "eos_token_id": [
+        127957
+    ],
+    "pad_token_id": 128009,
+    "do_sample": true,
+    "top_k": 1024,
+    "top_p": 0.95,
+    "temperature": 0.6,
+    "max_length": 22800,
+    "sequence_template": "instruct",
+    "diff_infer_steps": 8,
+    "diff_guidance_scale": 2.5,
+    "flow_shift": 3.0,
+    "use_system_prompt": "en_unified",
+    "drop_think": false,
+    "bot_task": "think_recaption",
+    "max_new_tokens": 2048,
+    "transformers_version": "4.50.0"
+}

hunyuan_image_3_pipeline.py

@@ -0,0 +1,913 @@
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/LICENSE
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================================
+
+import inspect
+import math
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+from PIL import Image
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput, logging
+from diffusers.utils.torch_utils import randn_tensor
+from .cache_utils import cache_init
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    """
+    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 rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    r"""
+    Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on
+    Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
+    Flawed](https://arxiv.org/pdf/2305.08891.pdf).
+
+    Args:
+        noise_cfg (`torch.Tensor`):
+            The predicted noise tensor for the guided diffusion process.
+        noise_pred_text (`torch.Tensor`):
+            The predicted noise tensor for the text-guided diffusion process.
+        guidance_rescale (`float`, *optional*, defaults to 0.0):
+            A rescale factor applied to the noise predictions.
+    Returns:
+        noise_cfg (`torch.Tensor`): The rescaled noise prediction tensor.
+    """
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    # rescale the results from guidance (fixes overexposure)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
+
+@dataclass
+class HunyuanImage3Text2ImagePipelineOutput(BaseOutput):
+    samples: Union[List[Any], np.ndarray]
+
+
+@dataclass
+class FlowMatchDiscreteSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+    """
+
+    prev_sample: torch.FloatTensor
+
+
+class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
+    """
+    Euler scheduler.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        timestep_spacing (`str`, defaults to `"linspace"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        shift (`float`, defaults to 1.0):
+            The shift value for the timestep schedule.
+        reverse (`bool`, defaults to `True`):
+            Whether to reverse the timestep schedule.
+    """
+
+    _compatibles = []
+    order = 1
+
+    @register_to_config
+    def __init__(
+            self,
+            num_train_timesteps: int = 1000,
+            shift: float = 1.0,
+            reverse: bool = True,
+            solver: str = "euler",
+            use_flux_shift: bool = False,
+            flux_base_shift: float = 0.5,
+            flux_max_shift: float = 1.15,
+            n_tokens: Optional[int] = None,
+    ):
+        sigmas = torch.linspace(1, 0, num_train_timesteps + 1)
+
+        if not reverse:
+            sigmas = sigmas.flip(0)
+
+        self.sigmas = sigmas
+        # the value fed to model
+        self.timesteps = (sigmas[:-1] * num_train_timesteps).to(dtype=torch.float32)
+        self.timesteps_full = (sigmas * num_train_timesteps).to(dtype=torch.float32)
+
+        self._step_index = None
+        self._begin_index = None
+
+        self.supported_solver = [
+            "euler",
+            "heun-2", "midpoint-2",
+            "kutta-4",
+        ]
+        if solver not in self.supported_solver:
+            raise ValueError(f"Solver {solver} not supported. Supported solvers: {self.supported_solver}")
+
+        # empty dt and derivative (for heun)
+        self.derivative_1 = None
+        self.derivative_2 = None
+        self.derivative_3 = None
+        self.dt = None
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+
+    @property
+    def state_in_first_order(self):
+        return self.derivative_1 is None
+
+    @property
+    def state_in_second_order(self):
+        return self.derivative_2 is None
+
+    @property
+    def state_in_third_order(self):
+        return self.derivative_3 is None
+
+    def get_timestep_r(self, timestep: Union[float, torch.FloatTensor]):
+        if self.step_index is None:
+            self._init_step_index(timestep)
+        return self.timesteps_full[self.step_index + 1]
+
+    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None,
+                      n_tokens: int = None):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            n_tokens (`int`, *optional*):
+                Number of tokens in the input sequence.
+        """
+        self.num_inference_steps = num_inference_steps
+
+        sigmas = torch.linspace(1, 0, num_inference_steps + 1)
+
+        # Apply timestep shift
+        if self.config.use_flux_shift:
+            assert isinstance(n_tokens, int), "n_tokens should be provided for flux shift"
+            mu = self.get_lin_function(y1=self.config.flux_base_shift, y2=self.config.flux_max_shift)(n_tokens)
+            sigmas = self.flux_time_shift(mu, 1.0, sigmas)
+        elif self.config.shift != 1.:
+            sigmas = self.sd3_time_shift(sigmas)
+
+        if not self.config.reverse:
+            sigmas = 1 - sigmas
+
+        self.sigmas = sigmas
+        self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(dtype=torch.float32, device=device)
+        self.timesteps_full = (sigmas * self.config.num_train_timesteps).to(dtype=torch.float32, device=device)
+
+        # empty dt and derivative (for kutta)
+        self.derivative_1 = None
+        self.derivative_2 = None
+        self.derivative_3 = None
+        self.dt = None
+
+        # Reset step index
+        self._step_index = None
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep):
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor:
+        return sample
+
+    @staticmethod
+    def get_lin_function(x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15):
+        m = (y2 - y1) / (x2 - x1)
+        b = y1 - m * x1
+        return lambda x: m * x + b
+
+    @staticmethod
+    def flux_time_shift(mu: float, sigma: float, t: torch.Tensor):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def sd3_time_shift(self, t: torch.Tensor):
+        return (self.config.shift * t) / (1 + (self.config.shift - 1) * t)
+
+    def step(
+            self,
+            model_output: torch.FloatTensor,
+            timestep: Union[float, torch.FloatTensor],
+            sample: torch.FloatTensor,
+            pred_uncond: torch.FloatTensor = None,
+            generator: Optional[torch.Generator] = None,
+            n_tokens: Optional[int] = None,
+            return_dict: bool = True,
+    ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            n_tokens (`int`, *optional*):
+                Number of tokens in the input sequence.
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
+                tuple.
+
+        Returns:
+            [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
+                returned, otherwise a tuple is returned where the first element is the sample tensor.
+        """
+
+        if (
+                isinstance(timestep, int)
+                or isinstance(timestep, torch.IntTensor)
+                or isinstance(timestep, torch.LongTensor)
+        ):
+            raise ValueError(
+                (
+                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                    " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                    " one of the `scheduler.timesteps` as a timestep."
+                ),
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+        model_output = model_output.to(torch.float32)
+        pred_uncond = pred_uncond.to(torch.float32) if pred_uncond is not None else None
+
+        # dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index]
+        sigma = self.sigmas[self.step_index]
+        sigma_next = self.sigmas[self.step_index + 1]
+
+        last_inner_step = True
+        if self.config.solver == "euler":
+            derivative, dt, sample, last_inner_step = self.first_order_method(model_output, sigma, sigma_next, sample)
+        elif self.config.solver in ["heun-2", "midpoint-2"]:
+            derivative, dt, sample, last_inner_step = self.second_order_method(model_output, sigma, sigma_next, sample)
+        elif self.config.solver == "kutta-4":
+            derivative, dt, sample, last_inner_step = self.fourth_order_method(model_output, sigma, sigma_next, sample)
+        else:
+            raise ValueError(f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}")
+
+        prev_sample = sample + derivative * dt
+
+        # Cast sample back to model compatible dtype
+        # prev_sample = prev_sample.to(model_output.dtype)
+
+        # upon completion increase step index by one
+        if last_inner_step:
+            self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample)
+
+    def first_order_method(self, model_output, sigma, sigma_next, sample):
+        derivative = model_output
+        dt = sigma_next - sigma
+        return derivative, dt, sample, True
+
+    def second_order_method(self, model_output, sigma, sigma_next, sample):
+        if self.state_in_first_order:
+            # store for 2nd order step
+            self.derivative_1 = model_output
+            self.dt = sigma_next - sigma
+            self.sample = sample
+
+            derivative = model_output
+            if self.config.solver == 'heun-2':
+                dt = self.dt
+            elif self.config.solver == 'midpoint-2':
+                dt = self.dt / 2
+            else:
+                raise NotImplementedError(f"Solver {self.config.solver} not supported.")
+            last_inner_step = False
+
+        else:
+            if self.config.solver == 'heun-2':
+                derivative = 0.5 * (self.derivative_1 + model_output)
+            elif self.config.solver == 'midpoint-2':
+                derivative = model_output
+            else:
+                raise NotImplementedError(f"Solver {self.config.solver} not supported.")
+
+            # 3. take prev timestep & sample
+            dt = self.dt
+            sample = self.sample
+            last_inner_step = True
+
+            # free dt and derivative
+            # Note, this puts the scheduler in "first order mode"
+            self.derivative_1 = None
+            self.dt = None
+            self.sample = None
+
+        return derivative, dt, sample, last_inner_step
+
+    def fourth_order_method(self, model_output, sigma, sigma_next, sample):
+        if self.state_in_first_order:
+            self.derivative_1 = model_output
+            self.dt = sigma_next - sigma
+            self.sample = sample
+            derivative = model_output
+            dt = self.dt / 2
+            last_inner_step = False
+
+        elif self.state_in_second_order:
+            self.derivative_2 = model_output
+            derivative = model_output
+            dt = self.dt / 2
+            last_inner_step = False
+
+        elif self.state_in_third_order:
+            self.derivative_3 = model_output
+            derivative = model_output
+            dt = self.dt
+            last_inner_step = False
+
+        else:
+            derivative = (1/6 * self.derivative_1 + 1/3 * self.derivative_2 + 1/3 * self.derivative_3 +
+                          1/6 * model_output)
+
+            # 3. take prev timestep & sample
+            dt = self.dt
+            sample = self.sample
+            last_inner_step = True
+
+            # free dt and derivative
+            # Note, this puts the scheduler in "first order mode"
+            self.derivative_1 = None
+            self.derivative_2 = None
+            self.derivative_3 = None
+            self.dt = None
+            self.sample = None
+
+        return derivative, dt, sample, last_inner_step
+
+    def __len__(self):
+        return self.config.num_train_timesteps
+
+
+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: Optional[torch.Tensor],
+            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
+
+
+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,
+    ):
+        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: Union[np.ndarray, torch.Tensor]) -> Union[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) shoudl 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,
+        sigmas: List[float] = None,
+        guidance_scale: float = 7.5,
+        meanflow: bool = False,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        guidance_rescale: float = 0.0,
+        callback_on_step_end: Optional[
+            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+        ] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        model_kwargs: Dict[str, Any] = 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
+
+
+        if not kwargs.get('cfg_distilled', False):
+            cfg_factor = 1 + self.do_classifier_free_guidance
+        else:
+            cfg_factor = 1
+        # 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,
+        )
+        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)
+
+        # Taylor cache 
+        cache_dic = None
+        if self.model.use_taylor_cache:
+            cache_dic = cache_init(cache_interval=self.model.taylor_cache_interval, max_order=self.model.taylor_cache_order, num_steps=len(timesteps),
+                                enable_first_enhance=self.model.taylor_cache_enable_first_enhance, first_enhance_steps=self.model.taylor_cache_first_enhance_steps, 
+                                enable_tailing_enhance=self.model.taylor_cache_enable_tailing_enhance,
+                                tailing_enhance_steps=self.model.taylor_cache_tailing_enhance_steps,
+                                low_freqs_order=self.model.taylor_cache_low_freqs_order,
+                                high_freqs_order=self.model.taylor_cache_high_freqs_order)
+        print(f"***use_taylor_cache: {self.model.use_taylor_cache}, cache_dic: {cache_dic}")
+
+        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)
+
+                if meanflow:
+                    r = self.scheduler.get_timestep_r(t)
+                    r_expand = r.repeat(latent_model_input.shape[0])
+                else:
+                    r_expand = None
+                model_kwargs["timesteps_r"] = r_expand
+
+                t_expand = t.repeat(latent_model_input.shape[0])
+
+                if self.model.use_taylor_cache:
+                    cache_dic['current_step'] = i
+                    model_kwargs['cache_dic'] = cache_dic
+                if kwargs.get('cfg_distilled', False):
+                    model_kwargs["guidance"] = torch.tensor(
+                        [1000.0*self._guidance_scale], device=self.device, dtype=torch.bfloat16
+                    )
+                model_inputs = self.model.prepare_inputs_for_generation(
+                    input_ids,
+                    images=latent_model_input,
+                    timesteps=t_expand,
+                    **model_kwargs,
+                )
+                with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+                    model_output = self.model(**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:
+                    if not kwargs.get('cfg_distilled', False):
+                        pred_cond, pred_uncond = pred.chunk(2)
+                        pred = self.cfg_operator(pred_cond, pred_uncond, self.guidance_scale, step=i)
+
+                if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
+                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                    pred = rescale_noise_cfg(pred, pred_cond, guidance_rescale=self.guidance_rescale)
+
+                # 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,
+                    )
+                    input_ids = None
+                    # if input_ids.shape[1] != model_kwargs["position_ids"].shape[1]:
+                    #     input_ids = torch.gather(input_ids, 1, index=model_kwargs["position_ids"])
+
+                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]
+
+        # b c t h w
+        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)

image_processor.py

@@ -0,0 +1,465 @@
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/LICENSE
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from dataclasses import dataclass, field, asdict
+from typing import Tuple, Optional, Callable, Union, Any
+import random
+import math
+
+import torch
+from PIL import Image
+from torchvision import transforms
+from transformers.image_processing_utils import BaseImageProcessor
+from transformers.image_utils import load_image
+from transformers.models.siglip2.image_processing_siglip2_fast import Siglip2ImageProcessorFast
+from transformers.generation.logits_process import LogitsProcessor, LogitsProcessorList
+
+from .tokenization_hunyuan_image_3 import ImageInfo, ImageTensor, CondImage, Resolution, ResolutionGroup
+
+InputImage = Union[Image.Image, str]
+
+
+class SliceVocabLogitsProcessor(LogitsProcessor):
+    """
+    [`LogitsProcessor`] that performs vocab slicing, i.e. restricting probabilities with in some range. This processor
+    is often used in multimodal discrete LLMs, which ensure that we only sample within one modality
+
+    Args:
+        vocab_start (`int`): start of slice, default None meaning from 0
+        vocab_end (`int`): end of slice, default None meaning to the end of list
+        when start and end are all None, this processor does noting
+
+    """
+
+    def __init__(self, vocab_start: int = None, vocab_end: int = None, **kwargs):
+        if vocab_start is not None and vocab_end is not None:
+            assert vocab_start < vocab_end, f"Ensure vocab_start {vocab_start} < vocab_end {vocab_end}"
+        self.vocab_start = vocab_start
+        self.vocab_end = vocab_end
+        self.other_slices = kwargs.get("other_slices", [])
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        scores_processed = scores[:, self.vocab_start: self.vocab_end]
+        for other_slice in self.other_slices:
+            scores_processed = torch.cat([scores_processed, scores[:, other_slice[0]: other_slice[1]]], dim=-1)
+        return scores_processed
+
+    def __repr__(self):
+        return f"SliceVocabLogitsWarper(vocab_start={self.vocab_start}, vocab_end={self.vocab_end}, other_slices={self.other_slices})"
+
+
+def resize_and_crop(image: Image.Image, target_size: Tuple[int, int], resample=Image.Resampling.LANCZOS, crop_type='center', crop_coords=None) -> Image.Image:
+    tw, th = target_size
+    w, h = image.size
+
+    tr = th / tw
+    r = h / w
+
+    if crop_type == "resize":
+        resize_width = tw
+        resize_height = th
+        crop_top = 0
+        crop_left = 0
+        image = image.resize((resize_width, resize_height), resample=resample)
+    else:
+        # maintain the aspect ratio
+        if r < tr:
+            resize_height = th
+            resize_width = int(round(th / h * w))
+        else:
+            resize_width = tw
+            resize_height = int(round(tw / w * h))
+
+        if crop_type == 'center':
+            crop_top = int(round((resize_height - th) / 2.0))
+            crop_left = int(round((resize_width - tw) / 2.0))
+        elif crop_type == 'random':
+            crop_top = random.randint(0, resize_height - th)
+            crop_left = random.randint(0, resize_width - tw)
+        elif crop_type == 'fixed':
+            assert crop_coords is not None, 'crop_coords should be provided when crop_type is fixed.'
+            crop_left, crop_top = crop_coords
+        else:
+            raise ValueError(f'crop_type must be center, random or fixed, but got {crop_type}')
+
+        image = image.resize((resize_width, resize_height), resample=resample)
+        image = image.crop((crop_left, crop_top, crop_left + tw, crop_top + th))
+
+    return image
+
+
+@dataclass
+class ResolutionGroupConfig:
+    base_size: int = None
+    step: Optional[int] = None
+    align: int = 16
+
+    def to_dict(self):
+        return asdict(self)
+
+
+@dataclass
+class VAEInfo:
+    encoder_type: str
+    down_h_factor: int = -1
+    down_w_factor: int = -1
+    patch_size: int = 1
+    h_factor: int = -1
+    w_factor: int = -1
+    image_type: str = None
+
+    def __post_init__(self):
+        self.h_factor = self.down_h_factor * self.patch_size
+        self.w_factor = self.down_w_factor * self.patch_size
+        if self.image_type is None:
+            self.image_type = "vae"
+
+
+@dataclass
+class ViTInfo:
+    encoder_type: str
+    h_factor: int = -1
+    w_factor: int = -1
+    max_token_length: int = 0   # pad to max_token_length
+    processor: Callable = field(default_factory=BaseImageProcessor)
+    image_type: str = None
+
+    def __post_init__(self):
+        if self.image_type is None:
+            self.image_type = self.encoder_type.split("-")[0]
+
+
+class HunyuanImage3ImageProcessor(object):
+    def __init__(self, config):
+        self.config = config
+
+        self.reso_group_config = ResolutionGroupConfig(base_size=config.image_base_size)
+        self.vae_reso_group = ResolutionGroup(
+            **self.reso_group_config.to_dict(),
+            extra_resolutions=[
+                Resolution("1024x768"),
+                Resolution("1280x720"),
+                Resolution("768x1024"),
+                Resolution("720x1280"),
+            ]
+        )
+        self.img_ratio_slice_logits_processor = None
+        self.pil_image_to_tensor = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5]),  # transform to [-1, 1]
+        ])
+        self.vae_info = VAEInfo(
+            encoder_type=config.vae_type,
+            down_h_factor=config.vae_downsample_factor[0], down_w_factor=config.vae_downsample_factor[0],
+            patch_size=config.patch_size,
+        )
+
+        if config.vit_type == "siglip2-so400m-patch16-naflex":
+            self.vit_processor = Siglip2ImageProcessorFast.from_dict(config.vit_processor)
+        else:
+            raise ValueError(f"Unsupported vit_type: {config.vit_type}")
+        self.vit_info = ViTInfo(
+            encoder_type=config.vit_type,
+            h_factor=self.vit_processor.patch_size,
+            w_factor=self.vit_processor.patch_size,
+            max_token_length=self.vit_processor.max_num_patches,
+            processor=self.vit_processor,
+        )
+        self.cond_token_attn_type = config.cond_token_attn_type
+        self.cond_image_type = config.cond_image_type
+
+    def build_gen_image_info(self, image_size, add_guidance_token=False, add_timestep_r_token=False) -> ImageInfo:
+        # parse image size (HxW, H:W, or <img_ratio_i>)
+        if isinstance(image_size, str):
+            if image_size.startswith("<img_ratio_"):
+                ratio_index = int(image_size.split("_")[-1].rstrip(">"))
+                reso = self.vae_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(':')]
+                assert len(image_size) == 2, f"`image_size` should be in the format of 'W:H', got {image_size}."
+                # Note that ratio is width:height
+                image_size = [image_size[1], image_size[0]]
+            else:
+                raise ValueError(
+                    f"`image_size` should be in the format of 'HxW', 'W:H' or <img_ratio_i>, 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', "
+                             f"got {image_size}.")
+        image_width, image_height = self.vae_reso_group.get_target_size(image_size[1], image_size[0])
+        token_height = image_height // self.vae_info.h_factor
+        token_width = image_width // self.vae_info.w_factor
+        base_size, ratio_idx = self.vae_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,
+            add_guidance_token=add_guidance_token, add_timestep_r_token=add_timestep_r_token,
+        )
+        return image_info
+
+    def as_image_tensor(self, image, image_type, **kwargs) -> ImageTensor:
+        if isinstance(image, Image.Image):
+            tensor = self.pil_image_to_tensor(image)
+        else:
+            tensor = image
+        
+        origin_size = kwargs["origin_size"]
+        ori_image_width = origin_size[0]
+        ori_image_height = origin_size[1]
+
+        if image_type == "vae":
+            assert tensor.ndim == 3 or tensor.ndim == 4
+            h, w = tensor.shape[-2], tensor.shape[-1]
+            assert (h % self.vae_info.h_factor == 0 and w % self.vae_info.w_factor == 0), \
+                (f"Image size should be divisible by ({self.vae_info.h_factor}, {self.vae_info.w_factor}), "
+                 f"but got ({h} x {w}).")
+            tk_height = h // self.vae_info.h_factor
+            tk_width = w // self.vae_info.w_factor
+            base_size, ratio_idx = self.vae_reso_group.get_base_size_and_ratio_index(w, h)
+            tensor.i = ImageInfo(
+                image_type=image_type,
+                image_width=w, image_height=h, token_width=tk_width, token_height=tk_height,
+                base_size=base_size, ratio_index=ratio_idx,
+                ori_image_width=ori_image_width,
+                ori_image_height=ori_image_height,
+            )
+            tensor.section_type = "cond_vae_image"
+        elif image_type == "siglip2":
+            spatial_shapes = kwargs["spatial_shapes"]  # 2  (h, w)
+            pixel_attention_mask = kwargs["pixel_attention_mask"]  # seq_len
+            tensor.i = ImageInfo(
+                image_type=image_type,
+                image_width=spatial_shapes[1].item() * self.vit_info.w_factor,
+                image_height=spatial_shapes[0].item() * self.vit_info.h_factor,
+                token_width=spatial_shapes[1].item(),
+                token_height=spatial_shapes[0].item(),
+                image_token_length=self.vit_info.max_token_length,
+                ori_image_width=ori_image_width,
+                ori_image_height=ori_image_height,
+            )
+            tensor.section_type = "cond_vit_image"
+            tensor.vision_encoder_kwargs = {
+                "spatial_shapes": spatial_shapes,
+                "pixel_attention_mask": pixel_attention_mask,
+            }
+        elif image_type == "anyres":
+            token_width = kwargs["resized_image_width"] // self.vit_info.w_factor
+            token_height = kwargs["resized_image_height"] // self.vit_info.h_factor
+            tensor.i = ImageInfo(
+                image_type=image_type,
+                image_width=kwargs["resized_image_width"],
+                image_height=kwargs["resized_image_height"],
+                token_width=token_width,
+                token_height=token_height,
+                image_token_length=token_height * (token_width + 1) + 2,
+            )
+            tensor.section_type = "cond_vit_image"
+        else:
+            raise ValueError(f"Unknown image type: {image_type}")
+        return tensor
+
+    def vae_process_image(self, image, target_size, random_crop: bool | str = False) -> ImageTensor:
+        origin_size = image.size
+        crop_type = random_crop if isinstance(random_crop, str) else ("random" if random_crop else "center")
+        resized_image = resize_and_crop(image, target_size, crop_type=crop_type)
+        return self.as_image_tensor(resized_image, image_type=self.vae_info.image_type, origin_size=origin_size)
+
+    def vit_process_image(self, image) -> ImageTensor:
+        origin_size = image.size
+        inputs = self.vit_info.processor(image)
+        image = inputs["pixel_values"].squeeze(0)   # (seq_len, dim)
+
+        remain_keys = set(inputs.keys()) - {"pixel_values"}
+        remain_kwargs = {}
+        for key in remain_keys:
+            if isinstance(inputs[key], torch.Tensor):
+                remain_kwargs[key] = inputs[key].squeeze(0)
+            else:
+                remain_kwargs[key] = inputs[key]
+
+        return self.as_image_tensor(image, image_type=self.vit_info.image_type, origin_size=origin_size, **remain_kwargs)
+
+    def get_image_with_size(
+            self,
+            src: InputImage,
+            random_crop: bool | str = False,
+            return_type: str = "vae",
+    ) -> tuple[ImageTensor | CondImage, bool]:
+        """ For various image generation tasks, dynamic image sizes """
+        image = load_image(src)
+        image_flag = "normal"
+        img_success = image_flag != "gray"
+        origin_size = image.size  # (w_ori, h_ori)
+
+        if "vae" in return_type:
+            target_size = self.vae_reso_group.get_target_size(*origin_size)
+            vae_image_tensor = self.vae_process_image(image, target_size, random_crop=random_crop)
+        else:
+            vae_image_tensor = None
+
+        if "vit" in return_type:
+            vit_image_tensor = self.vit_process_image(image)
+        else:
+            vit_image_tensor = None
+
+        if return_type == "vae":
+            image_tensor = vae_image_tensor
+        elif return_type == "vit":
+            image_tensor = vit_image_tensor
+        elif return_type == "vae_vit":
+            image_tensor = CondImage(image_type=return_type, vae_image=vae_image_tensor, vit_image=vit_image_tensor)
+        else:
+            raise ValueError(f"Unknown return_type: {return_type}")
+
+        return image_tensor, img_success
+
+    def build_cond_images(
+            self,
+            image_list: Optional[list[InputImage]] = None,
+            message_list: Optional[list[dict[str, Any]]] = None,
+            infer_align_image_size: bool = False,
+    ) -> Optional[list[CondImage]]:
+        if image_list is not None and message_list is not None:
+            raise ValueError("`image_list` and `message_list` cannot be provided at the same time.")
+        if message_list is not None:
+            image_list = []
+            for message in message_list:
+                visuals = [
+                    content
+                    for content in message["content"]
+                    if isinstance(content, dict) and content["type"] in ["image"]
+                ]
+                image_list.extend([
+                    vision_info[key]
+                    for vision_info in visuals
+                    for key in ["image", "url", "path", "base64"]
+                    if key in vision_info and vision_info["type"] == "image"
+                ])
+
+        if infer_align_image_size:
+            random_crop = "resize"
+        else:
+            random_crop = "center"
+
+        return [
+            self.get_image_with_size(src, return_type=self.cond_image_type, random_crop=random_crop)[0]
+            for src in image_list
+        ]
+
+    def prepare_full_attn_slices(self, output, batch_idx=None, with_gen=True):
+        """ Determine full attention image slices according to strategies. """
+        if self.cond_image_type == "vae":
+            cond_choices = dict(
+                causal=[],
+                full=output.vae_image_slices[batch_idx] if batch_idx is not None else output.vae_image_slices
+            )
+
+        elif self.cond_image_type == "vit":
+            cond_choices = dict(
+                causal=[],
+                full=output.vit_image_slices[batch_idx] if batch_idx is not None else output.vit_image_slices
+            )
+
+        elif self.cond_image_type == "vae_vit":
+            cond_choices = {
+                "causal": [],
+                "full": (
+                    output.vae_image_slices[batch_idx] + output.vit_image_slices[batch_idx]
+                    if batch_idx is not None
+                    else output.vae_image_slices + output.vit_image_slices
+                ),
+                "joint_full": (
+                    output.joint_image_slices[batch_idx]
+                    if batch_idx is not None
+                    else output.joint_image_slices
+                ),
+                "full_causal": (
+                    output.vae_image_slices[batch_idx]
+                    if batch_idx is not None
+                    else output.vae_image_slices
+                ),
+            }
+
+        else:
+            raise ValueError(f"Unknown cond_image_type: {self.cond_image_type}")
+        slices = cond_choices[self.cond_token_attn_type]
+
+        if with_gen:
+            gen_image_slices = (
+                output.gen_image_slices[batch_idx]
+                if batch_idx is not None
+                else output.gen_image_slices
+            )
+            slices = slices + gen_image_slices
+        return slices
+
+    def build_img_ratio_slice_logits_processor(self, tokenizer):
+        if self.img_ratio_slice_logits_processor is None:
+            self.img_ratio_slice_logits_processor = LogitsProcessorList()
+            self.img_ratio_slice_logits_processor.append(
+                SliceVocabLogitsProcessor(
+                    vocab_start=tokenizer.start_ratio_token_id,
+                    vocab_end=tokenizer.end_ratio_token_id + 1,
+                    other_slices=getattr(tokenizer, "ratio_token_other_slices", []),
+                )
+            )
+
+    def postprocess_outputs(self, outputs: list[Image.Image], batch_cond_images, infer_align_image_size: bool = False):
+        if infer_align_image_size:
+            target_area = self.vae_reso_group.base_size ** 2
+
+            for batch_index, (output_image, cond_images) in enumerate(zip(outputs, batch_cond_images)):
+                output_image_ratio_index = self.vae_reso_group.get_base_size_and_ratio_index(width=output_image.width, height=output_image.height)[1]
+                cond_images_ratio_index_list = []
+                cond_images_ori_width_list = []
+                cond_images_ori_height_list = []
+                for cond_image in cond_images:
+                    if isinstance(cond_image, ImageTensor):
+                        cond_images_ratio_index_list.append(cond_image.i.ratio_index)
+                        cond_images_ori_width_list.append(cond_image.i.ori_image_width)
+                        cond_images_ori_height_list.append(cond_image.i.ori_image_height)
+                    else: # CondImage
+                        cond_images_ratio_index_list.append(cond_image.vae_image.i.ratio_index)
+                        cond_images_ori_width_list.append(cond_image.vae_image.i.ori_image_width)
+                        cond_images_ori_height_list.append(cond_image.vae_image.i.ori_image_height)
+
+                if len(cond_images) == 0:
+                    continue
+                elif len(cond_images) == 1:
+                    if output_image_ratio_index == cond_images_ratio_index_list[0]:
+                        if abs(cond_images_ori_height_list[0] / cond_images_ori_width_list[0] - self.vae_reso_group[output_image_ratio_index].ratio) >= 0.01:
+                            scale = math.sqrt(target_area / (cond_images_ori_width_list[0] * cond_images_ori_height_list[0]))
+                            new_w = round(cond_images_ori_width_list[0] * scale)
+                            new_h = round(cond_images_ori_height_list[0] * scale)
+                            outputs[batch_index] = output_image.resize((new_w, new_h), resample=Image.Resampling.LANCZOS)
+                else:
+                    for cond_image_ratio_index, cond_image_ori_width, cond_image_ori_height in zip(cond_images_ratio_index_list, cond_images_ori_width_list, cond_images_ori_height_list):
+                        if output_image_ratio_index == cond_image_ratio_index:
+                            if abs(cond_image_ori_height / cond_image_ori_width - self.vae_reso_group[output_image_ratio_index].ratio) >= 0.01:
+                                scale = math.sqrt(target_area / (cond_image_ori_width * cond_image_ori_height))
+                                new_w = round(cond_image_ori_width * scale)
+                                new_h = round(cond_image_ori_height * scale)
+                                outputs[batch_index] = output_image.resize((new_w, new_h), resample=Image.Resampling.LANCZOS)
+                            break
+
+        return outputs
+
+__all__ = [
+    "HunyuanImage3ImageProcessor"
+]

model-0001-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f636cb156e0d38d0739bd8b879da5bcce568f6ff3121314bf54a744d70a3303e
+size 5348403192

model-0002-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d596c456ffb1c0369b296815dbf5aecae344658399f9dbd44b1d0a383cdb43f
+size 5344103080

model-0003-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:873864acf1cb6412ea2379cdf9a3517a64f64fbd584945ac225a5a5155a6b145
+size 5318937248

model-0004-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e1315b5aa43a77285047e916f67add4d8c63f08d894c723595edb8c7e3b3f09
+size 5353033424

model-0005-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e8a4e3175a906e0eb9b1e6b0ac45c9d26f4aad7a2f37f7d35c562f2a0154345
+size 5318937248

model-0006-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72ee1387a20f18cecced4fcd1dd19e6a5fe47473259f9901dc52d60a7b06a0cf
+size 5344103080

model-0007-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:735b9cc3f26729f3ac058d40dbab4dd6c4fc80fcef87814ecf979dde5b4809a7
+size 5318937256

model-0008-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:38ddb09d7f113bfd2b9b77d604251a27c1f5252fe161ef16138ddd3092c4a7b3
+size 5344103088

model-0009-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a670adad4ffc1c084887a144f40d83f2b737fb39d51b99e0a7ac719c0cd0e58
+size 5318937256

model-0010-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1bc6d5bd8e956687d7f1ecdc2dcf61731828f16758c6a0e61c6904a8e874f1ef
+size 5344103136

model-0011-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9bc082641d9f77d84c8b29a67cd4bdade9c754485741a4785fc550ecf17042a8
+size 5318937400

model-0012-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:493b7acccb94a8808c3851ab0fe585af0c39df1c1a36fd66305d6aea030f4292
+size 5344103232

model-0013-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a0ed8a1eadb8a750d238d490db84fd7546d31649991666b867811e5e2d59853
+size 5318937400

model-0014-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6023be614b51bbc4a987930cf1a40d5beb192b6cb33b916601813dcca8e0d98
+size 5344103232

model-0015-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf87b76628ed4c85a0d8d6259a224dfea5464eb2d59a7adea00a02e733574f72
+size 5318937400

model-0016-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4f83f99e3ba5038cb05e05a6419ff5bbc72f341f453d3be87cf975fdb798e8c
+size 5344103232

model-0017-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:791f22de90a61a7eab1bd8d00827920fac4b3591ff5f8c7b5ba3addcb23e564b
+size 5318937392

model-0018-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6239f3e75845c6fc21ce4b865b4d2bdcaece44363c549da1ca2209d3d49ef4c0
+size 5344636312

model-0019-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3eeffa236a37d693e6d2bc111d97a53f844a3d6f4adf99576841a2df8bc425d2
+size 5327334656

model-0020-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03a578aed706c5a11abeb8698a1eeb7aa1b05efac904087c0373c2218cc25a70
+size 5344103224

model-0021-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:443c86dd32419ef1b1891d876f0e5260dbf8a8758519bb70a41047394efdc53f
+size 5318937400

model-0022-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:99fef1fe82cdde878b6b7e022644571aba9710535be9fbbc9860f71eeb0d3efc
+size 5344103232

model-0023-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65c68d803db0691f317ef0a91a0937c84926be1dd9f5d7b622b4910b58ae9324
+size 5318937400

model-0024-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6323ef5306a06cae635a72aea64e2d34677cfd69beff18ffc90b248b2098257
+size 5344103232

model-0025-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a10239c412419c8cce8cc62bd6196775a037bc5c607b13018e55b0b13f1ecdef
+size 5318937400

model-0026-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94a48c3477b95f036ed5564e59e244400abd167da53d3b48af59870cb38d4c57
+size 5344103232

model-0027-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:77c33d36df1de8cac089fb38e56674d8824f20dae0609c18b0184a2b590b9881
+size 5318937400

model-0028-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:403c781276c72030dbe4b378cfb65bcbce7c40fb145bc42ba1295fe24c0ca59f
+size 5344103232

model-0029-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c4a8f2224fd91b6a2c1ba83c88e751d4a1cced6996c9c1e25bf92da64788c52
+size 5318937400

model-0030-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2223dd71e2a2826259094bf91c22cc5b70cdea5f26890fe332a7e9362cba3508
+size 5344103232

model-0031-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:459d8bbc4c0f52b2ebfeb79d74886a271216157c689748d1e6bbd97a8393a38b
+size 5302504652

model-0032-of-0032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65b72ec507661c7a146b3174fd0baaa2689a5c7e85a8736aa9778b1d146a4ec6
+size 3316299512

siglip2.py

@@ -0,0 +1,570 @@
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/LICENSE
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Optional, Tuple, Union
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+
+
+class Config(object):
+    def __init__(self, config):
+        if config is not None:
+            for key, value in config.items():
+                setattr(self, key, value)
+
+    def __getitem__(self, key):
+        return getattr(self, key, None)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+
+class Siglip2VisionEmbeddings(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.patch_size = config.patch_size
+
+        self.patch_embedding = nn.Linear(
+            in_features=config.num_channels * self.patch_size * self.patch_size,
+            out_features=self.embed_dim,
+        )
+
+        self.num_patches = config.num_patches
+        self.position_embedding_size = int(self.num_patches**0.5)
+        self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim)
+
+    @staticmethod
+    def resize_positional_embeddings(
+        positional_embeddings: torch.Tensor,
+        spatial_shapes: torch.LongTensor,
+        max_length: int,
+    ) -> torch.Tensor:
+        """
+        Resize positional embeddings to image-specific size and pad to a fixed size.
+
+        Args:
+            positional_embeddings (`torch.Tensor`):
+                Position embeddings of shape (height, width, embed_dim)
+            spatial_shapes (`torch.LongTensor`):
+                Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to
+            max_length (`int`):
+                Maximum length of the positional embeddings to pad resized positional embeddings to
+
+        Returns:
+            `torch.Tensor`: Embeddings of shape (batch_size, max_length, embed_dim)
+        """
+        batch_size = spatial_shapes.shape[0]
+        embed_dim = positional_embeddings.shape[-1]
+        source_dtype = positional_embeddings.dtype
+
+        resulted_positional_embeddings = torch.empty(
+            (batch_size, max_length, embed_dim),
+            device=positional_embeddings.device,
+            dtype=source_dtype,
+        )
+
+        # (height, width, embed_dim) -> (1, embed_dim, height, width) for interpolation
+        positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
+
+        # Upcast to float32 on CPU because antialias is not supported for bfloat16/float16 on CPU
+        if positional_embeddings.device.type == "cpu":
+            positional_embeddings = positional_embeddings.to(torch.float32)
+
+        for i in range(batch_size):
+            # (1, dim, height, width) -> (1, dim, target_height, target_width)
+            height, width = spatial_shapes[i]
+            resized_embeddings = F.interpolate(
+                positional_embeddings,
+                size=(height, width),
+                mode="bilinear",
+                align_corners=False,
+                antialias=True,
+            )
+
+            # (1, dim, target_height, target_width) -> (target_height * target_width, dim)
+            resized_embeddings = resized_embeddings.reshape(embed_dim, height * width).transpose(0, 1)
+
+            # Cast to original dtype
+            resized_embeddings = resized_embeddings.to(source_dtype)
+
+            resulted_positional_embeddings[i, : height * width] = resized_embeddings
+            resulted_positional_embeddings[i, height * width :] = resized_embeddings[0]
+
+        return resulted_positional_embeddings
+
+    def forward(self, pixel_values: torch.FloatTensor, spatial_shapes: torch.LongTensor) -> torch.Tensor:
+        """
+        Args:
+            pixel_values (`torch.FloatTensor`):
+                Pixel values of shape (batch_size, max_num_patches, num_channels * patch_size * patch_size)
+            spatial_shapes (`List[Tuple[int, int]]`):
+                Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to
+        """
+
+        # Apply patch embeddings to already patchified pixel values
+        target_dtype = self.patch_embedding.weight.dtype
+        patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
+
+        # Get positional resized and padded positional embeddings
+        positional_embeddings = self.position_embedding.weight.reshape(
+            self.position_embedding_size, self.position_embedding_size, -1
+        )
+        resized_positional_embeddings = self.resize_positional_embeddings(
+            positional_embeddings, spatial_shapes, max_length=pixel_values.shape[1]
+        )
+
+        # Add positional embeddings to patch embeddings
+        embeddings = patch_embeds + resized_positional_embeddings
+        return embeddings
+
+
+class Siglip2Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+        self.scale = self.head_dim**-0.5
+        self.dropout = config.attention_dropout
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Input shape: Batch x Time x Channel"""
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        k_v_seq_len = key_states.shape[-2]
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale
+
+        if attn_weights.size() != (batch_size, self.num_heads, q_len, k_v_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(batch_size, self.num_heads, q_len, k_v_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (batch_size, 1, q_len, k_v_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(batch_size, 1, q_len, k_v_seq_len)}, "
+                    f"but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (batch_size, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(batch_size, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights
+
+class Siglip2SdpaAttention(Siglip2Attention):
+    """
+    Siglip2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+    `Siglip2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt
+    to SDPA API.
+    """
+
+    is_causal = False
+
+    # Adapted from Siglip2Attention.forward and transformers.models.llama.modeling_llama.LlamaSdpaAttention.forward
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if output_attentions:
+            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"`
+            #  once this is implemented.
+            warnings.warn(
+                "Siglip2Model is using Siglip2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` "
+                "does not support `output_attentions=True`. Falling back to the manual attention implementation, "
+                'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. '
+                'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+            )
+            return super().forward(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                output_attentions=output_attentions,
+            )
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with
+        # custom attn_mask,
+        # Reference: https://github.com/pytorch/pytorch/issues/112577.
+        if query_states.device.type == "cuda" and attention_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+
+        # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an
+        # inline conditional assignment in SDPA to support both torch.compile's dynamic shapes and full graph options.
+        # An inline conditional prevents dynamic shapes from compiling.
+        is_causal = True if self.is_causal and q_len > 1 else False
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=is_causal,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(batch_size, q_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, None
+
+
+class Siglip2MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.activation_fn = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+class Siglip2EncoderLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.self_attn = Siglip2Attention(config=config)
+        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+        self.mlp = Siglip2MLP(config)
+        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+    # Ignore copy
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`):
+                Input to the layer of shape `(batch, seq_len, embed_dim)`.
+            attention_mask (`torch.FloatTensor`):
+                Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very
+                large negative values.
+            output_attentions (`bool`, *optional*, defaults to `False`):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+        """
+        residual = hidden_states
+
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+class Siglip2Encoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`Siglip2EncoderLayer`].
+
+    Args:
+        config: Siglip2Config
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([Siglip2EncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = True
+
+    # Ignore copy
+    def forward(
+        self,
+        inputs_embeds,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
+                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+                than the model's internal embedding lookup matrix.
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        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
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        hidden_states = inputs_embeds
+        for layer_index, encoder_layer in enumerate(self.layers): # len(self.layers): 27
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+
+            layer_outputs = encoder_layer(
+                hidden_states,
+                attention_mask,
+                output_attentions=output_attentions,
+            )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
+        )
+
+
+class Siglip2MultiheadAttentionPoolingHead(nn.Module):
+    """Multihead Attention Pooling."""
+
+    def __init__(self, config):
+        super().__init__()
+
+        self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
+        self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True)
+        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.mlp = Siglip2MLP(config)
+        self.num_heads = config.num_attention_heads
+
+    def forward(self, hidden_state: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
+        batch_size = hidden_state.shape[0]
+        probe = self.probe.repeat(batch_size, 1, 1)
+
+        if attention_mask is not None:
+            target_len, source_len = probe.shape[1], hidden_state.shape[1]
+            attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_state.dtype, target_len)
+            attention_mask = attention_mask.repeat(1, self.num_heads, target_len, 1)
+            attention_mask = attention_mask.reshape(-1, target_len, source_len)
+
+        hidden_state = self.attention(probe, hidden_state, hidden_state, attn_mask=attention_mask)[0]
+
+        residual = hidden_state
+        hidden_state = self.layernorm(hidden_state)
+        hidden_state = residual + self.mlp(hidden_state)
+
+        return hidden_state[:, 0]
+
+
+class Siglip2VisionTransformer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        config = Config(config)
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.embeddings = Siglip2VisionEmbeddings(config)
+        self.encoder = Siglip2Encoder(config)
+        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+        self.use_head = True if not hasattr(config, "vision_use_head") else config.vision_use_head
+        if self.use_head:
+            self.head = Siglip2MultiheadAttentionPoolingHead(config)
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        attention_mask: torch.Tensor,
+        spatial_shapes: torch.LongTensor,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+
+        """
+        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
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        hidden_states = self.embeddings(pixel_values, spatial_shapes)
+
+        if attention_mask is not None and not self._use_flash_attention_2:
+            # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
+            encoder_attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
+        else:
+            encoder_attention_mask = attention_mask
+
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        last_hidden_state = encoder_outputs[0]
+        last_hidden_state = self.post_layernorm(last_hidden_state)
+
+        pooler_output = self.head(last_hidden_state, attention_mask) if self.use_head else None
+        if not return_dict:
+            return (last_hidden_state, pooler_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooler_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+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
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+__all__ = [
+    "Siglip2VisionTransformer",
+    "LightProjector",
+]

system_prompt.py

@@ -0,0 +1,206 @@
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/LICENSE
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+t2i_system_prompt_en_vanilla = """
+You are an advanced AI text-to-image generation system. Given a detailed text prompt, your task is to create a high-quality, visually compelling image that accurately represents the described scene, characters, or objects. Pay careful attention to style, color, lighting, perspective, and any specific instructions provided.
+"""
+
+# 775
+t2i_system_prompt_en_recaption = """
+You are a world-class image generation prompt expert. Your task is to rewrite a user's simple description into a **structured, objective, and detail-rich** professional-level prompt.
+
+The final output must be wrapped in `<recaption>` tags.
+
+### **Universal Core Principles**
+
+When rewriting the prompt (inside the `<recaption>` tags), you must adhere to the following principles:
+
+1.  **Absolute Objectivity**: Describe only what is visually present. Avoid subjective words like "beautiful" or "sad". Convey aesthetic qualities through specific descriptions of color, light, shadow, and composition.
+2.  **Physical and Logical Consistency**: All scene elements (e.g., gravity, light, shadows, reflections, spatial relationships, object proportions) must strictly adhere to real-world physics and common sense. For example, tennis players must be on opposite sides of the net; objects cannot float without a cause.
+3.  **Structured Description**: Strictly follow a logical order: from general to specific, background to foreground, and primary to secondary elements. Use directional terms like "foreground," "mid-ground," "background," and "left side of the frame" to clearly define the spatial layout.
+4.  **Use Present Tense**: Describe the scene from an observer's perspective using the present tense, such as "A man stands..." or "Light shines on..."
+5.  **Use Rich and Specific Descriptive Language**: Use precise adjectives to describe the quantity, size, shape, color, and other attributes of objects, subjects, and text. Vague expressions are strictly prohibited.
+
+If the user specifies a style (e.g., oil painting, anime, UI design, text rendering), strictly adhere to that style. Otherwise, first infer a suitable style from the user's input. If there is no clear stylistic preference, default to an **ultra-realistic photographic style**. Then, generate the detailed rewritten prompt according to the **Style-Specific Creation Guide** below:
+
+### **Style-Specific Creation Guide**
+
+Based on the determined artistic style, apply the corresponding professional knowledge.
+
+**1. Photography and Realism Style**
+*   Utilize professional photography terms (e.g., lighting, lens, composition) and meticulously detail material textures, physical attributes of subjects, and environmental details.
+
+**2. Illustration and Painting Style**
+*   Clearly specify the artistic school (e.g., Japanese Cel Shading, Impasto Oil Painting) and focus on describing its unique medium characteristics, such as line quality, brushstroke texture, or paint properties.
+
+**3. Graphic/UI/APP Design Style**
+*   Objectively describe the final product, clearly defining the layout, elements, and color palette. All text on the interface must be enclosed in double quotes `""` to specify its exact content (e.g., "Login"). Vague descriptions are strictly forbidden.
+
+**4. Typographic Art**
+*   The text must be described as a complete physical object. The description must begin with the text itself. Use a straightforward front-on or top-down perspective to ensure the entire text is visible without cropping.
+
+### **Final Output Requirements**
+
+1.  **Output the Final Prompt Only**: Do not show any thought process, Markdown formatting, or line breaks.
+2.  **Adhere to the Input**: You must retain the core concepts, attributes, and any specified text from the user's input.
+3.  **Style Reinforcement**: Mention the core style 3-5 times within the prompt and conclude with a style declaration sentence.
+4.  **Avoid Self-Reference**: Describe the image content directly. Remove redundant phrases like "This image shows..." or "The scene depicts..."
+5.  **The final output must be wrapped in `<recaption>xxxx</recaption>` tags.**
+
+The user will now provide an input prompt. You will provide the expanded prompt.
+"""
+
+# 890
+t2i_system_prompt_en_think_recaption = """
+You will act as a top-tier Text-to-Image AI. Your core task is to deeply analyze the user's text input and transform it into a detailed, artistic, and fully user-intent-compliant image.
+
+Your workflow is divided into two phases:
+
+1. Thinking Phase (<think>): In the <think> tag, you need to conduct a structured thinking process, progressively breaking down and enriching the constituent elements of the image. This process must include, but is not limited to, the following dimensions:
+
+Subject: Clearly define the core character(s) or object(s) in the scene, including their appearance, posture, expression, and emotion.
+Composition: Set the camera angle and layout, such as close-up, long shot, bird's-eye view, golden ratio composition, etc.
+Environment/Background: Describe the scene where the subject is located, including the location, time of day, weather, and other elements in the background.
+Lighting: Define the type, direction, and quality of the light source, such as soft afternoon sunlight, cool tones of neon lights, dramatic Rembrandt lighting, etc., to create a specific atmosphere.
+Color Palette: Set the main color tone and color scheme of the image, such as vibrant and saturated, low-saturation Morandi colors, black and white, etc.
+Quality/Style: Determine the artistic style and technical details of the image. This includes user-specified styles (e.g., anime, oil painting) or the default realistic style, as well as camera parameters (e.g., focal length, aperture, depth of field).
+Details: Add minute elements that enhance the realism and narrative quality of the image, such as a character's accessories, the texture of a surface, dust particles in the air, etc.
+
+
+2. Recaption Phase (<recaption>): In the <recaption> tag, merge all the key details from the thinking process into a coherent, precise, and visually evocative final description. This description is the direct instruction for generating the image, so it must be clear, unambiguous, and organized in a way that is most suitable for an image generation engine to understand.
+
+Absolutely Objective: Describe only what is visually present. Avoid subjective words like "beautiful" or "sad." Convey aesthetic sense through concrete descriptions of colors, light, shadow, and composition.
+
+Physical and Logical Consistency: All scene elements (e.g., gravity, light and shadow, reflections, spatial relationships, object proportions) must strictly adhere to the physical laws of the real world and common sense. For example, in a tennis match, players must be on opposite sides of the net; objects cannot float without reason.
+
+Structured Description: Strictly follow a logical order: from whole to part, background to foreground, and primary to secondary. Use directional words like "foreground," "mid-ground," "background," "left side of the frame" to clearly define the spatial layout.
+
+Use Present Tense: Describe from an observer's perspective using the present tense, such as "a man stands," "light shines on..."
+Use Rich and Specific Descriptive Language: Use precise adjectives to describe the quantity, size, shape, color, and other attributes of objects/characters/text. Absolutely avoid any vague expressions.
+
+
+Output Format:
+<think>Thinking process</think><recaption>Refined image description</recaption>Generate Image
+
+
+You must strictly adhere to the following rules:
+
+1. Faithful to Intent, Reasonable Expansion: You can creatively add details to the user's description to enhance the image's realism and artistic quality. However, all additions must be highly consistent with the user's core intent and never introduce irrelevant or conflicting elements.
+2. Style Handling: When the user does not specify a style, you must default to an "Ultra-realistic, Photorealistic" style. If the user explicitly specifies a style (e.g., anime, watercolor, oil painting, cyberpunk, etc.), both your thinking process and final description must strictly follow and reflect that specified style.
+3. Text Rendering: If specific text needs to appear in the image (such as words on a sign, a book title), you must enclose this text in English double quotes (""). Descriptive text must not use double quotes.
+4. Design-related Images: You need to specify all text and graphical elements that appear in the image and clearly describe their design details, including font, color, size, position, arrangement, visual effects, etc.
+"""
+
+t2i_system_prompts = {
+    "en_vanilla": [t2i_system_prompt_en_vanilla],
+    "en_recaption": [t2i_system_prompt_en_recaption],
+    "en_think_recaption": [t2i_system_prompt_en_think_recaption]
+}
+
+
+unified_system_prompt_en = """You are an advanced multimodal model whose core mission is to analyze user intent and generate high-quality text and images.
+
+#### Four Core Capabilities
+1.  **Text-to-Text (T2T):** Generate coherent text responses from text prompts.
+2.  **Text-to-Image (T2I):** Generate high-quality images from text prompts.
+3.  **Text & Image to Text (TI2T):** Generate accurate text responses based on a combination of images and text.
+4.  **Text & Image to Image (TI2I):** Generate modified images based on a reference image and editing instructions.
+
+---
+### Image Generation Protocol (for T2I & TI2I)
+You will operate in one of two modes, determined by the user's starting tag:
+#### **<recaption> Mode (Prompt Rewriting)**:
+*   **Trigger:** Input begins with `<recaption>`.
+*   **Task:** Immediately rewrite the user's text into a structured, objective, and detail-rich professional-grade prompt.
+*   **Output:** Output only the rewritten prompt within `<recaption>` tags: `<recaption>Rewritten professional-grade prompt</recaption>`
+
+#### **<think> Mode (Think + Rewrite)**:
+*   **Trigger:** Input begins with `<think>`.
+*   **Task:** First, conduct a structured analysis of the request within `<think>` tags. Then, output the professional prompt, rewritten based on the analysis, within `<recaption>` tags.
+*   **Output:** Strictly adhere to the format: `<think>Analysis process</think><recaption>Rewritten prompt</recaption>`
+
+---
+### Execution Standards and Guidelines
+#### **`<think>` Phase: Analysis Guidelines**
+**For T2I (New Image Generation):**
+Deconstruct the user's request into the following core visual components:
+*   **Subject:** Key features of the main character/object, including appearance, pose, expression, and emotion.
+*   **Composition:** Camera angle, lens type, and layout.
+*   **Environment/Background:** The setting, time of day, weather, and background elements.
+*   **Lighting:** Technical details such as light source type, direction, and quality.
+*   **Color Palette:** The dominant hues and overall color scheme.
+*   **Style/Quality:** The artistic style, clarity, depth of field, and other technical details.
+*   **Text:** Identify any text to be rendered in the image, including its content, style, and position.
+*   **Details:** Small elements that add narrative depth and realism.
+
+**For TI2I (Image Editing):**
+Adopt a task-diagnostic approach:
+1.  **Diagnose Task:** Identify the edit type and analyze key requirements.
+2.  **Prioritize Analysis:**
+    *   **Adding:** Analyze the new element's position and appearance, ensuring seamless integration with the original image's lighting, shadows, and style.
+    *   **Removing:** Identify the target for removal and determine how to logically fill the resulting space using surrounding textures and lighting.
+    *   **Modifying:** Analyze what to change and what it should become, while emphasizing which elements must remain unchanged.
+    *   **Style Transfer:** Deconstruct the target style into specific features (e.g., brushstrokes, color palette) and apply them to the original image.
+    *   **Text Editing:** Ensure correct content and format. Consider the text's visual style (e.g., font, color, material) and how it adapts to the surface's perspective, curvature, and lighting.
+    *   **Reference Editing:** Extract specific visual elements (e.g., appearance, posture, composition, lines, depth) from the reference image to generate an image that aligns with the text description while also incorporating the referenced content.
+    *   **Inferential Editing:** Identify vague requests (e.g., "make it more professional") and translate them into concrete visual descriptions.
+
+#### `<recaption>` Phase: Professional-Grade Prompt Generation Rules
+**General Rewriting Principles (for T2I & TI2I):**
+1.  **Structure & Logic:** Start with a global description. Use positional words (e.g., "foreground", "background") to define the layout.
+2.  **Absolute Objectivity:** Avoid subjective terms. Convey aesthetics through precise descriptions of color, light, shadow, and materials.
+3.  **Physical & Logical Consistency:** Ensure all descriptions adhere to the laws of physics and common sense.
+4.  **Fidelity to User Intent:** Preserve the user's core concepts, subjects, and attributes. Text to be rendered in the image **must be enclosed in double quotes ("")**.
+5.  **Camera & Resolution:** Translate camera parameters into descriptions of visual effects. Convert resolution information into natural language.
+
+**T2I-Specific Guidelines:**
+*   **Style Adherence & Inference:** Strictly follow the specified style. If none is given, infer the most appropriate style and detail it using professional terminology.
+*   **Style Detailing:**
+    *   **Photography/Realism:** Use professional photography terms to describe lighting, lens effects, and material textures.
+    *   **Painting/Illustration:** Specify the art movement or medium's characteristics.
+    *   **UI/Design:** Objectively describe the final product. Define layout, elements, and typography. Text content must be specific and unambiguous.
+
+**TI2I-Specific Guidelines:**
+*   **Preserve Unchanged Elements:** Emphasize elements that **remain unchanged**. Unless explicitly instructed, never alter a character's identity/appearance, the core background, camera angle, or overall style.
+*   **Clear Editing Instructions:**
+    *   **Replacement:** Use the logic "**replace B with A**," and provide a detailed description of A.
+    *   **Addition:** Clearly state what to add, where, and what it looks like.
+*   **Unambiguous Referencing:** Avoid vague references (e.g., "that person"). Use specific descriptions of appearance.
+"""
+
+
+def get_system_prompt(sys_type, bot_task, system_prompt=None):
+    if sys_type == 'None':
+        return None
+    elif sys_type == "en_unified":
+        return unified_system_prompt_en
+    elif sys_type in ['en_vanilla', 'en_recaption', 'en_think_recaption']:
+        return t2i_system_prompts[sys_type][0]
+    elif sys_type == "dynamic":
+        if bot_task == "think":
+            return t2i_system_prompts["en_think_recaption"][0]
+        elif bot_task == "recaption":
+            return t2i_system_prompts["en_recaption"][0]
+        elif bot_task == "image":
+            return t2i_system_prompts["en_vanilla"][0].strip("\n")
+        else:
+            return system_prompt
+    elif sys_type == 'custom':
+        return system_prompt
+    else:
+        raise NotImplementedError(f"Unsupported system prompt type: {sys_type}")
+
+
+__all__ = [
+    "get_system_prompt"
+]

tokenization_hunyuan_image_3.py

@@ -0,0 +1,1773 @@
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/LICENSE
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import random
+from collections import defaultdict
+from copy import deepcopy
+from dataclasses import dataclass
+from enum import IntEnum, auto
+from typing import List, Tuple, Dict
+from typing import Optional, Union, Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from diffusers.utils import BaseOutput
+
+from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
+
+
+def default(value, default_value):
+    return value if value is not None else default_value
+
+
+def ensure_list(value):
+    if value is None:
+        return []
+    if isinstance(value, (list, tuple)):
+        return list(value)
+    return [value]
+
+
+class Resolution(object):
+    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(object):
+    def __init__(self, base_size=None, step=None, align=1, extra_resolutions=None):
+        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()
+
+        if extra_resolutions is not None:
+            for extra_resolution in extra_resolutions:
+                height, width = extra_resolution.height, extra_resolution.width
+                ratio = height / width
+                flag = True
+                for resolution in self.data:
+                    if resolution.ratio == ratio:
+                        flag = False
+                        break
+                if flag:
+                    self.data.append(extra_resolution)
+
+        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,
+            ori_image_width: int = None,
+            ori_image_height: int = None,
+            **kwargs,
+    ):
+        self.image_type = image_type
+        self.image_tensor = image_tensor
+        self.ori_image_width = ori_image_width
+        self.image_width = image_width
+        self.w = image_width
+        self.ori_image_height = ori_image_height
+        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)
+        self.add_timestep_r_token = kwargs.get("add_timestep_r_token", False)
+
+    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"ori_image_width={self.ori_image_width}, ori_image_height={self.ori_image_height}, "
+                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,
+                add_timestep_r_token=self.add_timestep_r_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,
+                ori_image_width=self.ori_image_width,
+                ori_image_height=self.ori_image_height,
+            )
+        elif self.image_type in ["vit", "siglip2"]:
+            return dict(
+                token_length=self.image_token_length,
+                use_front_boi_token=self.use_front_boi_token,
+                add_image_shape_token=self.add_image_shape_token,
+                # 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,
+                ori_image_width=self.ori_image_width,
+                ori_image_height=self.ori_image_height,
+            )
+        else:
+            raise ValueError(f"Unknown image type '{self.image_type}'")
+
+    @property
+    def num_special_tokens(self):
+        if self.args is None:
+            raise ValueError("meta_info requires `args` attribute to be set.")
+        if self.image_type in ["vae", "src_image", "gen_image"]:
+            count = (
+                    2 +  # <boi> + <eoi> or <src_boi> + <src_eoi>
+                    (1 if self.add_timestep_token else 0) +
+                    (1 if self.add_guidance_token else 0) +
+                    (1 if self.add_timestep_r_token else 0) +
+                    (2 if self.add_image_shape_token else 0)
+            )
+        else:
+            raise ValueError(f"Unknown image_type: {self.image_type}")
+        return count
+
+    def copy(self, copy_image_tensor=True):
+        if copy_image_tensor and self.image_tensor is None:
+            raise ValueError("image_tensor is None, cannot copy")
+        return ImageInfo(
+            image_type=self.image_type,
+            image_tensor=self.image_tensor.clone() if copy_image_tensor else None,
+            image_width=self.image_width,
+            image_height=self.image_height,
+            ori_image_width=self.ori_image_width,
+            ori_image_height=self.ori_image_height,
+            token_width=self.token_width,
+            token_height=self.token_height,
+            image_token_length=self.image_token_length,
+            base_size=self.base_size,
+            ratio_index=self.ratio_index,
+        )
+
+    def zeros_(self):
+        self.image_tensor = torch.zeros_like(self.image_tensor)
+
+
+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(object):
+    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 +  # <boi> + <eoi>
+                (1 if self.add_timestep_token else 0) +
+                (2 if self.add_image_shape_token else 0) +
+                1   # <joint_image_sep>
+        )
+
+    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 CondImage(object):
+    def __init__(self, image_type: str, vae_image: ImageTensor, vit_image: ImageTensor):
+        self.image_type = image_type
+        self.vae_image = vae_image
+        self.vit_image = vit_image
+
+        if image_type == "vae":
+            self.i = vae_image.i
+            self.section_type = "cond_vae_image"
+
+        elif image_type == "vit":
+            self.i = vit_image.i
+            self.section_type = "cond_vit_image"
+
+        elif image_type == "vae_vit":
+            self.i = JointImageInfo(vae_image.i, vit_image.i)
+            self.section_type = "cond_joint_image"
+
+        else:
+            raise ValueError(f"Unknown image_type: {image_type}")
+
+
+class TokenizerEncodeOutput(BaseOutput):
+    tokens: torch.Tensor = None
+    text_slices: Optional[list[slice]] = None
+    vae_image_slices: Optional[list[slice]] = None
+    gen_image_slices: Optional[list[slice]] = None
+    vit_image_slices: Optional[list[slice]] = None
+    joint_image_slices: Optional[list[slice]] = None
+    all_image_slices: Optional[list[slice]] = None
+    text_mask: Optional[torch.Tensor] = None
+    vae_image_mask: Optional[torch.Tensor] = None
+    gen_image_mask: Optional[torch.Tensor] = None
+    vit_image_mask: Optional[torch.Tensor] = None
+    real_pos: Optional[torch.Tensor] = None
+    guidance_scatter_index: Optional[torch.Tensor] = None
+    cond_timestep_scatter_index: Optional[torch.Tensor] = None
+    gen_timestep_scatter_index: Optional[torch.Tensor] = None
+    gen_timestep_r_scatter_index: Optional[torch.Tensor] = None
+
+
+class SeparatorStyle(IntEnum):
+    ADD_COLON_SPACE_SINGLE = auto()
+    NONE = auto()
+
+
+@dataclass
+class Conversation(object):
+    name: str
+    system_template: str = "{system_message}"
+    system_message: str = ""
+    roles: Tuple[str, str] = ("User", "Assistant")
+    messages: List[List[str]] = ()
+    sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SPACE_SINGLE
+    sep: str = "\n"
+    sep2: str = None
+    sep_sp: str = None
+    stop_token_ids: list[int] = None
+
+    def get_prompt(self, return_type="str", add_system=True):
+        system_prompt = self.system_template.format(system_message=self.system_message)
+        prompt_list = []
+
+        if self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
+            seps = [self.sep, self.sep2]
+            if add_system:
+                prompt_list.append(("System", system_prompt + self.sep_sp if system_prompt else ""))
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    prompt_list.append((role, f"{role}: {message}{seps[i % 2]}"))
+                else:
+                    prompt_list.append((role, f"{role}: "))
+
+        elif self.sep_style == SeparatorStyle.NONE:
+            seps = [self.sep, self.sep2]
+            if add_system:
+                prompt_list.append(("System", system_prompt + self.sep_sp if system_prompt else ""))
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    prompt_list.append((role, f"{role}{message}{seps[i % 2]}"))
+                else:
+                    prompt_list.append((role, f"{role}"))
+        else:
+            raise NotImplementedError(f"Unsupported sep_style: {self.sep_style}")
+
+        if return_type == "str":
+            prompt = "".join([msg for _, msg in prompt_list])
+        else:
+            prompt = prompt_list
+
+        return prompt
+
+    def get_role_prefix(self, role):
+        if self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
+            return f"{role}: "
+        elif self.sep_style == SeparatorStyle.NONE:
+            return f"{role}"
+        else:
+            raise NotImplementedError(f"Unsupported sep_style: {self.sep_style}")
+
+    def set_system_message(self, system_message: str):
+        """Set the system message."""
+        self.system_message = system_message
+
+    def add_message(self, role: str, message: str):
+        """Append a new message."""
+        self.messages.append([role, message])
+
+    def copy(self):
+        return deepcopy(self)
+
+    def empty(self, name=None):
+        """Return an empty conversation with the same template."""
+        return Conversation(
+            name=name or self.name,
+            system_template=self.system_template,
+            system_message="",
+            roles=self.roles,
+            messages=[],
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            sep_sp=self.sep_sp,
+            stop_token_ids=self.stop_token_ids,
+        )
+
+
+# A global registry for all conversation templates
+conv_templates: Dict[str, Conversation] = {}
+
+
+def register_conv_template(template: Conversation, override: bool = False):
+    """Register a new conversation template."""
+    if not override:
+        assert (
+            template.name not in conv_templates
+        ), f"{template.name} has been registered."
+
+    conv_templates[template.name] = template
+
+
+register_conv_template(
+    Conversation(
+        name="hunyuan-image-3",
+        system_template="{system_message}",
+        system_message="",
+        roles=("User", "Assistant"),
+        messages=[],
+        sep_style=SeparatorStyle.ADD_COLON_SPACE_SINGLE,
+        sep="\n\n",
+        sep2="<|endoftext|>",
+        sep_sp="\n\n",
+        stop_token_ids=[127957],
+    )
+)
+
+
+def get_conversation_template(name: str) -> Conversation:
+    """Get a conversation template."""
+    return conv_templates[name].copy()
+
+
+class HunyuanImage3TokenizerFast(PreTrainedTokenizerFast):
+    """
+    Tokenizer for Hunyuan Multimodal models, utilizing a fast tokenizer backend.
+    This tokenizer extends the PreTrainedTokenizerFast from Hugging Face Transformers
+    for multimodal tasks.
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # A convenience mapping for special tokens
+        special_tokens = self.special_tokens_map.get('additional_special_tokens', [])
+        if len(special_tokens) > 0:
+            special_token_ids = self.convert_tokens_to_ids(special_tokens)
+            self._sp_dict = dict(zip(special_tokens, special_token_ids))
+        else:
+            self._sp_dict = dict()
+
+        # Assign commonly used special tokens to attributes for easy access.
+        self.setup_special_tokens()
+
+        # Define decorator section
+        self.conversation_template = kwargs.get("conversation_template", "hunyuan-image-3")
+        self.conversation = get_conversation_template(self.conversation_template)
+        self.sequence_template = kwargs.get("sequence_template", "instruct")
+        self.decorator_section = DecoratorSections(
+            self,
+            conv=self.conversation,
+            sequence_template=self.sequence_template,
+        )
+
+    def setup_special_tokens(self):
+        # Define names for commonly used special tokens
+        predefined_name_mapping = dict(
+            boi="<boi>",
+            eoi="<eoi>",
+            boa="<boa>",
+            eoa="<eoa>",
+            bov="<bov>",
+            eov="<eov>",
+            img="<img>",
+            audio="<audio>",
+            video="<video>",
+            cfg="<cfg>",
+            timestep="<timestep>",
+            timestep_r="<timestep_r>",
+            guidance="<guidance>",
+            joint_img_sep="<joint_img_sep>",
+            answer="<answer>",
+            end_of_answer="</answer>",
+            # for extended cot types
+            think="<think>",
+            end_of_think="</think>",
+            recaption="<recaption>",
+            end_of_recaption="</recaption>",
+            # for grounding
+            ref="<ref>",
+            end_of_ref="</ref>",
+            quad="<quad>",
+            end_of_quad="</quad>",
+        )
+        for name, mapping in predefined_name_mapping.items():
+            setattr(self, f"{name}_token", mapping)
+            setattr(self, f"{name}_token_id", self.convert_tokens_to_ids(mapping))
+
+        if len(self._sp_dict) > 0:
+            name_mapping = dict()
+            for name, token in name_mapping.items():
+                setattr(self, name, token)
+                setattr(self, f"{name}_id", self._sp_dict[token])
+            
+
+        self.start_ratio_token_id = self.convert_tokens_to_ids("<img_ratio_0>")
+        self.end_ratio_token_id = self.convert_tokens_to_ids("<img_ratio_32>")
+        self.ratio_token_other_slices = [
+            (self.convert_tokens_to_ids("<img_ratio_33>"), self.convert_tokens_to_ids("<img_ratio_36>") + 1)
+        ]
+
+    @property
+    def max_token_id(self):
+        return self.vocab_size
+
+    def size_token(self, size: int):
+        return f"<img_size_{size}>"
+
+    def size_token_id(self, size: int):
+        return self.convert_tokens_to_ids(f"<img_size_{size}>")
+
+    def ratio_token(self, ratio_idx: int):
+        return f"<img_ratio_{ratio_idx}>"
+
+    def ratio_token_id(self, ratio_idx: int):
+        return self.convert_tokens_to_ids(f"<img_ratio_{ratio_idx}>")
+    
+    def get_all_ratio_token_ids(self):
+        return [self.ratio_token_id(i) for i in range(37)]
+
+    def relation_token(self, relation_idx: int, end: bool = False):
+        if end:
+            return f"</relation_{relation_idx}>"
+        return f"<relation_{relation_idx}>"
+
+    def relation_token_id(self, relation_idx: int, end: bool = False):
+        if end:
+            return self.convert_tokens_to_ids(f"</relation_{relation_idx}>")
+        return self.convert_tokens_to_ids(f"<relation_{relation_idx}>")
+
+    def x_token(self, pos: int):
+        return f"<pos_x_{pos}>"
+
+    def x_token_id(self, pos: int):
+        return self.convert_tokens_to_ids(f"<pos_x_{pos}>")
+
+    def y_token(self, pos: int):
+        return f"<pos_y_{pos}>"
+
+    def y_token_id(self, pos: int):
+        return self.convert_tokens_to_ids(f"<pos_y_{pos}>")
+
+    def z_token(self, pos: int):
+        return f"<pos_z_{pos}>"
+
+    def z_token_id(self, pos: int):
+        return self.convert_tokens_to_ids(f"<pos_z_{pos}>")
+
+    def get_img_token(self):
+        if hasattr(self, "img_token"):
+            return self.img_token
+        else:
+            return self.convert_ids_to_tokens(len(self) - 1)
+
+    def encode_text(
+            self,
+            *texts,
+            uncond_enabled: Optional[bool | list[bool]] = None,
+            uncond_p: Optional[float] = None,
+            max_length: Optional[int] = None,
+            pad: Optional[str] = None,
+            return_lengths: bool = False,
+    ):
+        """
+        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.
+        **<bos> token is always prepended to the text tokens.**
+
+        Parameters
+        ----------
+        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: Optional[str]
+            Padding method. Can be 'left' or 'right'.
+        return_lengths: bool
+            Whether to return the length of each encoded text.
+        """
+        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:
+                self.logger.warning(
+                    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 <bos> token in some LLMs.
+            if isinstance(text, str):
+                text_token = self.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 <pad> 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 {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: bool = False,
+            add_timestep_r_token: bool = False,
+            add_image_shape_token: bool = False,
+            add_guidance_token: bool = False,
+            base_size=None,
+            ratio_idx=None,
+            image_type=None,
+    ):
+        if add_image_shape_token:
+            token_seq.extend([self.size_token_id(base_size), self.ratio_token_id(ratio_idx)])
+            token_count += 2
+        if add_timestep_token:
+            token_seq.extend([self.timestep_token_id])
+            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 ["cond_joint_image", "cond_vae_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.guidance_token_id])
+            extra_token_pos['guidance'].append(token_count)
+            token_count += 1
+        if add_timestep_r_token:
+            token_seq.extend([self.timestep_r_token_id])
+            extra_token_pos['gen_timestep_r'].append(token_count)
+            token_count += 1
+        return token_count
+
+    def encode_sequence(
+            self,
+            template: str,
+            token_source: dict[str, list[list[int] | dict[str, Any]]],
+            total_length=None,
+            add_timestep_token=False,
+            add_timestep_r_token=False,
+            add_guidance_token=False,
+            add_eos=True,
+            add_pad=True,
+            add_bos=True,
+            drop_last: str | bool = 'auto',
+            add_image_shape_token=False,
+    ):
+        """
+        Encode a sequence of tokens based on the provided token source.
+
+        Args:
+            template: str
+                Template of the sequence. E.g., "text-image" means the sequence is composed of text and an image.
+                "text-image-image" means the sequence is composed of text and two images.
+            token_source (dict[str, list[list[int] | dict[str, Any]]]): Token source for each key in the template, in order.
+                - text: List[List[int]]. Each List[int] is a sequence of tokenized text tokens.
+                - gen_image: ImageInfo
+                - cond_joint_image: JointImageInfo
+                - cond_vae_image: ImageInfo
+                - cond_vit_image: ImageInfo
+            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 <iw><ih> token or <img_ratio_*><img_size_*> tokens)
+            add_guidance_token: bool
+                Whether to add guidance token before the image tokens.
+            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 <eos>.
+            use_front_boi_token: bool:
+                Whether to put the <boi> token at the front of iw, ih and timestep tokens.
+            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.
+            add_image_shape_token: bool
+                Whether to add image shape token before the image tokens. (Right before the <timestep> token)
+
+        Returns
+        -------
+        token_seq: list
+            Encoded token sequence.
+        extra_token_pos: dict
+            Positions of extra tokens. E.g., iw_ih, timestep.
+        """
+        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, face 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["<text>_start"].append(token_count)
+                token_count += len(source)
+                extra_token_pos["<text>_end"].append(token_count - 1)
+
+            elif key == "gen_image":
+                # 2 means boi and eoi
+                extra_count = \
+                    2 \
+                    + (1 if source.get('timestep', add_timestep_token) else 0) \
+                    + (1 if source.get('timestep_r', add_timestep_r_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
+                token_seq.append(self.boi_token_id)  # Use patched boi for Janus, otherwise using default <boi>
+                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_timestep_r_token=source.get('timestep_r', add_timestep_r_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,
+                )
+                token_seq.extend(
+                    [self.img_token_id] * source['length'] +  # token number
+                    [self.eoi_token_id]
+                )
+                extra_token_pos["<img>_start"].append(token_count)
+                extra_token_pos["<all_img>_start"].append(token_count)
+                token_count += source['length']
+                extra_token_pos["<img>_end"].append(token_count - 1)
+                extra_token_pos["<all_img>_end"].append(token_count - 1)
+                extra_token_pos["eoi"].append(token_count)
+                token_count += 1  # <eoi>
+
+            elif key == "cond_joint_image":
+                assert isinstance(source['length'], list) and len(
+                    source['length']) == 2, "cond_joint_image length should be a list of two integers"
+                # 2 + 1 means: boi, eoi, joint_img_sep
+                extra_count = \
+                    2 + 1 \
+                    + (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
+                token_seq.append(self.boi_token_id)  # Use patched boi for Janus, otherwise using default <boi>
+                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,
+                )
+                token_seq.extend(
+                    [self.img_token_id] * source['length'][0]
+                )
+                extra_token_pos["<vae_img>_start"].append(token_count)
+                extra_token_pos["<joint_img>_start"].append(token_count)
+                extra_token_pos["<all_img>_start"].append(token_count)
+                token_count += source['length'][0]
+                extra_token_pos["<vae_img>_end"].append(token_count - 1)
+                extra_token_pos["<all_img>_end"].append(token_count - 1)
+
+                token_seq.extend([self.joint_img_sep_token_id])
+                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["<vit_img>_start"].append(token_count)
+                extra_token_pos["<all_img>_start"].append(token_count)
+                token_count += source['length'][1]
+                extra_token_pos["<vit_img>_end"].append(token_count - 1)
+                extra_token_pos["<joint_img>_end"].append(token_count - 1)
+                extra_token_pos["<all_img>_end"].append(token_count - 1)
+
+                token_seq.extend(
+                    [self.eoi_token_id]
+                )
+                extra_token_pos["eoi"].append(token_count)
+                token_count += 1  # <eoi>
+
+            elif key == "cond_vae_image":
+                # 2 means: boi, eoi
+                extra_count = \
+                    2 \
+                    + (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 + source['length'] > total_length:
+                    drop_last_break = True
+                    break
+                token_seq.append(self.boi_token_id)  # Use patched boi for Janus, otherwise using default <boi>
+                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,
+                )
+                token_seq.extend(
+                    [self.img_token_id] * source['length']
+                )
+                extra_token_pos["<vae_img>_start"].append(token_count)
+                extra_token_pos["<all_img>_start"].append(token_count)
+                token_count += source['length']
+                extra_token_pos["<vae_img>_end"].append(token_count - 1)
+                extra_token_pos["<all_img>_end"].append(token_count - 1)
+                token_seq.extend(
+                    [self.eoi_token_id]
+                )
+                extra_token_pos["eoi"].append(token_count)
+                token_count += 1  # <eoi>
+
+            elif key == "cond_vit_image":
+                # 2 means: boi, eoi
+                extra_count = 2
+                if drop_last is True and token_count + extra_count + source['length'] > total_length:
+                    drop_last_break = True
+                    break
+
+                if hasattr(self, "boi_token_id"):
+                    token_seq.append(self.boi_token_id)
+                    extra_token_pos["boi"].append(token_count)
+                    token_count += 1
+
+                if hasattr(self, "img_token_id"):
+                    token_seq.extend([self.img_token_id] * source['length'])
+                else:
+                    # If not img_token_id defined, but we still need to fill the image tokens,
+                    # we use the last token id representing the image token.
+                    token_seq.extend([len(self) - 1] * source['length'])
+                extra_token_pos["<vit_img>_start"].append(token_count)
+                extra_token_pos["<all_img>_start"].append(token_count)
+                token_count += source['length']
+                extra_token_pos["<vit_img>_end"].append(token_count - 1)
+                extra_token_pos["<all_img>_end"].append(token_count - 1)
+
+                if hasattr(self, "eoi_token_id"):
+                    token_seq.append(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,
+                # src_image, und_image, face)
+                for start_key, end_key in [
+                    ("<img>_start", "<img>_end"), ("<vae_img>_start", "<vae_img>_end"),
+                    ("<vit_img>_start", "<vit_img>_end"), ("<joint>_start", "<joint>_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.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
+
+    @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: Optional[list[dict[str, Any]]] = None,
+        max_token_length: Optional[int] = None,
+        add_eos: bool | str = 'auto',
+        use_text_mask: bool = True,
+        add_pad: bool | str = 'auto',
+        add_bos: bool = True,
+        drop_last: bool | str = 'auto',
+    ):
+        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),
+                    timestep_r=section.get('add_timestep_r_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'] == 'cond_joint_image':
+                token_source['cond_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'),
+                ))
+            elif section['type'] == 'cond_vae_image':
+                token_source['cond_vae_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'),
+                ))
+            elif section['type'] == 'cond_vit_image':
+                token_source['cond_vit_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)
+
+        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_timestep_r_scatter_index = torch.tensor(extra_token_pos['gen_timestep_r'], dtype=torch.long) \
+            if 'gen_timestep_r' in extra_token_pos else None
+        gen_image_slices, gen_image_mask = self.parse_extra_token_pos(
+            extra_token_pos, 'img', full_seq_token_tensor)
+        vae_image_slices, vae_image_mask = self.parse_extra_token_pos(
+            extra_token_pos, 'vae_img', full_seq_token_tensor)
+        vit_image_slices, vit_image_mask = self.parse_extra_token_pos(
+            extra_token_pos, 'vit_img', full_seq_token_tensor)
+        joint_image_slices, _ = self.parse_extra_token_pos(
+            extra_token_pos, 'joint_img', full_seq_token_tensor)
+        # All image slices (src_image, gen_image, und_image)
+        all_image_slices = [
+            slice(start, end + 1)
+            for start, end in zip(extra_token_pos['<all_img>_start'], extra_token_pos['<all_img>_end'])
+        ] if '<all_img>_start' in extra_token_pos and '<all_img>_end' in extra_token_pos else []
+
+        # Text mask
+        text_slices = [
+            slice(start, end + 1)
+            for start, end in zip(extra_token_pos['<text>_start'], extra_token_pos['<text>_end'])
+        ] if '<text>_start' in extra_token_pos and '<text>_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 <pad> 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,
+            text_slices=text_slices,
+            gen_image_slices=gen_image_slices,
+            vae_image_slices=vae_image_slices,
+            vit_image_slices=vit_image_slices,
+            joint_image_slices=joint_image_slices,
+            all_image_slices=all_image_slices,
+            text_mask=text_mask,
+            gen_image_mask=gen_image_mask,
+            vae_image_mask=vae_image_mask,
+            vit_image_mask=vit_image_mask,
+            real_pos=real_pos,
+            guidance_scatter_index=guidance_scatter_index,
+            cond_timestep_scatter_index=cond_timestep_scatter_index,
+            gen_timestep_scatter_index=gen_timestep_scatter_index,
+            gen_timestep_r_scatter_index=gen_timestep_r_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 []
+        deco = self.decorator_section
+
+        if self.think_token in cot_text and self.end_of_think_token in cot_text:
+            before_think_sec = cot_text.split(self.think_token)[0]
+            after_think_sec = cot_text.split(self.end_of_think_token)[1]
+            think_sec = cot_text.split(self.think_token)[1].split(self.end_of_think_token)[0]
+            return self.get_cot_sections(before_think_sec, uncond_kwargs, drop_think=drop_think) + \
+                (deco.think(dict(type="text", text=think_sec, max_length=cot_max_length, **uncond_kwargs))
+                 if not drop_think else []) + \
+                self.get_cot_sections(after_think_sec, uncond_kwargs, drop_think=drop_think)
+
+        if self.recaption_token in cot_text and self.end_of_recaption_token in cot_text:
+            before_recaption_sec = cot_text.split(self.recaption_token)[0]
+            after_recaption_sec = cot_text.split(self.end_of_recaption_token)[1]
+            recaption_sec = cot_text.split(self.recaption_token)[1].split(self.end_of_recaption_token)[0]
+            return self.get_cot_sections(before_recaption_sec, uncond_kwargs, drop_think=drop_think) + \
+                (deco.recaption(dict(type="text", text=recaption_sec, max_length=cot_max_length, **uncond_kwargs))) + \
+                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=None,
+            uncond_p=0.0,
+            cfg_factor=1,
+            batchify=False,
+            image_base_size=None,
+            drop_think=False,
+    ):
+        if bot_task == "img_ratio":
+            assert image_base_size is not None, "image_base_size should be provided for img_ratio task."
+
+        # 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,
+            )
+
+        sequence_template = sequence_template or self.sequence_template
+        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 (self.recaption_token in text and self.end_of_recaption_token in text) or (
+                                self.think_token in text and self.end_of_think_token in text):
+                            _sub_sections.extend(self.get_cot_sections(text, uncond_kwargs, drop_think=drop_think))
+                        else:
+                            _sub_sections.append(dict(
+                            type="text", text=f"{answer_prefix}{text}{answer_suffix}", **uncond_kwargs))
+                    else:
+                        _sub_sections.append(dict(type="text", text=text, **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'] in ['cond_joint_image', 'cond_vae_image', 'cond_vit_image']:
+                    info = message['content']
+                    assert isinstance(info, (ImageInfo, JointImageInfo)), \
+                        f"Expected ImageInfo or 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 = self.answer_token, self.end_of_answer_token
+        else:
+            answer_prefix, answer_suffix = "", ""
+        if sequence_template == "pretrain":
+            system_suffix = ""
+            user_prefix = ""
+            user_suffix = ""
+            bot_prefix = ""
+            bot_suffix = ""
+        else:
+            conv = self.conversation
+            system_suffix = f"{conv.sep}"
+            user_prefix = conv.get_role_prefix(conv.roles[0])
+            user_suffix = f"{conv.sep}"
+            bot_prefix = conv.get_role_prefix(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 lastest message according to different tasks
+            bot_response_prefix = dict(
+                auto=_bot_prefix,
+                image="",
+                think=f"{_bot_prefix}{self.think_token}",
+                recaption=f"{_bot_prefix}{self.recaption_token}",
+                img_ratio=f"{_bot_prefix}{answer_prefix}{self.boi_token}{self.size_token(image_base_size)}",
+            )[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: Optional[list[str]] = None,
+            batch_message_list: Optional[list[list[dict[str, Any]]]] = None,
+            mode: str = "gen_text",
+            batch_gen_image_info: Optional[list[ImageInfo]] = None,
+            batch_cond_images: Optional[Union[list[CondImage], list[list[CondImage]]]] = None,
+            batch_system_prompt: Optional[list[str]] = None,
+            batch_cot_text: Optional[list[str]] = None,
+            max_length: Optional[int] = None,
+            bot_task: str = "auto",    # auto/image/think/recaption/img_ratio
+            image_base_size: Optional[int] = None,
+            sequence_template: str = "pretrain",
+            cfg_factor: int = 1,
+            add_assistant_prefix: Optional[bool] = 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_images is not None:
+                assert len(batch_cond_images) == batch_size, \
+                    (f"batch_cond_image_info should have the same length as batch_size ({batch_size}), "
+                     f"but got {len(batch_cond_images)}.")
+                batch_cond_images = [
+                    cond_images if isinstance(cond_images, list) else [cond_images]
+                    for cond_images in batch_cond_images
+                ]
+            else:
+                batch_cond_images = [[] 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_images in zip(
+                    batch_prompt, batch_system_prompt, batch_cot_text, batch_gen_image_info,
+                    batch_cond_images,
+            ):
+                message_list = []
+                # 1. system prompt section
+                if system_prompt:
+                    message_list.append(dict(role="system", type="text", content=system_prompt))
+                # 2. user inputs sections
+                #   2.1 image inputs
+                if len(cond_images) > 0:
+                    message_list.extend([
+                        dict(role="user", type=cond_image.section_type, content=cond_image.i)
+                        for cond_image in cond_images
+                    ])
+                #   2.2 text inputs
+                message_list.append(dict(role="user", type="text", content=prompt))
+                # 3. assistant answer sections
+                if cot_text is not None:
+                    message_list.append(dict(role="assistant", type="text", content=cot_text))
+                if mode == "gen_image":
+                    message_list.append(dict(
+                        role="assistant", type="gen_image", content=gen_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)
+
+    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 batch_gen_infer(
+            self,
+            infer_fn,
+            prompt_list: list,
+            negative_prompt_list: list = None,
+            infer_fn_kwargs_list: list[dict[str, int]] = None,
+            do_classifier_free_guidance=False,
+            condition_repeat_times: int = 1,
+            uncondition_repeat_times: int = 1,
+    ):
+        """
+        Batch inference for the AR-like model training of the text-to-image/instruction tuning tasks.
+
+        Parameters
+        ----------
+        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
+            List of negative prompts. Only used when do_classifier_free_guidance is True. If None, will use <cfg> token sequence as negative prompt.
+        infer_fn_kwargs_list: List[Dict[str, int]]
+            List of keyword arguments for the infer_fn.
+        do_classifier_free_guidance: bool
+            Whether to do classifier-free guidance.
+        condition_repeat_times: int
+        uncondition_repeat_times: int
+            Support multi-condition and multi-uncondition. e.g, [pred_cond, pred_uncond_text, pred_uncond_text_uncond_src]
+        """
+        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.pad_token_id
+                        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 == list:
+            return stacked_outputs
+        elif output_type == tuple:
+            return tuple(stacked_outputs)
+        elif output_num == 1:
+            return stacked_outputs[0]
+        else:
+            raise ValueError(f"Unsupported output type: {output_type}.")
+
+
+class DecoratorSections(object):
+    """ Define predefined sections in a multimodal template. """
+
+    def __init__(
+            self,
+            tokenizer: HunyuanImage3TokenizerFast,
+            conv: Conversation,
+            sequence_template: str,
+            ignore_start_tokens: Optional[set] = None,
+    ):
+        self.tokenizer = tokenizer
+        self.conv = conv
+        self.sequence_template = sequence_template
+        self.ignore_start_tokens = ignore_start_tokens or set()
+        self.roles = self.conv.roles
+
+        # Define sections based on the sequence template
+        if self.sequence_template == "pretrain":
+            self.user = []
+            self.user_sep = []
+            self.bot = []
+            self.bot_sep = []
+            self.answer_ = []
+            self._answer = []
+
+        elif self.sequence_template == "instruct":
+            self.user = [dict(type="text", text=self.conv.get_role_prefix(self.roles[0]), ignore=True)]
+            self.user_sep = [dict(type="text", text=self.conv.sep)]
+            self.bot = [dict(type="text", text=self.conv.get_role_prefix(self.roles[1]), ignore=True)]
+            self.bot_sep = [dict(type="text", text=self.conv.sep2)]
+            self.answer_ = [dict(type="text", text=self.tokenizer.answer_token,
+                                 ignore=(self.tokenizer.answer_token in self.ignore_start_tokens))]
+            self._answer = [dict(type="text", text=self.tokenizer.end_of_answer_token)]
+
+        else:
+            raise NotImplementedError(f"Unsupported sequence_template: {self.sequence_template}")
+
+        # Define eos token
+        eos_token = self.tokenizer.eos_token
+        if isinstance(eos_token, int):
+            eos_token = self.tokenizer.convert_ids_to_tokens(eos_token)
+        assert isinstance(eos_token, str), f"eos_token should be a string, got {type(eos_token)}."
+        self.eos = [dict(type="text", text=eos_token)]
+
+        # Define think sections
+        self.think_ = [dict(type="text", text=self.tokenizer.think_token,
+                            ignore=(self.tokenizer.think_token in self.ignore_start_tokens))]
+        self._think = [dict(type="text", text=self.tokenizer.end_of_think_token)]
+
+        # Define recaption sections
+        if hasattr(self.tokenizer, "recaption_token"):
+            self.recaption_ = [dict(type="text", text=self.tokenizer.recaption_token,
+                                    ignore=(self.tokenizer.recaption_token in self.ignore_start_tokens))]
+            self._recaption = [dict(type="text", text=self.tokenizer.end_of_recaption_token)]
+
+    def answer(self, section):
+        if isinstance(section, dict):
+            section = [section]
+        return self.answer_ + section + self._answer
+
+    def think(self, section):
+        if isinstance(section, dict):
+            section = [section]
+        return self.think_ + section + self._think
+
+    def recaption(self, section):
+        if not hasattr(self, "recaption_"):
+            raise AttributeError("This tokenizer does not support recaption sections.")
+        if isinstance(section, dict):
+            section = [section]
+        return self.recaption_ + section + self._recaption
+
+
+__all__ = [
+    "ResolutionGroup",
+    "ImageInfo",
+    "ImageTensor",
+    "JointImageInfo",
+    "CondImage",
+    "HunyuanImage3TokenizerFast",
+]

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2439418b9be76a54d8b5ff1ed89b37e36ec1735730f7da99b5aabb83a73db64
+size 25028750

tokenizer_config.json

@@ -0,0 +1,8 @@
+{
+    "bos_token": "<|startoftext|>",
+    "clean_up_tokenization_spaces": true,
+    "eos_token": "<|endoftext|>",
+    "model_max_length": 1048576,
+    "pad_token": "<pad>",
+    "tokenizer_class": "PreTrainedTokenizerFast"
+}

utils/__init__.py

@@ -0,0 +1,4 @@
+# Minimal utils module for hunyuan_image_3
+from .import_utils import _LazyModule, define_import_structure
+
+__all__ = ["_LazyModule", "define_import_structure"]