make style

Author: yiyixuxu

src/diffusers/__init__.py

@@ -184,9 +184,9 @@
             "AutoencoderKLAllegro",
             "AutoencoderKLCogVideoX",
             "AutoencoderKLCosmos",
-            "AutoencoderKLHunyuanVideo",
-            "AutoencoderKLHunyuanImageRefiner",
             "AutoencoderKLHunyuanImage",
+            "AutoencoderKLHunyuanImageRefiner",
+            "AutoencoderKLHunyuanVideo",
             "AutoencoderKLLTXVideo",
             "AutoencoderKLMagvit",
             "AutoencoderKLMochi",
@@ -872,9 +872,9 @@
             AutoencoderKLAllegro,
             AutoencoderKLCogVideoX,
             AutoencoderKLCosmos,
-            AutoencoderKLHunyuanVideo,
             AutoencoderKLHunyuanImage,
             AutoencoderKLHunyuanImageRefiner,
+            AutoencoderKLHunyuanVideo,
             AutoencoderKLLTXVideo,
             AutoencoderKLMagvit,
             AutoencoderKLMochi,
@@ -905,9 +905,9 @@
             HunyuanDiT2DControlNetModel,
             HunyuanDiT2DModel,
             HunyuanDiT2DMultiControlNetModel,
+            HunyuanImageTransformer2DModel,
             HunyuanVideoFramepackTransformer3DModel,
             HunyuanVideoTransformer3DModel,
-            HunyuanImageTransformer2DModel,
             I2VGenXLUNet,
             Kandinsky3UNet,
             LatteTransformer3DModel,

src/diffusers/models/__init__.py

@@ -73,7 +73,6 @@
     _import_structure["transformers.dit_transformer_2d"] = ["DiTTransformer2DModel"]
     _import_structure["transformers.dual_transformer_2d"] = ["DualTransformer2DModel"]
     _import_structure["transformers.hunyuan_transformer_2d"] = ["HunyuanDiT2DModel"]
-    _import_structure["transformers.transformer_hunyuanimage"] = ["HunyuanImageTransformer2DModel"]
     _import_structure["transformers.latte_transformer_3d"] = ["LatteTransformer3DModel"]
     _import_structure["transformers.lumina_nextdit2d"] = ["LuminaNextDiT2DModel"]
     _import_structure["transformers.pixart_transformer_2d"] = ["PixArtTransformer2DModel"]
@@ -93,6 +92,7 @@
     _import_structure["transformers.transformer_hidream_image"] = ["HiDreamImageTransformer2DModel"]
     _import_structure["transformers.transformer_hunyuan_video"] = ["HunyuanVideoTransformer3DModel"]
     _import_structure["transformers.transformer_hunyuan_video_framepack"] = ["HunyuanVideoFramepackTransformer3DModel"]
+    _import_structure["transformers.transformer_hunyuanimage"] = ["HunyuanImageTransformer2DModel"]
     _import_structure["transformers.transformer_ltx"] = ["LTXVideoTransformer3DModel"]
     _import_structure["transformers.transformer_lumina2"] = ["Lumina2Transformer2DModel"]
     _import_structure["transformers.transformer_mochi"] = ["MochiTransformer3DModel"]
@@ -132,9 +132,9 @@
             AutoencoderKLAllegro,
             AutoencoderKLCogVideoX,
             AutoencoderKLCosmos,
-            AutoencoderKLHunyuanVideo,
             AutoencoderKLHunyuanImage,
             AutoencoderKLHunyuanImageRefiner,
+            AutoencoderKLHunyuanVideo,
             AutoencoderKLLTXVideo,
             AutoencoderKLMagvit,
             AutoencoderKLMochi,

src/diffusers/models/autoencoders/__init__.py

@@ -5,12 +5,12 @@
 from .autoencoder_kl_cogvideox import AutoencoderKLCogVideoX
 from .autoencoder_kl_cosmos import AutoencoderKLCosmos
 from .autoencoder_kl_hunyuan_video import AutoencoderKLHunyuanVideo
