Merge branch 'main' into cuda_replace

diffsynth/utils/state_dict_converters/flux_dit.py

@@ -143,6 +143,8 @@ def FluxDiTStateDictConverterFromDiffusers(state_dict):
             suffix = ".weight" if name.endswith(".weight") else ".bias"
             prefix = name[:-len(suffix)]
             if prefix in global_rename_dict:
+                if global_rename_dict[prefix] == "final_norm_out.linear":
+                    param = torch.concat([param[3072:], param[:3072]], dim=0)
                 state_dict_[global_rename_dict[prefix] + suffix] = param
             elif prefix.startswith("transformer_blocks."):
                 names = prefix.split(".")

diffsynth/utils/state_dict_converters/z_image_text_encoder.py

@@ -0,0 +1,6 @@
+def ZImageTextEncoderStateDictConverter(state_dict):
+    state_dict_ = {}
+    for name in state_dict:
+        if name != "lm_head.weight":
+            state_dict_[name] = state_dict[name]
+    return state_dict_

diffsynth/utils/xfuser/xdit_context_parallel.py

@@ -50,7 +50,7 @@ def rope_apply(x, freqs, num_heads):
     sp_rank = get_sequence_parallel_rank()
     freqs = pad_freqs(freqs, s_per_rank * sp_size)
     freqs_rank = freqs[(sp_rank * s_per_rank):((sp_rank + 1) * s_per_rank), :, :]
-
+    freqs_rank = freqs_rank.to(torch.complex64) if freqs_rank.device == "npu" else freqs_rank
     x_out = torch.view_as_real(x_out * freqs_rank).flatten(2)
     return x_out.to(x.dtype)