Spaces:
Running on Zero
Running on Zero
File size: 3,428 Bytes
855e1f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """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)
@classmethod
def convert_state_dict_on_load(cls, state_dict):
return convert_comfy_vae_state_dict(state_dict)
|