+from .autoencoder_kl_hunyuanimage import AutoencoderKLHunyuanImage
+from .autoencoder_kl_hunyuanimage_refiner import AutoencoderKLHunyuanImageRefiner
 from .autoencoder_kl_ltx import AutoencoderKLLTXVideo
 from .autoencoder_kl_magvit import AutoencoderKLMagvit
 from .autoencoder_kl_mochi import AutoencoderKLMochi
 from .autoencoder_kl_qwenimage import AutoencoderKLQwenImage
-from .autoencoder_kl_hunyuanimage import AutoencoderKLHunyuanImage
-from .autoencoder_kl_hunyuanimage_refiner import AutoencoderKLHunyuanImageRefiner
 from .autoencoder_kl_temporal_decoder import AutoencoderKLTemporalDecoder
 from .autoencoder_kl_wan import AutoencoderKLWan
 from .autoencoder_oobleck import AutoencoderOobleck

src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py

@@ -24,7 +24,6 @@
 from ...utils import logging
 from ...utils.accelerate_utils import apply_forward_hook
 from ..activations import get_activation
-from ..attention_processor import Attention
 from ..modeling_outputs import AutoencoderKLOutput
 from ..modeling_utils import ModelMixin
 from .vae import DecoderOutput, DiagonalGaussianDistribution
@@ -126,8 +125,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         x = x.squeeze(1).reshape(batch_size, frames, height, width, channels).permute(0, 4, 1, 2, 3)
         x = self.proj_out(x)
-    
-        return x + identity 
+
+        return x + identity
 
 
 class HunyuanImageRefinerUpsampleDCAE(nn.Module):
@@ -143,11 +142,11 @@ def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: b
     def _dcae_upsample_rearrange(tensor, r1=1, r2=2, r3=2):
         """
         Convert (b, r1*r2*r3*c, f, h, w) -> (b, c, r1*f, r2*h, r3*w)
-        
+
         Args:
             tensor: Input tensor of shape (b, r1*r2*r3*c, f, h, w)
             r1: temporal upsampling factor
-            r2: height upsampling factor  
+            r2: height upsampling factor
             r3: width upsampling factor
         """
         b, packed_c, f, h, w = tensor.shape
@@ -187,12 +186,11 @@ def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample:
         self.add_temporal_downsample = add_temporal_downsample
         self.group_size = factor * in_channels // out_channels
 
-
     @staticmethod
     def _dcae_downsample_rearrange(self, tensor, r1=1, r2=2, r3=2):
         """
         Convert (b, c, r1*f, r2*h, r3*w) -> (b, r1*r2*r3*c, f, h, w)
-        
+
         This packs spatial/temporal dimensions into channels (opposite of upsample)
         """
         b, c, packed_f, packed_h, packed_w = tensor.shape
@@ -202,7 +200,6 @@ def _dcae_downsample_rearrange(self, tensor, r1=1, r2=2, r3=2):
         tensor = tensor.permute(0, 2, 4, 6, 1, 3, 5, 7)
         return tensor.reshape(b, r1 * r2 * r3 * c, f, h, w)
 
-
     def forward(self, x: torch.Tensor):
         r1 = 2 if self.add_temporal_downsample else 1
         h = self.conv(x)
