ruff

Author: hlky

scripts/convert_z_image_controlnet_to_diffusers.py

@@ -1,14 +1,15 @@
 import argparse
 from contextlib import nullcontext
 
-import torch
 import safetensors.torch
+import torch
 from accelerate import init_empty_weights
 from huggingface_hub import hf_hub_download
 
-from diffusers.utils.import_utils import is_accelerate_available
 from diffusers.models import ZImageTransformer2DModel
 from diffusers.models.controlnets.controlnet_z_image import ZImageControlNetModel
+from diffusers.utils.import_utils import is_accelerate_available
+
 
 """
 python scripts/convert_z_image_controlnet_to_diffusers.py  \
@@ -42,16 +43,28 @@ def load_original_checkpoint(args):
     original_state_dict = safetensors.torch.load_file(ckpt_path)
     return original_state_dict
 
+
 def load_z_image(args):
-    model = ZImageTransformer2DModel.from_pretrained(args.original_z_image_repo_id, subfolder="transformer", torch_dtype=torch.bfloat16)
+    model = ZImageTransformer2DModel.from_pretrained(
+        args.original_z_image_repo_id, subfolder="transformer", torch_dtype=torch.bfloat16
+    )
     return model.state_dict(), model.config
 
+
 def convert_z_image_controlnet_checkpoint_to_diffusers(z_image, original_state_dict):
     converted_state_dict = {}
 
     converted_state_dict.update(original_state_dict)
 
-    to_copy = {"all_x_embedder.", "noise_refiner.", "context_refiner.", "t_embedder.", "cap_embedder.", "x_pad_token", "cap_pad_token"}
+    to_copy = {
+        "all_x_embedder.",
+        "noise_refiner.",
+        "context_refiner.",
+        "t_embedder.",
+        "cap_embedder.",
+        "x_pad_token",
+        "cap_pad_token",
+    }
 
     for key in z_image.keys():
         for copy_key in to_copy:

src/diffusers/models/controlnets/controlnet_z_image.py

@@ -20,23 +20,26 @@
 
 from ...configuration_utils import ConfigMixin, register_to_config
 from ...loaders import PeftAdapterMixin
-from ...models.normalization import RMSNorm
 from ..controlnets.controlnet import zero_module
 from ..modeling_utils import ModelMixin
-from ..transformers.transformer_z_image import ZImageTransformer2DModel, ZImageTransformerBlock, RopeEmbedder, TimestepEmbedder, SEQ_MULTI_OF, ADALN_EMBED_DIM
+from ..transformers.transformer_z_image import (
+    SEQ_MULTI_OF,
+    ZImageTransformer2DModel,
+    ZImageTransformerBlock,
+)
 
 
 class ZImageControlTransformerBlock(ZImageTransformerBlock):
     def __init__(
-        self, 
+        self,
         layer_id: int,
         dim: int,
         n_heads: int,
         n_kv_heads: int,
         norm_eps: float,
         qk_norm: bool,
         modulation=True,
-        block_id=0
+        block_id=0,
     ):
         super().__init__(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation)
         self.block_id = block_id
@@ -57,7 +60,8 @@ def forward(self, c: torch.Tensor, x: torch.Tensor, **kwargs):
         all_c += [c_skip, c]
         c = torch.stack(all_c)
         return c
-    
+
+
 class ZImageControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
     _supports_gradient_checkpointing = True
 
@@ -72,7 +76,7 @@ def __init__(
         n_kv_heads=30,
         norm_eps=1e-5,
         qk_norm=True,
-        control_layers_places: List[int]=None,
+        control_layers_places: List[int] = None,
         control_in_dim=None,
     ):
         super().__init__()
@@ -84,15 +88,7 @@ def __init__(
         # control blocks
         self.control_layers = nn.ModuleList(
             [
-                ZImageControlTransformerBlock(
-                    i, 
-                    dim, 
-                    n_heads, 
-                    n_kv_heads, 
-                    norm_eps, 
-                    qk_norm,
-                    block_id=i
-                )
+                ZImageControlTransformerBlock(i, dim, n_heads, n_kv_heads, norm_eps, qk_norm, block_id=i)
                 for i in self.control_layers_places
             ]
         )
@@ -425,7 +421,9 @@ def forward(
 
         if torch.is_grad_enabled() and self.gradient_checkpointing:
             for layer in self.control_noise_refiner:
-                control_context = self._gradient_checkpointing_func(layer, control_context, x_attn_mask, x_freqs_cis, adaln_input)
+                control_context = self._gradient_checkpointing_func(
+                    layer, control_context, x_attn_mask, x_freqs_cis, adaln_input
+                )
         else:
             for layer in self.control_noise_refiner:
                 control_context = layer(control_context, x_attn_mask, x_freqs_cis, adaln_input)
@@ -440,14 +438,14 @@ def forward(
         control_context_unified = pad_sequence(control_context_unified, batch_first=True, padding_value=0.0)
         c = control_context_unified
 
-        new_kwargs = dict(x=unified, attn_mask=unified_attn_mask, freqs_cis=unified_freqs_cis, adaln_input=adaln_input)
-        
+        new_kwargs = {"x": unified, "attn_mask": unified_attn_mask, "freqs_cis": unified_freqs_cis, "adaln_input": adaln_input}
+
         for layer in self.control_layers:
             if torch.is_grad_enabled() and self.gradient_checkpointing:
                 c = self._gradient_checkpointing_func(layer, c, **new_kwargs)
             else:
                 c = layer(c, **new_kwargs)
- 
+
         hints = torch.unbind(c)[:-1] * conditioning_scale
         controlnet_block_samples = {}
         for layer_idx in range(self.n_layers):

src/diffusers/models/transformers/transformer_z_image.py

@@ -538,7 +538,7 @@ def forward(
         cap_feats: List[torch.Tensor],
         patch_size=2,
         f_patch_size=1,
-        controlnet_block_samples: Optional[dict[int, torch.Tensor]]=None,
+        controlnet_block_samples: Optional[dict[int, torch.Tensor]] = None,
         return_dict: bool = True,
     ):
         assert patch_size in self.all_patch_size

src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py

@@ -89,7 +89,6 @@ def retrieve_latents(
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
-
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
 def retrieve_timesteps(
     scheduler,
@@ -509,7 +508,7 @@ def __call__(
             num_images_per_prompt=num_images_per_prompt,
             device=device,
             dtype=self.vae.dtype,
-        )        
+        )
         height, width = control_image.shape[-2:]
         control_image = retrieve_latents(self.vae.encode(control_image), generator=generator)
         control_image = (control_image - self.vae.config.shift_factor) * self.vae.config.scaling_factor