Upload 9 files

autoencoder_kl_3d.py

@@ -0,0 +1,793 @@
+# 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
+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
+
+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):
+    """ Attention with torch sdpa implementation. """
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+        self.q = Conv3d(in_channels, in_channels, kernel_size=1)
+        self.k = Conv3d(in_channels, in_channels, kernel_size=1)
+        self.v = Conv3d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = Conv3d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        b, c, f, h, w = q.shape
+        q = rearrange(q, "b c f h w -> b 1 (f h w) c").contiguous()
+        k = rearrange(k, "b c f h w -> b 1 (f h w) c").contiguous()
+        v = rearrange(v, "b c f h w -> b 1 (f h w) c").contiguous()
+        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+
+        return rearrange(h_, "b 1 (f h w) c -> b c f h w", f=f, h=h, w=w, c=c, b=b)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = swish(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = swish(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + h
+
+
+class 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):
+    """
+    The encoder network of AutoencoderKLConv3D.
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        z_channels: int,
+        block_out_channels: Tuple[int, ...],
+        num_res_blocks: int,
+        ffactor_spatial: int,
+        ffactor_temporal: int,
+        downsample_match_channel: bool = True,
+    ):
+        super().__init__()
+        assert block_out_channels[-1] % (2 * z_channels) == 0
+
+        self.z_channels = z_channels
+        self.block_out_channels = block_out_channels
+        self.num_res_blocks = num_res_blocks
+
+        # downsampling
+        self.conv_in = Conv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1)
+
+        self.down = nn.ModuleList()
+        block_in = block_out_channels[0]
+        for i_level, ch in enumerate(block_out_channels):
+            block = nn.ModuleList()
+            block_out = ch
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+
+            add_spatial_downsample = bool(i_level < np.log2(ffactor_spatial))
+            add_temporal_downsample = (add_spatial_downsample and
+                                       bool(i_level >= np.log2(ffactor_spatial // ffactor_temporal)))
+            if add_spatial_downsample or add_temporal_downsample:
+                assert i_level < len(block_out_channels) - 1
+                block_out = block_out_channels[i_level + 1] if downsample_match_channel else block_in
+                down.downsample = DownsampleDCAE(block_in, block_out, add_temporal_downsample)
+                block_in = block_out
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = Conv3d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, x: Tensor) -> Tensor:
+        use_checkpointing = bool(self.training and self.gradient_checkpointing)
+
+        # downsampling
+        h = self.conv_in(x)
+        for i_level in range(len(self.block_out_channels)):
+            for i_block in range(self.num_res_blocks):
+                h = forward_with_checkpointing(
+                    self.down[i_level].block[i_block], h, use_checkpointing=use_checkpointing)
+            if hasattr(self.down[i_level], "downsample"):
+                h = forward_with_checkpointing(self.down[i_level].downsample, h, use_checkpointing=use_checkpointing)
+
+        # middle
+        h = forward_with_checkpointing(self.mid.block_1, h, use_checkpointing=use_checkpointing)
+        h = forward_with_checkpointing(self.mid.attn_1, h, use_checkpointing=use_checkpointing)
+        h = forward_with_checkpointing(self.mid.block_2, h, use_checkpointing=use_checkpointing)
+
+        # end
+        group_size = self.block_out_channels[-1] // (2 * self.z_channels)
+        shortcut = rearrange(h, "b (c r) f h w -> b c r f h w", r=group_size).mean(dim=2)
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        h += shortcut
+        return h
+
+
+class Decoder(nn.Module):
+    """
+    The decoder network of AutoencoderKLConv3D.
+    """
+    def __init__(
+        self,
+        z_channels: int,
+        out_channels: int,
+        block_out_channels: Tuple[int, ...],
+        num_res_blocks: int,
+        ffactor_spatial: int,
+        ffactor_temporal: int,
+        upsample_match_channel: bool = True,
+    ):
+        super().__init__()
+        assert block_out_channels[0] % z_channels == 0
+
+        self.z_channels = z_channels
+        self.block_out_channels = block_out_channels
+        self.num_res_blocks = num_res_blocks
+
+        # z to block_in
+        block_in = block_out_channels[0]
+        self.conv_in = Conv3d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level, ch in enumerate(block_out_channels):
+            block = nn.ModuleList()
+            block_out = ch
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+
+            add_spatial_upsample = bool(i_level < np.log2(ffactor_spatial))
+            add_temporal_upsample = bool(i_level < np.log2(ffactor_temporal))
+            if add_spatial_upsample or add_temporal_upsample:
+                assert i_level < len(block_out_channels) - 1
+                block_out = block_out_channels[i_level + 1] if upsample_match_channel else block_in
+                up.upsample = UpsampleDCAE(block_in, block_out, add_temporal_upsample)
+                block_in = block_out
+            self.up.append(up)
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = Conv3d(block_in, out_channels, kernel_size=3, stride=1, padding=1)
+
+        self.gradient_checkpointing = False
+
+    def forward(self, z: Tensor) -> Tensor:
+        use_checkpointing = bool(self.training and self.gradient_checkpointing)
+
+        # z to block_in
+        repeats = self.block_out_channels[0] // (self.z_channels)
+        h = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1)
+
+        # middle
+        h = forward_with_checkpointing(self.mid.block_1, h, use_checkpointing=use_checkpointing)
+        h = forward_with_checkpointing(self.mid.attn_1, h, use_checkpointing=use_checkpointing)
+        h = forward_with_checkpointing(self.mid.block_2, h, use_checkpointing=use_checkpointing)
+
+        # upsampling
+        for i_level in range(len(self.block_out_channels)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = forward_with_checkpointing(self.up[i_level].block[i_block], h, use_checkpointing=use_checkpointing)
+            if hasattr(self.up[i_level], "upsample"):
+                h = forward_with_checkpointing(self.up[i_level].upsample, h, use_checkpointing=use_checkpointing)
+
+        # end
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+
+class AutoencoderKLConv3D(ModelMixin, ConfigMixin):
+    """
+    Autoencoder model with KL-regularized latent space based on 3D convolutions.
+    """
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        latent_channels: int,
+        block_out_channels: Tuple[int, ...],
+        layers_per_block: int,
+        ffactor_spatial: int,
+        ffactor_temporal: int,
+        sample_size: int,
+        sample_tsize: int,
+        scaling_factor: float = None,
+        shift_factor: Optional[float] = None,
+        downsample_match_channel: bool = True,
+        upsample_match_channel: bool = True,
+        only_encoder: bool = False,     # only build encoder for saving memory
+        only_decoder: bool = False,     # only build decoder for saving memory
+    ):
+        super().__init__()
+        self.ffactor_spatial = ffactor_spatial
+        self.ffactor_temporal = ffactor_temporal
+        self.scaling_factor = scaling_factor
+        self.shift_factor = shift_factor
+
+        # build model
+        if not only_decoder:
+            self.encoder = Encoder(
+                in_channels=in_channels,
+                z_channels=latent_channels,
+                block_out_channels=block_out_channels,
+                num_res_blocks=layers_per_block,
+                ffactor_spatial=ffactor_spatial,
+                ffactor_temporal=ffactor_temporal,
+                downsample_match_channel=downsample_match_channel,
+            )
+        if not only_encoder:
+            self.decoder = Decoder(
+                z_channels=latent_channels,
+                out_channels=out_channels,
+                block_out_channels=list(reversed(block_out_channels)),
+                num_res_blocks=layers_per_block,
+                ffactor_spatial=ffactor_spatial,
+                ffactor_temporal=ffactor_temporal,
+                upsample_match_channel=upsample_match_channel,
+            )
+
+        # slicing and tiling related
+        self.use_slicing = False
+        self.slicing_bsz = 1
+        self.use_spatial_tiling = False
+        self.use_temporal_tiling = False
+        self.use_tiling_during_training = False
+
+        # only relevant if vae tiling is enabled
+        self.tile_sample_min_size = sample_size
+        self.tile_latent_min_size = sample_size // ffactor_spatial
+        self.tile_sample_min_tsize = sample_tsize
+        self.tile_latent_min_tsize = sample_tsize // ffactor_temporal
+        self.tile_overlap_factor = 0.25
+
+        # use torch.compile for faster encode speed
+        self.use_compile = False
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (Encoder, Decoder)):
+            module.gradient_checkpointing = value
+
+    def 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):
+        """ spatial tailing for frames """
+        B, C, T, H, W = x.shape
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))  # 256 * (1 - 0.25) = 192
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)  # 8 * 0.25 = 2
+        row_limit = self.tile_latent_min_size - blend_extent  # 8 - 2 = 6
+
+        rows = []
+        for i in range(0, H, overlap_size):
+            row = []
+            for j in range(0, W, overlap_size):
+                tile = x[:, :, :, i: i + self.tile_sample_min_size, j: j + self.tile_sample_min_size]
+                tile = self.encoder(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=-1))
+        moments = torch.cat(result_rows, dim=-2)
+        return moments
+
+    def temporal_tiled_encode(self, x: torch.Tensor):
+        """ temporal tailing for frames """
+        B, C, T, H, W = x.shape
+        overlap_size = int(self.tile_sample_min_tsize * (1 - self.tile_overlap_factor))  # 64 * (1 - 0.25) = 48
+        blend_extent = int(self.tile_latent_min_tsize * self.tile_overlap_factor)  # 8 * 0.25 = 2
+        t_limit = self.tile_latent_min_tsize - blend_extent  # 8 - 2 = 6
+
+        row = []
+        for i in range(0, T, overlap_size):
+            tile = x[:, :, i: i + self.tile_sample_min_tsize, :, :]
+            if self.use_spatial_tiling and (
+                    tile.shape[-1] > self.tile_sample_min_size or tile.shape[-2] > self.tile_sample_min_size):
+                tile = self.spatial_tiled_encode(tile)
+            else:
+                tile = self.encoder(tile)
+            row.append(tile)
+        result_row = []
+        for i, tile in enumerate(row):
+            if i > 0:
+                tile = self.blend_t(row[i - 1], tile, blend_extent)
+            result_row.append(tile[:, :, :t_limit, :, :])
+        moments = torch.cat(result_row, dim=-3)
+        return moments
+
+    def spatial_tiled_decode(self, z: torch.Tensor):
+        """ spatial tailing for frames """
+        B, C, T, H, W = z.shape
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))  # 8 * (1 - 0.25) = 6
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)  # 256 * 0.25 = 64
+        row_limit = self.tile_sample_min_size - blend_extent  # 256 - 64 = 192
+
+        rows = []
+        for i in range(0, H, overlap_size):
+            row = []
+            for j in range(0, W, overlap_size):
+                tile = z[:, :, :, i: i + self.tile_latent_min_size, j: j + self.tile_latent_min_size]
+                decoded = self.decoder(tile)
+                row.append(decoded)
+            rows.append(row)
+
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=-1))
+        dec = torch.cat(result_rows, dim=-2)
+        return dec
+
+    def temporal_tiled_decode(self, z: torch.Tensor):
+        """ temporal tailing for frames """
+        B, C, T, H, W = z.shape
+        overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor))  # 8 * (1 - 0.25) = 6
+        blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor)  # 64 * 0.25 = 16
+        t_limit = self.tile_sample_min_tsize - blend_extent  # 64 - 16 = 48
+        assert 0 < overlap_size < self.tile_latent_min_tsize
+
+        row = []
+        for i in range(0, T, overlap_size):
+            tile = z[:, :, i: i + self.tile_latent_min_tsize, :, :]
+            if self.use_spatial_tiling and (
+                    tile.shape[-1] > self.tile_latent_min_size or tile.shape[-2] > self.tile_latent_min_size):
+                decoded = self.spatial_tiled_decode(tile)
+            else:
+                decoded = self.decoder(tile)
+            row.append(decoded)
+
+        result_row = []
+        for i, tile in enumerate(row):
+            if i > 0:
+                tile = self.blend_t(row[i - 1], tile, blend_extent)
+            result_row.append(tile[:, :, :t_limit, :, :])
+        dec = torch.cat(result_row, dim=-3)
+        return dec
+
+    def encode(self, x: Tensor, return_dict: bool = True):
+        """
+        Encodes the input by passing through the encoder network.
+        Support slicing and tiling for memory efficiency.
+        """
+        def _encode(x):
+            if self.use_temporal_tiling and x.shape[-3] > self.tile_sample_min_tsize:
+                return self.temporal_tiled_encode(x)
+            if self.use_spatial_tiling and (
+                    x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
+                return self.spatial_tiled_encode(x)
+
+            if self.use_compile:
+                @torch.compile
+                def encoder(x):
+                    return self.encoder(x)
+                return encoder(x)
+            return self.encoder(x)
+
+        if len(x.shape) != 5:  # (B, C, T, H, W)
+            x = x[:, :, None]
+        assert len(x.shape) == 5  # (B, C, T, H, W)
+        if x.shape[2] == 1:
+            x = x.expand(-1, -1, self.ffactor_temporal, -1, -1)
+        else:
+            assert x.shape[2] != self.ffactor_temporal and x.shape[2] % self.ffactor_temporal == 0
+
+        if self.use_slicing and x.shape[0] > 1:
+            if self.slicing_bsz == 1:
+                encoded_slices = [_encode(x_slice) for x_slice in x.split(1)]
+            else:
+                sections = [self.slicing_bsz] * (x.shape[0] // self.slicing_bsz)
+                if x.shape[0] % self.slicing_bsz != 0:
+                    sections.append(x.shape[0] % self.slicing_bsz)
+                encoded_slices = [_encode(x_slice) for x_slice in x.split(sections)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = _encode(x)
+        posterior = DiagonalGaussianDistribution(h)
+
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def decode(self, z: Tensor, return_dict: bool = True, generator=None):
+        """
+        Decodes the input by passing through the decoder network.
+        Support slicing and tiling for memory efficiency.
+        """
+        def _decode(z):
+            if self.use_temporal_tiling and z.shape[-3] > self.tile_latent_min_tsize:
+                return self.temporal_tiled_decode(z)
+            if self.use_spatial_tiling and (
+                    z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
+                return self.spatial_tiled_decode(z)
+            return self.decoder(z)
+
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [_decode(z_slice) for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = _decode(z)
+
+        if z.shape[-3] == 1:
+            decoded = decoded[:, :, -1:]
+
+        if not return_dict:
+            return (decoded,)
+
+        return DecoderOutput(sample=decoded)
+
+    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
+
+        # Use fixed shape here
+        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()

configuration_hunyuan.py

@@ -0,0 +1,285 @@
+# 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
+
+
+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=290943,
+            hidden_size=4096,
+            intermediate_size: int=11008,
+            moe_intermediate_size: Union[int, List]=None,
+            num_hidden_layers=32,
+            num_attention_heads=32,
+            num_key_value_heads=None,
+            attention_head_dim=None,
+            hidden_act="silu",
+            max_position_embeddings=2048,
+            initializer_range=0.02,
+            rms_norm_eps=1e-5,
+            use_cache=True,
+            pad_token_id=0,
+            bos_token_id=1,
+            eos_token_id=2,
+            eod_token_id=3,
+            im_start_id=4,
+            im_end_id=5,
+            text_start_id=6,
+            text_end_id=7,
+            image_token_id=8,
+            video_start_id=9,
+            video_end_id=10,
+            im_newline_id=11,
+            mask_init_id=12,
+            pretraining_tp=1,
+            tie_word_embeddings=False,
+            rope_theta=10000.0,
+            rope_scaling=None,
+            attention_bias=False,
+            mlp_bias=False,
+            attention_dropout=0.0,
+            use_qk_norm=False,
+            use_rotary_pos_emb=True,
+            use_cla=False,
+            cla_share_factor=1,
+            norm_type="hf_rms",
+            num_experts: 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,
+            vae=None,
+            vit=None,
+            vit_processor=None,
+            vit_aligner=None,
+            **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.moe_intermediate_size = moe_intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.moe_impl = moe_impl
+        self.num_experts = num_experts
+        self.use_mixed_mlp_moe = use_mixed_mlp_moe
+        self.num_shared_expert = num_shared_expert
+        self.moe_topk = moe_topk
+        self.capacity_factor = capacity_factor
+        self.moe_drop_tokens = moe_drop_tokens
+        self.moe_random_routing_dropped_token = moe_random_routing_dropped_token
+
+        if attention_head_dim is not None:
+            self.attention_head_dim = attention_head_dim
+        else:
+            self.attention_head_dim = self.hidden_size // num_attention_heads
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.pretraining_tp = pretraining_tp
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.mlp_bias = mlp_bias
+        self.attention_dropout = attention_dropout
+        self.use_qk_norm = use_qk_norm
+        self.use_rotary_pos_emb = use_rotary_pos_emb
+        self.use_cla = use_cla
+        self.cla_share_factor = cla_share_factor
+        self.norm_type = norm_type
+        # MLA args
+        self.use_mla = use_mla
+        self.kv_lora_rank = kv_lora_rank
+        self.q_lora_rank = q_lora_rank
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.v_head_dim = v_head_dim
+
+        # DeepSeek related args
+        self.moe_layer_num_skipped = moe_layer_num_skipped
+        self.norm_topk_prob = norm_topk_prob
+        self.routed_scaling_factor = routed_scaling_factor
+        self.group_limited_greedy = group_limited_greedy
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.add_classification_head = add_classification_head
+        self.class_num = class_num
+        self.pool_type = pool_type
+        self.pad_id = pad_id
+
+        if self.class_num is not None:
+            self.dense_list = [self.hidden_size, self.class_num]
+
+        # ViT args
+        self.vit = vit
+        self.vit_processor = vit_processor
+        self.vit_aligner = vit_aligner
+
+        # Image Gen args
+        self.vae = vae
+        self.vae_downsample_factor = vae_downsample_factor
+        self.img_proj_type = img_proj_type
+        self.patch_size = patch_size
+        self.patch_embed_hidden_dim = patch_embed_hidden_dim
+        self.image_base_size = image_base_size
+
+        # token id
+        self.eod_token_id = eod_token_id
+        self.im_start_id = im_start_id
+        self.im_end_id = im_end_id
+        self.text_start_id = text_start_id
+        self.text_end_id = text_end_id
+        self.image_token_id = image_token_id
+        self.video_start_id = video_start_id
+        self.video_end_id = video_end_id
+        self.im_newline_id = im_newline_id
+        self.mask_init_id = mask_init_id
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

hunyuan_image_3_pipeline.py

@@ -0,0 +1,879 @@
+# 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
+
+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 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,
+        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
+
+        cfg_factor = 1 + self.do_classifier_free_guidance
+
+        # Define call parameters
+        device = self._execution_device
+
+        # Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, timesteps, sigmas,
+        )
+
+        # Prepare latent variables
+        latents = self.prepare_latents(
+            batch_size=batch_size,
+            latent_channel=self.model.config.vae["latent_channels"],
+            image_size=image_size,
+            dtype=torch.bfloat16,
+            device=device,
+            generator=generator,
+            latents=latents,
+        )
+
+        # Prepare extra step kwargs.
+        _scheduler_step_extra_kwargs = self.prepare_extra_func_kwargs(
+            self.scheduler.step, {"generator": generator}
+        )
+
+        # Prepare model kwargs
+        input_ids = model_kwargs.pop("input_ids")
+        attention_mask = self.model._prepare_attention_mask_for_generation(     # noqa
+            input_ids, self.model.generation_config, model_kwargs=model_kwargs,
+        )
+        model_kwargs["attention_mask"] = attention_mask.to(latents.device)
+
+        # Sampling loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * cfg_factor)
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                t_expand = t.repeat(latent_model_input.shape[0])
+
+                model_inputs = self.model.prepare_inputs_for_generation(
+                    input_ids,
+                    images=latent_model_input,
+                    timestep=t_expand,
+                    **model_kwargs,
+                )
+
+                with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+                    model_output = self.model(**model_inputs, first_step=(i == 0))
+                    pred = model_output["diffusion_prediction"]
+                pred = pred.to(dtype=torch.float32)
+
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    pred_cond, pred_uncond = pred.chunk(2)
+                    pred = self.cfg_operator(pred_cond, pred_uncond, self.guidance_scale, step=i)
+
+                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,
+                    )
+                    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,125 @@
+# 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 typing import Tuple
+
+from PIL import Image
+from torchvision import transforms
+from transformers import Siglip2ImageProcessorFast
+
+from .tokenizer_wrapper import ImageInfo, JointImageInfo, ResolutionGroup
+
+
+def resize_and_crop(image: Image.Image, target_size: Tuple[int, int]) -> Image.Image:
+    tw, th = target_size
+    w, h = image.size
+
+    tr = th / tw
+    r = h / w
+
+    # resize
+    if r < tr:
+        resize_height = th
+        resize_width = int(round(th / h * w))
+    else:
+        resize_width = tw
+        resize_height = int(round(tw / w * h))
+
+    image = image.resize((resize_width, resize_height), resample=Image.Resampling.LANCZOS)
+
+    # center crop
+    crop_top = int(round((resize_height - th) / 2.0))
+    crop_left = int(round((resize_width - tw) / 2.0))
+
+    image = image.crop((crop_left, crop_top, crop_left + tw, crop_top + th))
+    return image
+
+
+class HunyuanImage3ImageProcessor(object):
+    def __init__(self, config):
+        self.config = config
+
+        self.reso_group = ResolutionGroup(base_size=config.image_base_size)
+        self.vae_processor = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5]),  # transform to [-1, 1]
+        ])
+        self.vision_encoder_processor = Siglip2ImageProcessorFast.from_dict(config.vit_processor)
+
+    def build_image_info(self, image_size):
+        # parse image size (HxW, H:W, or <img_ratio_i>)
+        if isinstance(image_size, str):
+            if image_size.startswith("<img_ratio_"):
+                ratio_index = int(image_size.split("_")[-1].rstrip(">"))
+                reso = self.reso_group[ratio_index]
+                image_size = reso.height, reso.width
+            elif 'x' in image_size:
+                image_size = [int(s) for s in image_size.split('x')]
+            elif ':' in image_size:
+                image_size = [int(s) for s in image_size.split(':')]
+            else:
+                raise ValueError(
+                    f"`image_size` should be in the format of 'HxW', 'H:W' or <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.reso_group.get_target_size(image_size[1], image_size[0])
+        token_height = image_height // (self.config.vae_downsample_factor[0] * self.config.patch_size)
+        token_width = image_width // (self.config.vae_downsample_factor[1] * self.config.patch_size)
+        base_size, ratio_idx = self.reso_group.get_base_size_and_ratio_index(image_size[1], image_size[0])
+        image_info = ImageInfo(
+            image_type="gen_image", image_width=image_width, image_height=image_height,
+            token_width=token_width, token_height=token_height, base_size=base_size, ratio_index=ratio_idx,
+        )
+        return image_info
+
+    def preprocess(self, image: Image.Image):
+        # ==== VAE processor ====
+        image_width, image_height = self.reso_group.get_target_size(image.width, image.height)
+        resized_image = resize_and_crop(image, (image_width, image_height))
+        image_tensor = self.vae_processor(resized_image)
+        token_height = image_height // (self.config.vae_downsample_factor[0] * self.config.patch_size)
+        token_width = image_width // (self.config.vae_downsample_factor[1] * self.config.patch_size)
+        base_size, ratio_index = self.reso_group.get_base_size_and_ratio_index(width=image_width, height=image_height)
+        vae_image_info = ImageInfo(
+            image_type="vae",
+            image_tensor=image_tensor.unsqueeze(0),     # include batch dim
+            image_width=image_width, image_height=image_height,
+            token_width=token_width, token_height=token_height,
+            base_size=base_size, ratio_index=ratio_index,
+        )
+
+        # ==== ViT processor ====
+        inputs = self.vision_encoder_processor(image)
+        image = inputs["pixel_values"].squeeze(0)  # seq_len x dim
+        pixel_attention_mask = inputs["pixel_attention_mask"].squeeze(0)  # seq_len
+        spatial_shapes = inputs["spatial_shapes"].squeeze(0)  # 2  (h, w)
+        vision_encoder_kwargs = dict(
+            pixel_attention_mask=pixel_attention_mask,
+            spatial_shapes=spatial_shapes,
+        )
+        vision_image_info = ImageInfo(
+            image_type="vit",
+            image_tensor=image.unsqueeze(0),  # 1 x seq_len x dim
+            image_width=spatial_shapes[1].item() * self.config.vit_processor["patch_size"],
+            image_height=spatial_shapes[0].item() * self.config.vit_processor["patch_size"],
+            token_width=spatial_shapes[1].item(),
+            token_height=spatial_shapes[0].item(),
+            image_token_length=self.config.vit_processor["max_num_patches"],
+            # may not equal to token_width * token_height
+        )
+        return JointImageInfo(vae_image_info, vision_image_info, vision_encoder_kwargs)

siglip2.py

@@ -0,0 +1,564 @@
+# 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)

system_prompt.py

@@ -0,0 +1,128 @@
+# 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]
+}
+
+
+def get_system_prompt(sys_type, bot_task, system_prompt=None):
+    if sys_type == 'None':
+        return None
+    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}")

tokenizer_wrapper.py

@@ -0,0 +1,1425 @@
+# 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 warnings
+import random
+from typing import List, Optional, Union, Dict, Any
+from collections import defaultdict
+from copy import deepcopy
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+from diffusers.utils import BaseOutput
+
+
+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):
+        self.align = align
+        self.base_size = base_size
+        assert base_size % align == 0, f'base_size {base_size} is not divisible by align {align}'
+        if base_size is not None and not isinstance(base_size, int):
+            raise ValueError(f'base_size must be None or int, but got {type(base_size)}')
+        if step is None:
+            step = base_size // 16
+        if step is not None and step > base_size // 2:
+            raise ValueError(f'step must be smaller than base_size // 2, but got {step} > {base_size // 2}')
+
+        self.step = step
+        self.data = self._calc_by_step()
+
+        self.ratio = np.array([x.ratio for x in self.data])
+        self.attr = ['' for _ in range(len(self.data))]
+        self.prefix_space = 0
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+    def __repr__(self):
+        prefix = self.prefix_space * ' '
+        prefix_close = (self.prefix_space - 4) * ' '
+        res_str = f'ResolutionGroup(base_size={self.base_size}, step={self.step}, data='
+        attr_maxlen = max([len(x) for x in self.attr] + [5])
+        res_str += \
+            f'\n{prefix}ID: height width   ratio {" " * max(0, attr_maxlen - 4)}count  h/16 w/16    tokens\n{prefix}'
+        res_str += \
+            ('\n' + prefix).join([f'{i:2d}: ({x.h:4d}, {x.w:4d})  {self.ratio[i]:.4f}  {self.attr[i]:>{attr_maxlen}s}  '
+                                  f'({x.h // 16:3d}, {x.w // 16:3d})  {x.h // 16 * x.w // 16:6d}'
+                                  for i, x in enumerate(self.data)])
+        res_str += f'\n{prefix_close})'
+        return res_str
+
+    def _calc_by_step(self):
+        assert self.align <= self.step, f'align {self.align} must be smaller than step {self.step}'
+
+        min_height = self.base_size // 2
+        min_width = self.base_size // 2
+        max_height = self.base_size * 2
+        max_width = self.base_size * 2
+
+        resolutions = [Resolution(self.base_size, self.base_size)]
+
+        cur_height, cur_width = self.base_size, self.base_size
+        while True:
+            if cur_height >= max_height and cur_width <= min_width:
+                break
+
+            cur_height = min(cur_height + self.step, max_height)
+            cur_width = max(cur_width - self.step, min_width)
+            resolutions.append(Resolution(cur_height // self.align * self.align, cur_width // self.align * self.align))
+
+        cur_height, cur_width = self.base_size, self.base_size
+        while True:
+            if cur_height <= min_height and cur_width >= max_width:
+                break
+
+            cur_height = max(cur_height - self.step, min_height)
+            cur_width = min(cur_width + self.step, max_width)
+            resolutions.append(Resolution(cur_height // self.align * self.align, cur_width // self.align * self.align))
+
+        resolutions = sorted(resolutions, key=lambda x: x.ratio)
+
+        return resolutions
+
+    def get_target_size(self, width, height):
+        ratio = height / width
+        idx = np.argmin(np.abs(self.ratio - ratio))
+        reso = self.data[idx]
+        return reso.w, reso.h
+
+    def get_base_size_and_ratio_index(self, width, height):
+        ratio = height / width
+        idx = np.argmin(np.abs(self.ratio - ratio))
+        return self.base_size, idx
+
+
+class ImageInfo:
+    """ Class to store image information for processing and generation. """
+
+    def __init__(
+            self,
+            image_type: str = None,
+            image_tensor: torch.Tensor = None,
+            image_width: int = None,
+            image_height: int = None,
+            token_width: int = None,
+            token_height: int = None,
+            image_token_length: int = None,
+            base_size: int = None,
+            ratio_index: int = None,
+            **kwargs,
+    ):
+        self.image_type = image_type
+        self.image_tensor = image_tensor
+        self.image_width = image_width
+        self.w = image_width
+        self.image_height = image_height
+        self.h = image_height
+        self.token_width = token_width
+        self.tk_w = token_width
+        self.token_height = token_height
+        self.tk_h = token_height
+        self.image_token_length = default(
+            image_token_length,
+            token_width * token_height if token_width is not None and token_height is not None else None
+        )
+        self.base_size = base_size
+        self.ratio_index = ratio_index
+
+        self.add_timestep_token = kwargs.get("add_timestep_token", True)
+        self.add_guidance_token = kwargs.get("add_guidance_token", False)
+        self.use_front_boi_token = kwargs.get("use_front_boi_token", True)
+        self.add_image_shape_token = kwargs.get("add_image_shape_token", True)
+
+    def __getitem__(self, key: str) -> Any:
+        """Allow dictionary-like access to attributes."""
+        if hasattr(self, key):
+            return getattr(self, key)
+        raise KeyError(f"Key '{key}' not found in ImageInfo")
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        """Allow dictionary-like assignment to attributes."""
+        if hasattr(self, key):
+            setattr(self, key, value)
+        else:
+            raise KeyError(f"Key '{key}' not found in ImageInfo")
+
+    def __contains__(self, key: str) -> bool:
+        """Check if the key exists in the ImageInfo object."""
+        return hasattr(self, key)
+
+    def __repr__(self):
+        return (f"ImageInfo(image_type={self.image_type}, image_tensor={self.image_tensor}, "
+                f"image_width={self.image_width}, image_height={self.image_height}, "
+                f"token_width={self.token_width}, token_height={self.token_height}, "
+                f"image_token_length={self.image_token_length}, "
+                f"base_size={self.base_size}, ratio_index={self.ratio_index}")
+
+    @property
+    def meta_info(self):
+        # Used for image sections of tkwrapper.encode_general()
+        if self.image_type in ["vae", "gen_image"]:
+            return dict(
+                token_length=self.image_token_length,
+                add_timestep_token=self.add_timestep_token,
+                add_guidance_token=self.add_guidance_token,
+                use_front_boi_token=self.use_front_boi_token,
+                add_image_shape_token=self.add_image_shape_token,
+                base_size=self.base_size,
+                ratio_idx=self.ratio_index,
+                # for rope 2d
+                token_height=self.token_height,
+                token_width=self.token_width,
+                # for bc
+                image_height=self.image_height,
+                image_width=self.image_width,
+            )
+        elif self.image_type in ["vit"]:
+            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,
+            )
+        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) +
+                    (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,
+            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 JointImage(object):
+    def __init__(self, vae_image: ImageTensor, vision_image: ImageTensor):
+        self.vae_image = vae_image
+        self.vision_image = vision_image
+        self.i = JointImageInfo(vae_image.i, vision_image.i)
+
+
+class TokenizerEncodeOutput(BaseOutput):
+    tokens: torch.Tensor = None
+    timestep_scatter_index: Optional[torch.Tensor] = None
+    guidance_scatter_index: Optional[torch.Tensor] = None
+    text_slices: Optional[List[slice]] = None
+    gen_image_slices: Optional[List[slice]] = None
+    joint_image_slices: Optional[List[slice]] = None
+    cond_vae_image_slices: Optional[List[slice]] = None
+    cond_vit_image_slices: Optional[List[slice]] = None
+    text_mask: Optional[torch.Tensor] = None
+    gen_image_mask: Optional[torch.Tensor] = None
+    cond_vae_image_mask: Optional[torch.Tensor] = None
+    cond_vit_image_mask: Optional[torch.Tensor] = None
+    real_pos: Optional[torch.Tensor] = None
+    all_image_slices: Optional[List[slice]] = None
+    cond_timestep_scatter_index: Optional[torch.Tensor] = None
+    gen_timestep_scatter_index: Optional[torch.Tensor] = None
+
+
+class Conversation:
+    roles: List[str] = ["User", "Assistant"]
+    sep: str = "\n\n"
+
+
+class TokenizerWrapper(object):
+    def __init__(self, tokenizer):
+        if isinstance(tokenizer, str):
+            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)
+        else:
+            self.tokenizer = tokenizer
+
+        # Define short names
+        self.bos_token_id = self.tokenizer.bos_token_id
+        self.eos_token_id = self.tokenizer.eos_token_id
+        self.pad_token_id = self.tokenizer.pad_token_id
+        self.boi_token_id = self.tokenizer.convert_tokens_to_ids("<boi>")
+        self.eoi_token_id = self.tokenizer.convert_tokens_to_ids("<eoi>")
+        self.img_token_id = self.tokenizer.convert_tokens_to_ids("<img>")
+        self.cfg_token_id = self.tokenizer.convert_tokens_to_ids("<cfg>")
+        self.end_answer_token_id = self.tokenizer.convert_tokens_to_ids("</answer>")
+        self.end_recaption_token_id = self.tokenizer.convert_tokens_to_ids("</recaption>")
+        self.ratio_token_offset = self.tokenizer.convert_tokens_to_ids("<img_ratio_0>")
+        self.special_token_map = self.tokenizer.added_tokens_encoder
+
+    def pad(self, tensor_list, dim=0, pad_val=None):
+        if pad_val is None:
+            pad_val = self.pad_token_id
+        max_len = max([t.shape[dim] for t in tensor_list])
+        padded_tensor_list = []
+        for t in tensor_list:
+            if t.shape[dim] < max_len:
+                assert pad_val is not False, "Not allowed pad."
+                t = F.pad(t, (0, max_len - t.shape[dim]), value=pad_val)
+            padded_tensor_list.append(t)
+        return padded_tensor_list
+
+    def encode(self, *args, **kwargs):
+        return self.tokenizer.encode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        return self.tokenizer.decode(*args, **kwargs)
+
+    def encode_text(
+            self,
+            *texts,
+            uncond_enabled: Optional[Union[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:
+                warnings.warn(
+                    f"Text length exceeds the max_length({max_length}). The remaining texts will be ignored: "
+                    f"{text[:80]}..."
+                )
+                break
+            # Set add_special_tokens=False to avoid adding <bos> token in some LLMs.
+            if isinstance(text, str):
+                text_token = self.tokenizer.encode(text, add_special_tokens=False)
+            else:
+                text_token = text
+            if uncond_flag and do_uncond_drop:
+                text_token = [self.cfg_token_id] * len(text_token)
+            # Cutoff the text by max_length if necessary
+            if max_length is not None and (cum_length + len(text_token)) > max_length:
+                text_token = text_token[:max_length - cum_length]
+            text_tokens.extend(text_token)
+            lengths.append(len(text_token))
+            cum_length += len(text_token)
+
+        # Prepend/Append <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 {set(keys)} and {list(data.keys())}."
+        )
+        key_counts = {k: 0 for k in keys}
+        for key in keys:
+            key_counts[key] += 1
+        for key, count in key_counts.items():
+            assert len(data[key]) == count, (
+                f"Number of `{key}` in the token source should be matched with the template, but got "
+                f"{data[key]}({len(data[key])}) and {count}."
+            )
+
+    def _add_image_meta_info_token(self, token_seq, token_count, extra_token_pos, add_timestep_token=False,
+                                   add_image_shape_token=False, base_size=None, ratio_idx=None, image_type=None,
+                                   add_guidance_token=False):
+        if add_image_shape_token:
+            token_seq.extend([
+                self.special_token_map[f"<img_size_{base_size}>"],
+                self.special_token_map[f"<img_ratio_{ratio_idx}>"]
+            ])
+            token_count += 2
+        if add_timestep_token:
+            token_seq.extend([self.special_token_map["<timestep>"]])
+            extra_token_pos['timestep'].append(token_count)
+            if image_type is not None:
+                if image_type == "gen_image":
+                    extra_token_pos['gen_timestep'].append(token_count)
+                elif image_type in ["joint_image"]:
+                    extra_token_pos['cond_timestep'].append(token_count)
+                else:
+                    raise ValueError(f"Unsupported image type: {image_type}.")
+            token_count += 1
+        if add_guidance_token:
+            token_seq.extend([self.special_token_map["<guidance>"]])
+            extra_token_pos['guidance'].append(token_count)
+            token_count += 1
+        return token_count
+
+    @staticmethod
+    def _shorten_text(text):
+        import re
+        text = re.sub(r"(<img>)+", lambda m: f"[<img>]{{{len(m.group(0)) // 5}}}", text)
+        text = re.sub(r"(<pad>)+", lambda m: f"[<pad>]{{{len(m.group(0)) // 5}}}", text)
+        return text
+
+    def encode_sequence(
+            self,
+            template: str,
+            token_source: Dict[str, List],
+            total_length=None,
+            add_timestep_token=False,
+            add_guidance_token=False,
+            last_key_only_prefix=False,
+            add_eos=True,
+            use_front_boi_token=True,
+            add_pad=True,
+            add_bos=True,
+            drop_last: Union[str, bool] = 'auto',
+            add_image_shape_token=False,
+    ):
+        """
+        Encode a sequence based on the template (e.g., `text-image` for t2i, `text-image-image` for instruction tuning)
+        and token source.
+
+        Parameters
+        ----------
+        template: str
+            Template of the sequence. E.g., "text-gen_image" means the sequence is composed of text and an image.
+            "text-text-gen_image" means the sequence is composed of two sections of text and an image.
+        token_source: Dict[str, List]
+            Token source for each key in the template, in order.
+            - text: List[Dict].
+            - gen_image: List[Dict].
+            - joint_image: List[Dict].
+        total_length: int
+            Total length of the encoded sequence, include padding tokens.
+        add_timestep_token: bool
+            Whether to add timestep token before the image tokens.
+            (Right after the <img_ratio_*><img_size_*> tokens)
+        add_guidance_token: bool
+            Whether to add guidance token before the image tokens.
+        last_key_only_prefix: bool
+            Whether to only use the modal prefix in the last key.
+        add_eos: bool or 'auto'
+            Whether to add eos token at the end of the sequence. If True, always add eos token. If 'auto',
+            add eos token only when the total_length is not reached and the last token is not <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.
+        """
+        if last_key_only_prefix:
+            assert add_eos is not True, "add_eos should not be True when last_key_only_prefix is True."
+        if drop_last is True and total_length is None:
+            raise ValueError("total_length should be provided when drop_last is True.")
+
+        keys = template.split('-')
+        modal_length = len(keys)
+        index_indicator = {k: 0 for k in token_source}
+        for k, v in token_source.items():
+            assert isinstance(v, (list, tuple)), (
+                f"Value of `{k}` in the token source should be a list or tuple, but got {type(v)}."
+            )
+        self._check_key_number_matched(keys, token_source)
+
+        token_seq = []
+        token_count = 0
+        extra_token_pos = defaultdict(list)
+        if add_bos:
+            token_seq.append(self.bos_token_id)
+            token_count += 1
+        # If drop_last is True, we check the token_count on the fly and exit the loop if the total_length is reached.
+        # This check is only applied to the block tokens. Block tokens mean the tokens that are unsplittable, like
+        # image tokens. Text tokens are splittable, so we don't need to check the token_count for text.
+        # If the loop is broken by drop_last, we don't add the eos token at the end because the sequence is not
+        # complete.
+        drop_last_break = False
+        for i, key in enumerate(keys):
+            source = token_source[key][index_indicator[key]]
+            if key == "text":
+                token_seq.extend(source)  # text token sequence
+                extra_token_pos["<text>_start"].append(token_count)
+                token_count += len(source)
+                extra_token_pos["<text>_end"].append(token_count - 1)
+
+            elif key == "gen_image":
+                if isinstance(source, int):
+                    source = {'length': source}
+                extra_count = 2 + (
+                    1 if source.get('timestep', add_timestep_token) else 0) + (
+                    1 if source.get('guidance', add_guidance_token) else 0) + (
+                    2 if source.get('image_shape', add_image_shape_token) else 0
+                )
+                if drop_last is True and token_count + extra_count + source['length'] > total_length:
+                    drop_last_break = True
+                    break
+                if source.get('front_boi', use_front_boi_token):
+                    token_seq.append(self.boi_token_id)
+                    extra_token_pos["boi"].append(token_count)
+                    token_count += 1
+                token_count = self._add_image_meta_info_token(
+                    token_seq=token_seq,
+                    token_count=token_count,
+                    extra_token_pos=extra_token_pos,
+                    add_timestep_token=source.get('timestep', add_timestep_token),
+                    add_guidance_token=source.get('guidance', add_guidance_token),
+                    add_image_shape_token=source.get('image_shape', add_image_shape_token),
+                    base_size=source.get('base_size'),
+                    ratio_idx=source.get('ratio_idx'),
+                    image_type=key,
+                )
+                if not source.get('front_boi', use_front_boi_token):
+                    token_seq.append(self.boi_token_id)
+                    extra_token_pos["boi"].append(token_count)
+                    token_count += 1
+                if last_key_only_prefix and i == modal_length - 1:
+                    pass  # for AR inference
+                else:
+                    token_seq.extend(
+                        [self.img_token_id] * source['length'] +  # token number
+                        [self.eoi_token_id]
+                    )
+                    extra_token_pos["<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 == "joint_image":
+                assert isinstance(source['length'], list) and len(
+                    source['length']) == 2, "joint_image length should be a list of two integers"
+                extra_count = 2 + 1 + (  # boi, eoi, joint_img_sep
+                    1 if source.get('timestep', add_timestep_token) else 0) + (
+                    2 if source.get('image_shape', add_image_shape_token) else 0
+                )
+                if drop_last is True and token_count + extra_count + sum(source['length']) > total_length:
+                    drop_last_break = True
+                    break
+                if source.get('front_boi', use_front_boi_token):
+                    token_seq.append(self.boi_token_id)  # Use patched boi for Janus, otherwise useing default <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,
+                )
+                if not source.get('front_boi', use_front_boi_token):
+                    token_seq.append(self.boi_token_id)
+                    extra_token_pos["boi"].append(token_count)
+                    token_count += 1
+                if last_key_only_prefix and i == modal_length - 1:
+                    pass  # for AR inference
+                else:
+                    token_seq.extend(
+                        [self.img_token_id] * source['length'][0]
+                    )
+                    extra_token_pos["<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.special_token_map["<joint_img_sep>"]]
+                    )
+                    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>
+
+            else:
+                raise ValueError(f"Not supported key: {key}")
+            index_indicator[key] += 1
+
+        if add_eos is True and not drop_last_break:
+            # Typically used for t2i task.
+            token_seq.append(self.eos_token_id)
+            extra_token_pos["eos"].append(token_count)
+            token_count += 1
+        elif add_eos == 'auto' and not drop_last_break:
+            # Typically used for lm and mmu task.
+            if token_seq[-1] != self.eos_token_id and (total_length is None or token_count < total_length):
+                token_seq.append(self.eos_token_id)
+                extra_token_pos["eos"].append(token_count)
+                token_count += 1
+
+        if total_length:
+            # Check token count and clip sequence if necessary
+            if token_count > total_length and drop_last:
+                # Assert clip position is not in the middle of the block-wise tokens (gen_image, joint_image)
+                for start_key, end_key in [
+                    ("<img>_start", "<img>_end"), ("<joint_img>_start", "<joint_img>_end"),
+                    ("<vae_img>_start", "<vae_img>_end"), ("<vit_img>_start", "<vit_img>_end"),
+                ]:
+                    if start_key in extra_token_pos and end_key in extra_token_pos:
+                        assert all(
+                            (start > total_length or end + 1 < total_length)
+                            for start, end in zip(extra_token_pos[start_key], extra_token_pos[end_key])
+                        ), ("Clip position should not be in the middle of the image tokens.\n"
+                            f"Below is the text:\n{self._shorten_text(self.tokenizer.decode(token_seq))}")
+                token_seq = token_seq[:total_length]
+
+            # Pad the sequence if necessary
+            pad_num = max(0, total_length - len(token_seq))
+            if add_pad and pad_num:
+                token_seq.extend([self.pad_token_id] * pad_num)
+                extra_token_pos["first_pad"].append(token_count)
+
+        return token_seq, extra_token_pos
+
+    def batch_gen_infer(
+            self,
+            infer_fn,
+            prompt_list: list,
+            negative_prompt_list: list = None,
+            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
+            Support multi-condition.
+        uncondition_repeat_times: int
+            Support multi-uncondition.
+        """
+        if infer_fn_kwargs_list is None:
+            infer_fn_kwargs_list = [{} for _ in prompt_list]
+
+        # [n_output, bsz]
+        cond_results_list = None
+        uncond_results_list = None
+        output_type_list = []
+
+        for prompt_idx, (prompt, infer_fn_kwargs) in enumerate(zip(prompt_list, infer_fn_kwargs_list)):
+            if not isinstance(prompt, (list, tuple)):
+                prompt = [prompt]
+            cond_kwargs = {"uncond_p": 0.0} if do_classifier_free_guidance else {}
+            results = infer_fn(
+                *prompt,
+                **infer_fn_kwargs,
+                **cond_kwargs,
+            )
+            output_type_list.append((type(results), len(results) if isinstance(results, (list, tuple)) else 1))
+            if isinstance(results, dict):
+                raise ValueError("Make batch on dict is not supported. Please return list or tuple for infer_fn.")
+            if not isinstance(results, (list, tuple)):
+                results = (results,)
+            if cond_results_list is None:
+                cond_results_list = [[] for _ in results]
+                uncond_results_list = [[] for _ in results]
+            for i, result in enumerate(results):
+                cond_results_list[i].append(result)
+
+            if do_classifier_free_guidance:
+                if negative_prompt_list is None:
+                    uncond_kwargs = {"uncond_p": 1.0}
+                    uncond_results = infer_fn(
+                        *prompt,
+                        **infer_fn_kwargs,
+                        **uncond_kwargs,
+                    )
+                else:
+                    negative_prompt = negative_prompt_list[prompt_idx]
+                    if not isinstance(negative_prompt, (list, tuple)):
+                        negative_prompt = [negative_prompt]
+                    uncond_results = infer_fn(
+                        *negative_prompt,
+                        **infer_fn_kwargs,
+                    )
+                if isinstance(uncond_results, TokenizerEncodeOutput):
+                    uncond_results_list.append(uncond_results)
+                else:
+                    for i, result in enumerate(uncond_results):
+                        uncond_results_list[i].append(result)
+
+        assert all(output_type_list[0] == n for n in output_type_list), \
+            f"Number of outputs should be equal for all samples, but got {output_type_list}."
+        output_type, output_num = output_type_list[0]
+
+        def make_batch(batch_cond_item, batch_uncond_item):
+            # Process each output item to make batch
+            first = batch_cond_item[0]  # The first element in the batch
+            if isinstance(first, torch.Tensor):
+                stacked_item = torch.stack(self.pad(
+                    batch_cond_item * condition_repeat_times + batch_uncond_item * uncondition_repeat_times,
+                ))
+
+            elif first is None:
+                assert all(item is None for item in batch_cond_item + batch_uncond_item), \
+                    (f"The first cond item is None, but some items are not None:\n\n"
+                     f"condition: {batch_cond_item}\n\n"
+                     f"uncondition: {batch_uncond_item}")
+                stacked_item = None
+
+            elif isinstance(first, (list, tuple)):
+                # If the output item is a list or tuple, we treat it as a whole, and won't make nested batch any more.
+                stacked_item = batch_cond_item * condition_repeat_times + batch_uncond_item * uncondition_repeat_times
+
+            elif isinstance(first, TokenizerEncodeOutput):
+                stacked_item = {}
+                # Traverse not-None attributes
+                for key in list(first.keys()):
+                    merged_list = [cond_item[key] for cond_item in batch_cond_item] * condition_repeat_times + \
+                                  [uncond_item[key] for uncond_item in batch_uncond_item] * uncondition_repeat_times
+                    if isinstance(first[key], torch.Tensor):
+                        if 'mask' in key:
+                            pad_val = 0.0
+                        elif key == 'tokens':
+                            pad_val = self.special_token_map["<pad>"]
+                        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}.")
+
+    @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='auto',
+            use_text_mask=True,
+            add_pad='auto',
+            add_bos=True,
+            drop_last='auto',
+    ):
+        """
+        General encode function to encode a sequence with multiple sections of text and images.
+        Each section is a dict with a `type` key and other keys depending on the type.
+        Supported section types:
+        - text: dict with keys:
+            - text: str or List[int], text to be encoded. Either `text` or `tokens` should be provided.
+            - tokens: List[int], pre-encoded text tokens. Either `text` or `tokens` should be provided.
+            - uncond_enabled: bool, whether to enable uncondition for this text section.
+            - uncond_p: float, probability to drop the text section for uncondition.
+            - max_length: int, maximum length of the text section.
+            - ignore: bool, whether to ignore this text section in the text mask.
+            - start_offset: int, start offset of the text mask.
+            - end_offset: int, end offset of the text mask.
+        - gen_image: dict with keys:
+            - token_length: int, number of image tokens.
+            - add_timestep_token: bool, whether to add timestep token before the image tokens.
+            - add_guidance_token: bool, whether to add guidance token before the image tokens.
+            - use_front_boi_token: bool, whether to put the <boi> token at the front of size, ratio and timestep tokens.
+            - add_image_shape_token: bool, whether to add image shape token before the image tokens.
+            - base_size: int, base size of the image.
+            - ratio_idx: int, ratio index of the image.
+        - joint_image: dict with keys:
+            - token_length: List[int], number of image tokens for the two images.
+            - add_timestep_token: bool, whether to add timestep token before the image tokens.
+            - use_front_boi_token: bool, whether to put the <boi> token at the front of size, ratio and timestep tokens.
+            - add_image_shape_token: bool, whether to add image shape token before the image tokens.
+            - base_size: int, base size of the image.
+            - ratio_idx: int, ratio index of the image.
+
+        Parameters
+        ----------
+        sections: List[Dict[str, Any]]
+            List of sections to be encoded.
+        max_token_length: int
+            Maximum length of the encoded token sequence.
+        add_eos: bool or 'auto'
+            Whether to add eos token at the end of the sequence. 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_text_mask: bool
+            Whether to generate text mask.
+        add_pad: bool or 'auto'
+            Whether to add padding tokens to the sequence. If True and total_length is not reached,
+            add padding tokens.
+        add_bos: bool
+            Whether to add bos token at the beginning of the sequence.
+        drop_last: bool or 'auto'
+            - If auto, drop last tokens exceeding the total_length if the total_length is provided.
+            If cut point is in the middle of the image tokens, an error will raised.
+            - If True, drop last tokens exceeding the total_length. If cut point is in the
+            middle of the image tokens, all the successive image tokens will be dropped.
+            - If False, keep the last tokens exceeding the total_length, even if the total_length
+            is reached.
+
+        Returns
+        -------
+        TokenizerEncodeOutput
+            Encoded token sequence and extra information.
+        """
+        if sections is None:
+            raise ValueError("sections must be provided.")
+        template = '-'.join([section['type'] for section in sections])
+
+        sections = deepcopy(sections)
+        token_source = defaultdict(list)
+        text_mask_specs = []
+        for section in sections:
+            if section['type'] == 'text':
+                text = self.encode_text(
+                    section['text'] if 'text' in section else section['tokens'],
+                    uncond_enabled=section.get('uncond_enabled'),
+                    uncond_p=section.get('uncond_p'),
+                    max_length=section.get('max_length'),
+                )
+                token_source['text'].append(text)
+                text_mask_specs.append(dict(
+                    ignore=section.get('ignore', False),
+                    start_offset=section.get('start_offset', 0),
+                    end_offset=section.get('end_offset', 0),
+                ))
+            elif section['type'] == 'gen_image':
+                token_source['gen_image'].append(dict(
+                    length=section['token_length'],
+                    timestep=section.get('add_timestep_token', False),
+                    guidance=section.get('add_guidance_token', False),
+                    front_boi=section.get('use_front_boi_token', False),
+                    image_shape=section.get('add_image_shape_token', False),
+                    base_size=section.get('base_size'),
+                    ratio_idx=section.get('ratio_idx'),
+                ))
+            elif section['type'] == 'joint_image':
+                token_source['joint_image'].append(dict(
+                    length=section['token_length'],
+                    timestep=section.get('add_timestep_token', False),
+                    front_boi=section.get('use_front_boi_token', False),
+                    image_shape=section.get('add_image_shape_token', False),
+                    base_size=section.get('base_size'),
+                    ratio_idx=section.get('ratio_idx'),
+                ))
+            else:
+                raise ValueError(f"Invalid section type: {section['type']}")
+
+        # Combine text and image tokens
+        full_token_seq, extra_token_pos = self.encode_sequence(
+            template=template,
+            token_source=dict(token_source),
+            total_length=max_token_length,
+            add_eos=add_eos,
+            add_pad=add_pad,
+            add_bos=add_bos,
+            drop_last=drop_last,
+        )
+        full_seq_token_tensor = torch.tensor(full_token_seq, dtype=torch.long)
+
+        timestep_scatter_index = torch.tensor(extra_token_pos['timestep'], dtype=torch.long) \
+            if 'timestep' in extra_token_pos else None
+        guidance_scatter_index = torch.tensor(extra_token_pos['guidance'], dtype=torch.long) \
+            if 'guidance' in extra_token_pos else None
+        cond_timestep_scatter_index = torch.tensor(extra_token_pos['cond_timestep'], dtype=torch.long) \
+            if 'cond_timestep' in extra_token_pos else None
+        gen_timestep_scatter_index = torch.tensor(extra_token_pos['gen_timestep'], dtype=torch.long) \
+            if 'gen_timestep' in extra_token_pos else None
+
+        # Gen image mask
+        gen_image_slices, gen_image_mask = self.parse_extra_token_pos(extra_token_pos, 'img', full_seq_token_tensor)
+        # Joint image
+        joint_image_slices, _ = self.parse_extra_token_pos(extra_token_pos, 'joint_img', full_seq_token_tensor)
+        # Conditional vae image
+        cond_vae_image_slices, cond_vae_image_mask = self.parse_extra_token_pos(
+            extra_token_pos, 'vae_img', full_seq_token_tensor)
+        # Conditional vit image
+        cond_vit_image_slices, cond_vit_image_mask = self.parse_extra_token_pos(
+            extra_token_pos, 'vit_img', full_seq_token_tensor)
+        # All image slices (gen_image, joint_image)
+        all_image_slices = [
+            slice(start, end + 1)
+            for start, end in zip(extra_token_pos['<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,
+            timestep_scatter_index=timestep_scatter_index,
+            guidance_scatter_index=guidance_scatter_index,
+            text_slices=text_slices,
+            gen_image_slices=gen_image_slices,
+            joint_image_slices=joint_image_slices,
+            cond_vae_image_slices=cond_vae_image_slices,
+            cond_vit_image_slices=cond_vit_image_slices,
+            text_mask=text_mask,
+            gen_image_mask=gen_image_mask,
+            cond_vae_image_mask=cond_vae_image_mask,
+            cond_vit_image_mask=cond_vit_image_mask,
+            real_pos=real_pos,
+            all_image_slices=all_image_slices,
+            cond_timestep_scatter_index=cond_timestep_scatter_index,
+            gen_timestep_scatter_index=gen_timestep_scatter_index,
+        )
+
+    def get_cot_sections(self, cot_text, uncond_kwargs, cot_max_length=None, drop_think=False):
+        if not cot_text:  # None or empty
+            return []
+        if '<think>' in cot_text and '</think>' in cot_text:
+            before_think_sec = cot_text.split('<think>')[0]
+            after_think_sec = cot_text.split('</think>')[1]
+            think_sec = cot_text.split('<think>')[1].split('</think>')[0]
+            return self.get_cot_sections(before_think_sec, uncond_kwargs, drop_think=drop_think) + \
+                ([
+                    dict(type="text", text="<think>"),
+                    dict(type="text", text=think_sec, max_length=cot_max_length, **uncond_kwargs),
+                    dict(type="text", text="</think>")
+                ] if not drop_think else []) + \
+                self.get_cot_sections(after_think_sec, uncond_kwargs, drop_think=drop_think)
+
+        if '<recaption>' in cot_text and '</recaption>' in cot_text:
+            before_recaption_sec = cot_text.split('<recaption>')[0]
+            after_recaption_sec = cot_text.split('</recaption>')[1]
+            recaption_sec = cot_text.split('<recaption>')[1].split('</recaption>')[0]
+            return self.get_cot_sections(before_recaption_sec, uncond_kwargs, drop_think=drop_think) + \
+                [
+                    dict(type="text", text="<recaption>"),
+                    dict(type="text", text=recaption_sec, max_length=cot_max_length, **uncond_kwargs),
+                    dict(type="text", text="</recaption>")
+                ] + \
+                self.get_cot_sections(after_recaption_sec, uncond_kwargs, drop_think=drop_think)
+
+        return [
+            dict(type="text", text=cot_text, **uncond_kwargs),
+        ]
+
+    def apply_general_template(
+            self,
+            message_list,
+            max_length=None,
+            add_assistant_prefix=False,
+            answer="auto",
+            bot_task="auto",
+            sequence_template="instruct",
+            uncond_p=0.0,
+            cfg_factor=1,
+            batchify=False,
+            image_base_size=1024,
+            drop_think=False,
+    ):
+        # If cfg_factor > 1, we need to repeat the unconditioned part
+        if batchify:
+            assert isinstance(message_list[0], list), \
+                f"When batchify is True, message_list should be a list of list, but got [{type(message_list[0])}, ...]."
+            return self.batch_gen_infer(
+                infer_fn=self.apply_general_template,
+                prompt_list=[[]],
+                infer_fn_kwargs_list=[dict(
+                    message_list=message_list_i,
+                    max_length=max_length,
+                    add_assistant_prefix=add_assistant_prefix,
+                    answer=answer,
+                    bot_task=bot_task,
+                    sequence_template=sequence_template,
+                    image_base_size=image_base_size,
+                    drop_think=drop_think,
+                ) for message_list_i in message_list],
+                do_classifier_free_guidance=cfg_factor > 1,
+                condition_repeat_times=1,
+                uncondition_repeat_times=cfg_factor - 1,
+            )
+
+        conv = Conversation()
+        uncond_kwargs = dict(uncond_enabled=uncond_p == 1.0, uncond_p=uncond_p)
+
+        def process_successive_message(_message_list, _cur_message_idx, role, prefix, suffix,
+                                       answer_prefix="", answer_suffix=""):
+            _sub_sections = []
+            while _cur_message_idx < len(message_list) and _message_list[_cur_message_idx]['role'] == role:
+                message = _message_list[_cur_message_idx]
+                if message['type'] == 'text':
+                    text = message['content']
+                    if role == "system":
+                        _sub_sections.append(dict(type="text", text=text))
+                    elif role == "assistant":
+                        if ("<recaption>" in text and "</recaption>" in text) or (
+                                "<think>" in text and "</think>" in text):
+                            _sub_sections.extend(self.get_cot_sections(text, uncond_kwargs, drop_think=drop_think))
+                        else:
+                            _sub_sections.append(dict(type="text", text=text, **uncond_kwargs))
+                    else:
+                        _sub_sections.append(dict(
+                            type="text", text=f"{answer_prefix}{text}{answer_suffix}", **uncond_kwargs))
+                elif message['type'] == 'gen_image':
+                    info = message['content']
+                    assert isinstance(info, ImageInfo), f"Expected ImageInfo, but got {type(info)}"
+                    if role == "assistant":
+                        _sub_sections.append(dict(type="text", text=answer_prefix))
+                    _sub_sections.append(dict(type=message['type'], **info.meta_info))
+                    if role == "assistant":
+                        _sub_sections.append(dict(type="text", text=answer_suffix))
+                elif message['type'] == 'joint_image':
+                    info = message['content']
+                    assert isinstance(info, JointImageInfo), f"Expected JointImageInfo, but got {type(info)}"
+                    _sub_sections.append(dict(type=message['type'], **info.meta_info))
+                else:
+                    raise ValueError(f"Unknown message type: {message['type']}")
+                _cur_message_idx += 1
+            if len(_sub_sections) > 0:
+                # Add role prefix and suffix
+                _sub_sections.insert(0, dict(type='text', text=prefix))
+                _sub_sections.append(dict(type='text', text=suffix))
+            return _sub_sections, _cur_message_idx
+
+        # Define assistant prefix and suffix
+        if (answer == "auto" and sequence_template == "instruct") or answer is True:
+            answer_prefix, answer_suffix = "<answer>", "</answer>"
+        else:
+            answer_prefix, answer_suffix = "", ""
+        if sequence_template == "pretrain":
+            system_suffix = ""
+            user_prefix = ""
+            user_suffix = ""
+            bot_prefix = ""
+            bot_suffix = ""
+        else:
+            system_suffix = f"{conv.sep}"
+            user_prefix = f"{conv.roles[0]}: "
+            user_suffix = f"{conv.sep}"
+            bot_prefix = f"{conv.roles[1]}: "
+            bot_suffix = f"{conv.sep}"
+
+        # Process successive user and assistant messages
+        sections = []
+        cur_message_idx = 0
+        final_role = None
+        while cur_message_idx < len(message_list):
+            # Process successive system messages
+            sub_sections, cur_message_idx = process_successive_message(
+                message_list, cur_message_idx, role="system", prefix="", suffix=system_suffix)
+            # Add to the template and sections
+            sections.extend(sub_sections)
+            if len(sub_sections) > 0:
+                final_role = "system"
+
+            # Process successive user messages
+            sub_sections, cur_message_idx = process_successive_message(
+                message_list, cur_message_idx, role="user", prefix=user_prefix, suffix=user_suffix)
+            # Add to the template and sections
+            sections.extend(sub_sections)
+            if len(sub_sections) > 0:
+                final_role = "user"
+
+            # Process successive assistant messages
+            sub_sections, cur_message_idx = process_successive_message(
+                message_list, cur_message_idx, role="assistant", prefix=bot_prefix, suffix=bot_suffix,
+                answer_prefix=answer_prefix, answer_suffix=answer_suffix,
+            )
+            # Add to the template and sections
+            sections.extend(sub_sections)
+            if len(sub_sections) > 0:
+                final_role = "assistant"
+
+        if add_assistant_prefix:
+            if final_role == "assistant":
+                # Avoid adding prefix twice
+                _bot_prefix = ""
+                # Remove the final bot_suffix
+                if len(sections) > 0 and sections[-1]['type'] == 'text' and sections[-1]['text'] == bot_suffix:
+                    sections = sections[:-1]
+            else:
+                _bot_prefix = bot_prefix
+            # We can add special tokens for the bot lastest message according to different tasks
+            bot_response_prefix = dict(
+                auto=_bot_prefix,
+                think=f"{_bot_prefix}<think>",
+                recaption=f"{_bot_prefix}<recaption>",
+                img_ratio=f"{_bot_prefix}{answer_prefix}<boi><img_size_{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_image_info: Optional[Union[List[JointImageInfo], List[List[JointImageInfo]]]] = 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/think/recaption/img_ratio
+            image_base_size: int = 1024,
+            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 ["auto", "think", "recaption", "img_ratio"], \
+            f"bot_task should be one of ['auto', 'think', 'recaption', 'img_ratio'], but got {bot_task}."
+
+        if batch_message_list is None:
+            # Simple text-to-image or text-cot-to-image task
+            batch_size = len(batch_prompt)
+
+            # Batchify inputs
+            if not isinstance(batch_system_prompt, list):
+                batch_system_prompt = [batch_system_prompt] * batch_size
+            if not isinstance(batch_gen_image_info, list):
+                batch_gen_image_info = [batch_gen_image_info] * batch_size
+            if batch_cot_text is not None:
+                assert len(batch_cot_text) == batch_size, \
+                    (f"batch_cot_text should have the same length as batch_size ({batch_size}), "
+                     f"but got {len(batch_cot_text)}.")
+            else:
+                batch_cot_text = [None] * batch_size
+            if batch_cond_image_info is not None:
+                assert len(batch_cond_image_info) == batch_size, \
+                    (f"batch_cond_image_info should have the same length as batch_size ({batch_size}), "
+                     f"but got {len(batch_cond_image_info)}.")
+                batch_cond_image_info = [
+                    cond_image_info if isinstance(cond_image_info, list) else [cond_image_info]
+                    for cond_image_info in batch_cond_image_info
+                ]
+            else:
+                batch_cond_image_info = [[] for _ in range(batch_size)]
+
+            # Convert single round materials into standard message list
+            batch_message_list = []
+            for prompt, system_prompt, cot_text, gen_image_info, cond_image_info_list in zip(
+                    batch_prompt, batch_system_prompt, batch_cot_text, batch_gen_image_info,
+                    batch_cond_image_info,
+            ):
+                message_list = []
+                # 1. system prompt section
+                if system_prompt:
+                    message_list.append(dict(
+                        role="system", type="text", content=system_prompt, context_type="str"))
+                # 2. user inputs sections
+                #   2.1 image inputs
+                if len(cond_image_info_list) > 0:
+                    message_list.extend([
+                        dict(role="user", type="joint_image", content=cond_image_info, context_type="image_info")
+                        for cond_image_info in cond_image_info_list
+                    ])
+                #   2.2 text inputs
+                message_list.append(dict(
+                    role="user", type="text", content=prompt, context_type="str"))
+                # 3. assistant answer sections
+                if cot_text is not None:
+                    message_list.append(dict(role="assistant", type="text", content=cot_text, context_type="str"))
+                if mode == "gen_image":
+                    message_list.append(dict(
+                        role="assistant", type="gen_image", content=gen_image_info, context_type="image_info"))
+                # ---
+                batch_message_list.append(message_list)
+
+        output, sections = self.apply_general_template(
+            message_list=batch_message_list,
+            max_length=max_length,
+            add_assistant_prefix=default(add_assistant_prefix, mode != "gen_image"),
+            bot_task=bot_task,
+            sequence_template=sequence_template,
+            cfg_factor=cfg_factor,
+            batchify=True,
+            image_base_size=image_base_size,
+            drop_think=drop_think,
+        )
+        return dict(output=output, sections=sections)