make style

Author: a-r-r-o-w

scripts/convert_dcae_to_diffusers.py

@@ -99,7 +99,7 @@ def convert_ae(config_name: str, dtype: torch.dtype):
     hub_id = f"mit-han-lab/{config_name}"
     ckpt_path = hf_hub_download(hub_id, "model.safetensors")
     original_state_dict = get_state_dict(load_file(ckpt_path))
-    
+
     ae = AutoencoderDC(**config).to(dtype=dtype)
 
     for key in list(original_state_dict.keys()):
@@ -122,8 +122,22 @@ def get_ae_config(name: str):
     if name in ["dc-ae-f32c32-sana-1.0"]:
         config = {
             "latent_channels": 32,
-            "encoder_block_types": ("ResBlock", "ResBlock", "ResBlock", "EfficientViTBlock", "EfficientViTBlock", "EfficientViTBlock"),
-            "decoder_block_types": ("ResBlock", "ResBlock", "ResBlock", "EfficientViTBlock", "EfficientViTBlock", "EfficientViTBlock"),
+            "encoder_block_types": (
+                "ResBlock",
+                "ResBlock",
+                "ResBlock",
+                "EfficientViTBlock",
+                "EfficientViTBlock",
+                "EfficientViTBlock",
+            ),
+            "decoder_block_types": (
+                "ResBlock",
+                "ResBlock",
+                "ResBlock",
+                "EfficientViTBlock",
+                "EfficientViTBlock",
+                "EfficientViTBlock",
+            ),
             "encoder_block_out_channels": (128, 256, 512, 512, 1024, 1024),
             "decoder_block_out_channels": (128, 256, 512, 512, 1024, 1024),
             "encoder_qkv_multiscales": ((), (), (), (5,), (5,), (5,)),

src/diffusers/models/autoencoders/autoencoder_dc.py

@@ -88,7 +88,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = self.conv1(hidden_states)
         hidden_states = self.nonlinearity(hidden_states)
         hidden_states = self.conv2(hidden_states)
-        
+
         if self.norm_type == "rms_norm":
             hidden_states = self.norm(hidden_states.movedim(1, -1)).movedim(-1, 1)
         else:
@@ -230,7 +230,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             hidden_states = self.quadratic_attention(qkv)
 
         hidden_states = self.proj_out(hidden_states)
-        
+
         if self.norm_type == "rms_norm":
             hidden_states = self.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
         else: