From e270142f9a4066fd76bca3cd259b60e8b1051bb3 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Sat, 14 Dec 2024 21:12:39 +0100
Subject: [PATCH] Formatting

---
 comfy/ldm/pixart/pixart.py |  2 --
 comfy/utils.py             | 36 +++++++++++++-----------------------
 2 files changed, 13 insertions(+), 25 deletions(-)

diff --git a/comfy/ldm/pixart/pixart.py b/comfy/ldm/pixart/pixart.py
index d2e0f45d0..cd572efce 100644
--- a/comfy/ldm/pixart/pixart.py
+++ b/comfy/ldm/pixart/pixart.py
@@ -3,7 +3,6 @@
 # https://github.com/PixArt-alpha/PixArt-sigma [Apache 2.0 license]
 import torch
 import torch.nn as nn
-import numpy as np
 
 from .blocks import (
     t2i_modulate,
@@ -11,7 +10,6 @@ from .blocks import (
     AttentionKVCompress,
     MultiHeadCrossAttention,
     T2IFinalLayer,
-    TimestepEmbedder,
 )
 from comfy.ldm.modules.diffusionmodules.mmdit import PatchEmbed, TimestepEmbedder, Mlp, get_1d_sincos_pos_embed_from_grid_torch
 
diff --git a/comfy/utils.py b/comfy/utils.py
index 1e021ac7a..a0cf8d3fd 100644
--- a/comfy/utils.py
+++ b/comfy/utils.py
@@ -385,52 +385,42 @@ def mmdit_to_diffusers(mmdit_config, output_prefix=""):
     return key_map
 
 PIXART_MAP_BASIC = {
-    # Resolution
     ("csize_embedder.mlp.0.weight", "adaln_single.emb.resolution_embedder.linear_1.weight"),
-    ("csize_embedder.mlp.0.bias",   "adaln_single.emb.resolution_embedder.linear_1.bias"),
+    ("csize_embedder.mlp.0.bias", "adaln_single.emb.resolution_embedder.linear_1.bias"),
     ("csize_embedder.mlp.2.weight", "adaln_single.emb.resolution_embedder.linear_2.weight"),
-    ("csize_embedder.mlp.2.bias",   "adaln_single.emb.resolution_embedder.linear_2.bias"),
-    # Aspect ratio
+    ("csize_embedder.mlp.2.bias", "adaln_single.emb.resolution_embedder.linear_2.bias"),
     ("ar_embedder.mlp.0.weight", "adaln_single.emb.aspect_ratio_embedder.linear_1.weight"),
-    ("ar_embedder.mlp.0.bias",   "adaln_single.emb.aspect_ratio_embedder.linear_1.bias"),
+    ("ar_embedder.mlp.0.bias", "adaln_single.emb.aspect_ratio_embedder.linear_1.bias"),
     ("ar_embedder.mlp.2.weight", "adaln_single.emb.aspect_ratio_embedder.linear_2.weight"),
-    ("ar_embedder.mlp.2.bias",   "adaln_single.emb.aspect_ratio_embedder.linear_2.bias"),
-    # Patch embeddings
+    ("ar_embedder.mlp.2.bias", "adaln_single.emb.aspect_ratio_embedder.linear_2.bias"),
     ("x_embedder.proj.weight", "pos_embed.proj.weight"),
     ("x_embedder.proj.bias", "pos_embed.proj.bias"),
-    # Caption projection
     ("y_embedder.y_embedding", "caption_projection.y_embedding"),
     ("y_embedder.y_proj.fc1.weight", "caption_projection.linear_1.weight"),
     ("y_embedder.y_proj.fc1.bias", "caption_projection.linear_1.bias"),
     ("y_embedder.y_proj.fc2.weight", "caption_projection.linear_2.weight"),
     ("y_embedder.y_proj.fc2.bias", "caption_projection.linear_2.bias"),
-    # AdaLN-single LN
     ("t_embedder.mlp.0.weight", "adaln_single.emb.timestep_embedder.linear_1.weight"),
     ("t_embedder.mlp.0.bias", "adaln_single.emb.timestep_embedder.linear_1.bias"),
     ("t_embedder.mlp.2.weight", "adaln_single.emb.timestep_embedder.linear_2.weight"),
     ("t_embedder.mlp.2.bias", "adaln_single.emb.timestep_embedder.linear_2.bias"),
-    # Shared norm
     ("t_block.1.weight", "adaln_single.linear.weight"),
     ("t_block.1.bias", "adaln_single.linear.bias"),
-    # Final block
     ("final_layer.linear.weight", "proj_out.weight"),
     ("final_layer.linear.bias", "proj_out.bias"),
     ("final_layer.scale_shift_table", "scale_shift_table"),
 }
 
 PIXART_MAP_BLOCK = {
-    (f"scale_shift_table", f"scale_shift_table"),
-    # Projection
-    (f"attn.proj.weight", f"attn1.to_out.0.weight"),
-    (f"attn.proj.bias",   f"attn1.to_out.0.bias"),
-    # Feed-forward
-    (f"mlp.fc1.weight", f"ff.net.0.proj.weight"),
-    (f"mlp.fc1.bias",   f"ff.net.0.proj.bias"),
-    (f"mlp.fc2.weight", f"ff.net.2.weight"),
-    (f"mlp.fc2.bias",   f"ff.net.2.bias"),
-    # Cross-attention (proj)
-    (f"cross_attn.proj.weight" ,f"attn2.to_out.0.weight"),
-    (f"cross_attn.proj.bias"   ,f"attn2.to_out.0.bias"),
+    ("scale_shift_table", "scale_shift_table"),
+    ("attn.proj.weight", "attn1.to_out.0.weight"),
+    ("attn.proj.bias", "attn1.to_out.0.bias"),
+    ("mlp.fc1.weight", "ff.net.0.proj.weight"),
+    ("mlp.fc1.bias", "ff.net.0.proj.bias"),
+    ("mlp.fc2.weight", "ff.net.2.weight"),
+    ("mlp.fc2.bias", "ff.net.2.bias"),
+    ("cross_attn.proj.weight" ,"attn2.to_out.0.weight"),
+    ("cross_attn.proj.bias"   ,"attn2.to_out.0.bias"),
 }
 
 def pixart_to_diffusers(mmdit_config, output_prefix=""):