Spaces:
Running on Zero
Running on Zero
| """The Ming-Image VAE: the Qwen-Image (Wan-style video) VAE with four image | |
| channels. The ComfyUI repack ships it in Comfy's Wan-native key layout, which | |
| the load converter maps onto the diffusers module names.""" | |
| from toolkit.models.v2.vae.qwen_image import QwenImageVAE | |
| from .checkpoints import BASE_REPO, COMFY_REPO, COMFY_VAE_FILES, comfy_weight_names | |
| # comfy `<block>.residual.<n>` / `shortcut` -> the diffusers resnet submodule | |
| _RESNET_PARTS = { | |
| "residual.0": "norm1", | |
| "residual.2": "conv1", | |
| "residual.3": "norm2", | |
| "residual.6": "conv2", | |
| "shortcut": "conv_shortcut", | |
| } | |
| # comfy's `middle` Sequential is resnet, attention, resnet | |
| _MID_PARTS = {"0": "resnets.0", "1": "attentions.0", "2": "resnets.1"} | |
| def _resnet_suffix(inner): | |
| consumed = 2 if inner[0] == "residual" else 1 | |
| return [_RESNET_PARTS[".".join(inner[:consumed])]] + inner[consumed:] | |
| def comfy_to_diffusers_key(key: str, up_block_len: int) -> str: | |
| """One ComfyUI (Wan-native) VAE parameter name -> its diffusers name. | |
| `up_block_len` is the number of decoder entries per up block (its resnets | |
| plus the upsampler).""" | |
| parts = key.split(".") | |
| if parts[0] == "conv1": | |
| return ".".join(["quant_conv"] + parts[1:]) | |
| if parts[0] == "conv2": | |
| return ".".join(["post_quant_conv"] + parts[1:]) | |
| side, rest = parts[0], parts[1:] | |
| if rest[0] == "conv1": | |
| return ".".join([side, "conv_in"] + rest[1:]) | |
| if rest[0] == "head": | |
| # head.0 is the output norm, head.2 the output conv (head.1 is SiLU) | |
| tail = "norm_out" if rest[1] == "0" else "conv_out" | |
| return ".".join([side, tail] + rest[2:]) | |
| if rest[0] == "middle": | |
| part, inner = _MID_PARTS[rest[1]], rest[2:] | |
| if part.startswith("resnets"): | |
| inner = _resnet_suffix(inner) | |
| return ".".join([side, "mid_block", part] + inner) | |
| index, inner = int(rest[1]), rest[2:] | |
| if side == "encoder": | |
| # encoder.downsamples.<i>: a flat list of resnets and resamplers in | |
| # the same order as diffusers' down_blocks | |
| if inner[0] in ("resample", "time_conv"): | |
| return ".".join([side, "down_blocks", str(index)] + inner) | |
| return ".".join([side, "down_blocks", str(index)] + _resnet_suffix(inner)) | |
| # decoder.upsamples.<i> is flat too; diffusers nests each up block's | |
| # resnets under it with the upsampler last | |
| block, j = divmod(index, up_block_len) | |
| if inner[0] in ("resample", "time_conv"): | |
| return ".".join([side, "up_blocks", str(block), "upsamplers", "0"] + inner) | |
| return ".".join([side, "up_blocks", str(block), "resnets", str(j)] + _resnet_suffix(inner)) | |
| def convert_comfy_vae_state_dict(state_dict): | |
| if "encoder.conv1.weight" not in state_dict: | |
| return state_dict | |
| # the first decoder resampler sits right after a block's resnets | |
| first_resample = min( | |
| int(k.split(".")[2]) | |
| for k in state_dict | |
| if k.startswith("decoder.upsamples.") and ".resample." in k | |
| ) | |
| up_block_len = first_resample + 1 | |
| return {comfy_to_diffusers_key(k, up_block_len): v for k, v in state_dict.items()} | |
| class MingImageVAE(QwenImageVAE): | |
| aitk_config_repo = BASE_REPO | |
| aitk_comfy_repo = COMFY_REPO | |
| aitk_comfy_weight_names = comfy_weight_names(COMFY_VAE_FILES) | |
| def convert_state_dict_on_load(cls, state_dict): | |
| return convert_comfy_vae_state_dict(state_dict) | |