@@ -304,16 +301,13 @@ def __init__(
         self.gradient_checkpointing = False
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-
         hidden_states = self.resnets[0](hidden_states)
 
-
         for attn, resnet in zip(self.attentions, self.resnets[1:]):
             if attn is not None:
                 hidden_states = attn(hidden_states)
             hidden_states = resnet(hidden_states)
 
-
         return hidden_states
 
 
@@ -356,7 +350,6 @@ def __init__(
         self.gradient_checkpointing = False
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-
         for resnet in self.resnets:
             hidden_states = resnet(hidden_states)
 
@@ -461,7 +454,6 @@ def __init__(
                 )
                 input_channel = output_channel
             else:
-
                 add_temporal_downsample = i >= np.log2(spatial_compression_ratio // temporal_compression_ratio)
                 downsample_out_channels = block_out_channels[i + 1] if downsample_match_channel else output_channel
                 down_block = HunyuanImageRefinerDownBlock3D(
@@ -518,7 +510,7 @@ class HunyuanImageRefinerDecoder3D(nn.Module):
     def __init__(
         self,
         in_channels: int = 32,
-        out_channels: int = 3, 
+        out_channels: int = 3,
         block_out_channels: Tuple[int, ...] = (1024, 1024, 512, 256, 128),
         layers_per_block: int = 2,
         spatial_compression_ratio: int = 16,
@@ -574,10 +566,8 @@ def __init__(
         self.gradient_checkpointing = False
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-
         hidden_states = self.conv_in(hidden_states) + hidden_states.repeat_interleave(repeats=self.repeat, dim=1)
 
-
         if torch.is_grad_enabled() and self.gradient_checkpointing:
             hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states)
 
@@ -598,8 +588,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 class AutoencoderKLHunyuanImageRefiner(ModelMixin, ConfigMixin):
     r"""
-    A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos.
-    Used for HunyuanImage-2.1 Refiner..
+    A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Used for
+    HunyuanImage-2.1 Refiner..
 
     This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
     for all models (such as downloading or saving).
@@ -621,7 +611,7 @@ def __init__(
         upsample_match_channel: bool = True,
         scaling_factor: float = 1.03682,
     ) -> None:
-        super().__init__() 
+        super().__init__()
 
         self.encoder = HunyuanImageRefinerEncoder3D(
             in_channels=in_channels,
@@ -655,7 +645,6 @@ def __init__(
         # intermediate tiles together, the memory requirement can be lowered.
         self.use_tiling = False
 
-
         # The minimal tile height and width for spatial tiling to be used
         self.tile_sample_min_height = 256
         self.tile_sample_min_width = 256
@@ -763,7 +752,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor:
 
         if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height):
             return self.tiled_decode(z)
-        
+
         dec = self.decoder(z)
 
         return dec
@@ -829,7 +818,7 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
                 The latent representation of the encoded videos.
         """
         _, _, _, height, width = x.shape
-        
+
         tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
         tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
         overlap_height = int(tile_latent_min_height * (1 - self.tile_overlap_factor))  # 256 * (1 - 0.25) = 192
@@ -922,7 +911,6 @@ def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
 
         return dec
 
-
     def forward(
         self,
         sample: torch.Tensor,

src/diffusers/models/transformers/__init__.py

@@ -8,7 +8,6 @@
     from .dit_transformer_2d import DiTTransformer2DModel
     from .dual_transformer_2d import DualTransformer2DModel
     from .hunyuan_transformer_2d import HunyuanDiT2DModel
-    from .transformer_hunyuanimage import HunyuanImageTransformer2DModel
     from .latte_transformer_3d import LatteTransformer3DModel
     from .lumina_nextdit2d import LuminaNextDiT2DModel
     from .pixart_transformer_2d import PixArtTransformer2DModel
@@ -28,6 +27,7 @@
     from .transformer_hidream_image import HiDreamImageTransformer2DModel
     from .transformer_hunyuan_video import HunyuanVideoTransformer3DModel
     from .transformer_hunyuan_video_framepack import HunyuanVideoFramepackTransformer3DModel
+    from .transformer_hunyuanimage import HunyuanImageTransformer2DModel
     from .transformer_ltx import LTXVideoTransformer3DModel
     from .transformer_lumina2 import Lumina2Transformer2DModel
     from .transformer_mochi import MochiTransformer3DModel

src/diffusers/models/transformers/transformer_hunyuanimage.py

@@ -12,9 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, List, Optional, Tuple, Union
-
 import math
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -217,7 +216,9 @@ def __init__(
             self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
 
     def forward(
-        self, timestep: torch.Tensor, guidance: Optional[torch.Tensor] = None,
+        self,
+        timestep: torch.Tensor,
+        guidance: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         timesteps_proj = self.time_proj(timestep)
         timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=timestep.dtype))  # (N, D)
@@ -381,13 +382,15 @@ def forward(
 
 
 class HunyuanImageRotaryPosEmbed(nn.Module):
-    def __init__(self, patch_size: Union[Tuple, List[int]], rope_dim: Union[Tuple, List[int]], theta: float = 256.0) -> None:
+    def __init__(
+        self, patch_size: Union[Tuple, List[int]], rope_dim: Union[Tuple, List[int]], theta: float = 256.0
+    ) -> None:
         super().__init__()
 
-        if not isinstance(patch_size, (tuple, list)) or not len(patch_size) in [2, 3]:
+        if not isinstance(patch_size, (tuple, list)) or len(patch_size) not in [2, 3]:
             raise ValueError(f"patch_size must be a tuple or list of length 2 or 3, got {patch_size}")
-        
-        if not isinstance(rope_dim, (tuple, list)) or not len(rope_dim) in [2, 3]:
+
+        if not isinstance(rope_dim, (tuple, list)) or len(rope_dim) not in [2, 3]:
             raise ValueError(f"rope_dim must be a tuple or list of length 2 or 3, got {rope_dim}")
 
         if not len(patch_size) == len(rope_dim):
@@ -398,7 +401,6 @@ def __init__(self, patch_size: Union[Tuple, List[int]], rope_dim: Union[Tuple, L
         self.theta = theta
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-
         if hidden_states.ndim == 5:
             _, _, frame, height, width = hidden_states.shape
             patch_size_frame, patch_size_height, patch_size_width = self.patch_size
@@ -805,7 +807,7 @@ def forward(
             sizes = (frame, height, width)
         else:
             raise ValueError(f"hidden_states must be a 4D or 5D tensor, got {hidden_states.shape}")
-        
+
         post_patch_sizes = tuple(d // p for d, p in zip(sizes, self.config.patch_size))
 
         # 1. RoPE
@@ -816,7 +818,7 @@ def forward(
         temb = self.time_guidance_embed(timestep, guidance)
         hidden_states = self.x_embedder(hidden_states)
         encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask)
-        
+
         if self.context_embedder_2 is not None and encoder_hidden_states_2 is not None:
             encoder_hidden_states_2 = self.context_embedder_2(encoder_hidden_states_2)
 
@@ -912,7 +914,7 @@ def forward(
         hidden_states = hidden_states.reshape(*reshape_dims)
 
         # create permutation pattern: batch, channels, then interleave post_patch and patch dims
-        # For 4D: [0, 3, 1, 4, 2, 5] -> batch, channels, post_patch_height, patch_size_height, post_patch_width, patch_size_width  
+        # For 4D: [0, 3, 1, 4, 2, 5] -> batch, channels, post_patch_height, patch_size_height, post_patch_width, patch_size_width
         # For 5D: [0, 4, 1, 5, 2, 6, 3, 7] -> batch, channels, post_patch_frame, patch_size_frame, post_patch_height, patch_size_height, post_patch_width, patch_size_width
         ndim = len(post_patch_sizes)
         permute_pattern = [0, ndim + 1]  # batch, channels
@@ -922,7 +924,9 @@ def forward(
 
         # flatten patch dimensions: flatten each (post_patch_size, patch_size) pair
         # batch_size, channels, post_patch_sizes[0] * patch_sizes[0], post_patch_sizes[1] * patch_sizes[1], ...
-        final_dims = [batch_size, out_channels] + [post_patch * patch for post_patch, patch in zip(post_patch_sizes, self.config.patch_size)]
+        final_dims = [batch_size, out_channels] + [
+            post_patch * patch for post_patch, patch in zip(post_patch_sizes, self.config.patch_size)
+        ]
         hidden_states = hidden_states.reshape(*final_dims)
 
         if USE_PEFT_BACKEND: