diff --git a/comfy/lora.py b/comfy/lora.py index f36ddb046..76f5eb0a9 100644 --- a/comfy/lora.py +++ b/comfy/lora.py @@ -426,9 +426,11 @@ def calculate_weight(patches, weight, key, intermediate_dtype=torch.float32, ori if output is None: logging.warning("Calculate Weight Failed: {} {}".format(v.name, key)) else: - weight = output if old_weight is not None: + old_weight.narrow(offset[0], offset[1], offset[2]).copy_(output) weight = old_weight + else: + weight = output continue if len(v) == 1: diff --git a/comfy/model_detection.py b/comfy/model_detection.py index 3faa950ca..ebccc2b46 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -1064,6 +1064,12 @@ def convert_diffusers_mmdit(state_dict, output_prefix=""): num_blocks = count_blocks(state_dict, 'transformer_blocks.{}.') depth = state_dict["pos_embed.proj.weight"].shape[0] // 64 sd_map = comfy.utils.mmdit_to_diffusers({"depth": depth, "num_blocks": num_blocks}, output_prefix=output_prefix) + elif 'all_x_embedder.2-1.weight' in state_dict: #Z-Image (HuggingFace format) + w = state_dict.get('cap_embedder.1.weight') + hidden_size = w.shape[0] if w is not None else 3840 + n_layers = count_blocks(state_dict, 'layers.{}.') + n_refiner = count_blocks(state_dict, 'noise_refiner.{}.') + sd_map = comfy.utils.z_image_to_diffusers({"n_layers": n_layers, "n_refiner_layers": n_refiner, "dim": hidden_size}, output_prefix=output_prefix) else: return None diff --git a/comfy/sd.py b/comfy/sd.py index a9ad7c2d2..a2e8a3704 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -1713,6 +1713,12 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None, disable if model_config is not None: new_sd = sd + # Z-Image from HuggingFace uses diffusers-style key names that need conversion + if model_config.unet_config.get('z_image_modulation', False) and 'all_x_embedder.2-1.weight' in new_sd: + sd_copy = dict(new_sd) + converted = model_detection.convert_diffusers_mmdit(sd_copy, "") + if converted is not None: + new_sd = {**sd_copy, **converted} else: new_sd = model_detection.convert_diffusers_mmdit(sd, "") if new_sd is not None: #diffusers mmdit