diffsynth-engine 0.6.0__tar.gz → 0.6.1.dev2__tar.gz

{diffsynth_engine-0.6.0 → diffsynth_engine-0.6.1.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.6.0
+Version: 0.6.1.dev2
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

{diffsynth_engine-0.6.0 → diffsynth_engine-0.6.1.dev2}/diffsynth_engine/models/flux/flux_dit.py

@@ -243,7 +243,9 @@ class FluxDoubleTransformerBlock(nn.Module):
         self.norm_msa_a = AdaLayerNormZero(dim, device=device, dtype=dtype)
         self.norm_mlp_a = AdaLayerNormZero(dim, device=device, dtype=dtype)
         self.ff_a = nn.Sequential(
-            nn.Linear(dim, dim * 4), nn.GELU(approximate="tanh"), nn.Linear(dim * 4, dim, device=device, dtype=dtype)
+            nn.Linear(dim, dim * 4, device=device, dtype=dtype),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(dim * 4, dim, device=device, dtype=dtype)
         )
         # Text
         self.norm_msa_b = AdaLayerNormZero(dim, device=device, dtype=dtype)
@@ -313,10 +315,10 @@ class FluxSingleTransformerBlock(nn.Module):
         self.norm = AdaLayerNormZero(dim, device=device, dtype=dtype)
         self.attn = FluxSingleAttention(dim, num_heads, attn_kwargs=attn_kwargs, device=device, dtype=dtype)
         self.mlp = nn.Sequential(
-            nn.Linear(dim, dim * 4),
+            nn.Linear(dim, dim * 4, device=device, dtype=dtype),
             nn.GELU(approximate="tanh"),
         )
-        self.proj_out = nn.Linear(dim * 5, dim)
+        self.proj_out = nn.Linear(dim * 5, dim, device=device, dtype=dtype)
 
     def forward(self, x, t_emb, rope_emb, image_emb=None):
         h, gate = self.norm(x, emb=t_emb)

{diffsynth_engine-0.6.0 → diffsynth_engine-0.6.1.dev2}/diffsynth_engine/models/wan/wan_text_encoder.py

@@ -198,22 +198,22 @@ class WanTextEncoderStateDictConverter(StateDictConverter):
 
     def _from_diffusers(self, state_dict):
         rename_dict = {
-            "enc.output_norm.weight": "norm.weight",
-            "token_embd.weight": "token_embedding.weight",
+            "shared.weight": "token_embedding.weight",
+            "encoder.final_layer_norm.weight": "norm.weight",
         }
         for i in range(self.num_encoder_layers):
             rename_dict.update(
                 {
-                    f"enc.blk.{i}.attn_q.weight": f"blocks.{i}.attn.q.weight",
-                    f"enc.blk.{i}.attn_k.weight": f"blocks.{i}.attn.k.weight",
-                    f"enc.blk.{i}.attn_v.weight": f"blocks.{i}.attn.v.weight",
-                    f"enc.blk.{i}.attn_o.weight": f"blocks.{i}.attn.o.weight",
-                    f"enc.blk.{i}.ffn_up.weight": f"blocks.{i}.ffn.fc1.weight",
-                    f"enc.blk.{i}.ffn_down.weight": f"blocks.{i}.ffn.fc2.weight",
-                    f"enc.blk.{i}.ffn_gate.weight": f"blocks.{i}.ffn.gate.0.weight",
-                    f"enc.blk.{i}.attn_norm.weight": f"blocks.{i}.norm1.weight",
-                    f"enc.blk.{i}.ffn_norm.weight": f"blocks.{i}.norm2.weight",
-                    f"enc.blk.{i}.attn_rel_b.weight": f"blocks.{i}.pos_embedding.embedding.weight",
+                    f"encoder.block.{i}.layer.0.SelfAttention.q.weight": f"blocks.{i}.attn.q.weight",
+                    f"encoder.block.{i}.layer.0.SelfAttention.k.weight": f"blocks.{i}.attn.k.weight",
+                    f"encoder.block.{i}.layer.0.SelfAttention.v.weight": f"blocks.{i}.attn.v.weight",
+                    f"encoder.block.{i}.layer.0.SelfAttention.o.weight": f"blocks.{i}.attn.o.weight",
+                    f"encoder.block.{i}.layer.0.SelfAttention.relative_attention_bias.weight": f"blocks.{i}.pos_embedding.embedding.weight",
+                    f"encoder.block.{i}.layer.0.layer_norm.weight": f"blocks.{i}.norm1.weight",
+                    f"encoder.block.{i}.layer.1.DenseReluDense.wi_0.weight": f"blocks.{i}.ffn.gate.0.weight",
+                    f"encoder.block.{i}.layer.1.DenseReluDense.wi_1.weight": f"blocks.{i}.ffn.fc1.weight",
+                    f"encoder.block.{i}.layer.1.DenseReluDense.wo.weight": f"blocks.{i}.ffn.fc2.weight",
+                    f"encoder.block.{i}.layer.1.layer_norm.weight": f"blocks.{i}.norm2.weight",
                 }
             )
 
@@ -224,7 +224,7 @@ class WanTextEncoderStateDictConverter(StateDictConverter):
         return new_state_dict
 
     def convert(self, state_dict):
-        if "enc.output_norm.weight" in state_dict:
+        if "encoder.final_layer_norm.weight" in state_dict:
             logger.info("use diffusers format state dict")
             return self._from_diffusers(state_dict)
         return state_dict

{diffsynth_engine-0.6.0 → diffsynth_engine-0.6.1.dev2}/diffsynth_engine.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.6.0
+Version: 0.6.1.dev2
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent