fix code format style

comfy/ldm/hydit/controlnet.py

@@ -6,7 +6,12 @@ import torch.nn.functional as F
 
 from torch.utils import checkpoint
 
-from comfy.ldm.modules.diffusionmodules.mmdit import Mlp, TimestepEmbedder, PatchEmbed, RMSNorm
+from comfy.ldm.modules.diffusionmodules.mmdit import (
+    Mlp,
+    TimestepEmbedder,
+    PatchEmbed,
+    RMSNorm,
+)
 from comfy.ldm.modules.diffusionmodules.util import timestep_embedding
 from .poolers import AttentionPool
 
@@ -15,11 +20,13 @@ from .models import HunYuanDiTBlock
 
 from .posemb_layers import get_2d_rotary_pos_embed, get_fill_resize_and_crop
 
+
 def zero_module(module):
     for p in module.parameters():
         nn.init.zeros_(p)
     return module
 
+
 def calc_rope(x, patch_size, head_size):
     th = (x.shape[2] + (patch_size // 2)) // patch_size
     tw = (x.shape[3] + (patch_size // 2)) // patch_size
@@ -60,7 +67,9 @@ class HunYuanControlNet(nn.Module):
     log_fn: callable
         The logging function.
     """
-    def __init__(self,
+
+    def __init__(
+        self,
         input_size: tuple = 128,
         patch_size: int = 2,
         in_channels: int = 4,
@@ -68,15 +77,15 @@ class HunYuanControlNet(nn.Module):
         depth: int = 40,
         num_heads: int = 16,
         mlp_ratio: float = 4.3637,
-                 text_states_dim = 1024,
+        text_states_dim=1024,
-                 text_states_dim_t5 = 2048,
+        text_states_dim_t5=2048,
-                 text_len = 77,
+        text_len=77,
-                 text_len_t5 = 256,
+        text_len_t5=256,
-                 qk_norm = True,# See http://arxiv.org/abs/2302.05442 for details.
+        qk_norm=True,  # See http://arxiv.org/abs/2302.05442 for details.
-                 size_cond = False,
+        size_cond=False,
-                 use_style_cond = False,
+        use_style_cond=False,
-                 learn_sigma = True,
+        learn_sigma=True,
-                 norm = "layer",
+        norm="layer",
         log_fn: callable = print,
         attn_precision=None,
         dtype=None,
@@ -103,19 +112,44 @@ class HunYuanControlNet(nn.Module):
         self.dtype = dtype
         self.latent_format = comfy.latent_formats.SDXL
 
-
         self.mlp_t5 = nn.Sequential(
-            nn.Linear(self.text_states_dim_t5, self.text_states_dim_t5 * 4, bias=True, dtype=dtype, device=device),
+            nn.Linear(
+                self.text_states_dim_t5,
+                self.text_states_dim_t5 * 4,
+                bias=True,
+                dtype=dtype,
+                device=device,
+            ),
             nn.SiLU(),
-            nn.Linear(self.text_states_dim_t5 * 4, self.text_states_dim, bias=True, dtype=dtype, device=device),
+            nn.Linear(
+                self.text_states_dim_t5 * 4,
+                self.text_states_dim,
+                bias=True,
+                dtype=dtype,
+                device=device,
+            ),
         )
         # learnable replace
         self.text_embedding_padding = nn.Parameter(
-            torch.randn(self.text_len + self.text_len_t5, self.text_states_dim, dtype=dtype, device=device))
+            torch.randn(
+                self.text_len + self.text_len_t5,
+                self.text_states_dim,
+                dtype=dtype,
+                device=device,
+            )
+        )
 
         # Attention pooling
         pooler_out_dim = 1024
-        self.pooler = AttentionPool(self.text_len_t5, self.text_states_dim_t5, num_heads=8, output_dim=pooler_out_dim, dtype=dtype, device=device, operations=operations)
+        self.pooler = AttentionPool(
+            self.text_len_t5,
+            self.text_states_dim_t5,
+            num_heads=8,
+            output_dim=pooler_out_dim,
+            dtype=dtype,
+            device=device,
+            operations=operations,
+        )
 
         # Dimension of the extra input vectors
         self.extra_in_dim = pooler_out_dim
@@ -126,24 +160,42 @@ class HunYuanControlNet(nn.Module):
 
         if self.use_style_cond:
             # Here we use a default learned embedder layer for future extension.
-            self.style_embedder = nn.Embedding(1, hidden_size, dtype=dtype, device=device)
+            self.style_embedder = nn.Embedding(
+                1, hidden_size, dtype=dtype, device=device
+            )
             self.extra_in_dim += hidden_size
 
         # Text embedding for `add`
-        self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, dtype=dtype, device=device, operations=operations)
+        self.x_embedder = PatchEmbed(
-        self.t_embedder = TimestepEmbedder(hidden_size, dtype=dtype, device=device, operations=operations)
+            input_size,
+            patch_size,
+            in_channels,
+            hidden_size,
+            dtype=dtype,
+            device=device,
+            operations=operations,
+        )
+        self.t_embedder = TimestepEmbedder(
+            hidden_size, dtype=dtype, device=device, operations=operations
+        )
         self.extra_embedder = nn.Sequential(
-            operations.Linear(self.extra_in_dim, hidden_size * 4, dtype=dtype, device=device),
+            operations.Linear(
+                self.extra_in_dim, hidden_size * 4, dtype=dtype, device=device
+            ),
             nn.SiLU(),
-            operations.Linear(hidden_size * 4, hidden_size, bias=True, dtype=dtype, device=device),
+            operations.Linear(
+                hidden_size * 4, hidden_size, bias=True, dtype=dtype, device=device
+            ),
         )
 
         # Image embedding
         num_patches = self.x_embedder.num_patches
 
         # HUnYuanDiT Blocks
-        self.blocks = nn.ModuleList([
+        self.blocks = nn.ModuleList(
-            HunYuanDiTBlock(hidden_size=hidden_size,
+            [
+                HunYuanDiTBlock(
+                    hidden_size=hidden_size,
                     c_emb_size=hidden_size,
                     num_heads=num_heads,
                     mlp_ratio=mlp_ratio,
@@ -157,20 +209,31 @@ class HunYuanControlNet(nn.Module):
                     operations=operations,
                 )
                 for _ in range(19)
-        ])
+            ]
+        )
 
         # Input zero linear for the first block
-        self.before_proj = zero_module(nn.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device))
+        self.before_proj = zero_module(
+            nn.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device)
+        )
 
         # Output zero linear for the every block
         self.after_proj_list = nn.ModuleList(
-            [zero_module(nn.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device)) for _ in range(len(self.blocks))]
+            [
+                zero_module(
+                    nn.Linear(
+                        self.hidden_size, self.hidden_size, dtype=dtype, device=device
+                    )
+                )
+                for _ in range(len(self.blocks))
+            ]
         )
 
-    def forward(self,
+    def forward(
+        self,
         x: torch.Tensor,
         t: torch.Tensor = None,
-                condition = None,
+        condition=None,
         encoder_hidden_states: Optional[torch.Tensor] = None,
         text_embedding_mask=None,
         encoder_hidden_states_t5=None,
@@ -178,7 +241,7 @@ class HunYuanControlNet(nn.Module):
         image_meta_size=None,
         style=None,
         transformer_options=None,
-                **kwarg
+        **kwarg,
     ):
         """
         Forward pass of the encoder.
@@ -218,22 +281,31 @@ class HunYuanControlNet(nn.Module):
 
         padding = comfy.ops.cast_to_input(self.text_embedding_padding, text_states)
 
-        text_states[:,-self.text_len:] = torch.where(text_states_mask[:,-self.text_len:].unsqueeze(2), text_states[:,-self.text_len:], padding[:self.text_len])
+        text_states[:, -self.text_len :] = torch.where(
-        text_states_t5[:,-self.text_len_t5:] = torch.where(text_states_t5_mask[:,-self.text_len_t5:].unsqueeze(2), text_states_t5[:,-self.text_len_t5:], padding[self.text_len:])
+            text_states_mask[:, -self.text_len :].unsqueeze(2),
+            text_states[:, -self.text_len :],
+            padding[: self.text_len],
+        )
+        text_states_t5[:, -self.text_len_t5 :] = torch.where(
+            text_states_t5_mask[:, -self.text_len_t5 :].unsqueeze(2),
+            text_states_t5[:, -self.text_len_t5 :],
+            padding[self.text_len :],
+        )
 
         text_states = torch.cat([text_states, text_states_t5], dim=1)  # 2,205，1024
 
-        _, _, oh, ow = x.shape
+        # _, _, oh, ow = x.shape
-        th, tw = oh // self.patch_size, ow // self.patch_size
+        # th, tw = oh // self.patch_size, ow // self.patch_size
 
         # Get image RoPE embedding according to `reso`lution.
-        freqs_cis_img = calc_rope(x, self.patch_size, self.hidden_size // self.num_heads) # (cos_cis_img, sin_cis_img)
+        freqs_cis_img = calc_rope(
+            x, self.patch_size, self.hidden_size // self.num_heads
+        )  # (cos_cis_img, sin_cis_img)
 
         # ========================= Build time and image embedding =========================
         t = self.t_embedder(t, dtype=self.dtype)
         x = self.x_embedder(x)
 
-
         # ========================= Concatenate all extra vectors =========================
         # Build text tokens with pooling
         extra_vec = self.pooler(encoder_hidden_states_t5)
@@ -264,5 +336,4 @@ class HunYuanControlNet(nn.Module):
             x = block(x, c, text_states, freqs_cis_img)
             controls.append(self.after_proj_list[layer](x))  # zero linear for output
 
-        return {'output': controls}
+        return {"output": controls}
-