Mister-Link
466e2b5734
Reproduces app.py's TeleStyleV2 style-transfer pipeline as native ComfyUI nodes (Qwen TeleStyle Loader + Qwen TeleStyle Apply), loading weights directly from ComfyUI's own model folders instead of bridging live MODEL/CLIP/VAE objects (which doubled memory and OOM'd on this box).
199 lines
8.3 KiB
Python
199 lines
8.3 KiB
Python
"""Convert a ComfyUI-native Wan-style Qwen-Image VAE state dict (comfy.ldm.wan.vae.WanVAE,
|
|
keys like `encoder.downsamples.0.residual.0.gamma`) into the key layout expected by
|
|
diffusers' AutoencoderKLQwenImage (keys like `encoder.down_blocks.0.norm1.gamma`).
|
|
|
|
Both classes are literal re-implementations of the same architecture (diffusers'
|
|
own docstring: "QwenImageVAE is further fine-tuned from the Wan Video VAE"), so the
|
|
mapping below mirrors each class's construction loop exactly rather than hardcoding
|
|
per-key tables. Every call is validated against the real diffusers module's own
|
|
state_dict() keys/shapes before use, so any topology mismatch (e.g. an unexpected
|
|
config) raises loudly instead of silently loading misaligned weights.
|
|
"""
|
|
|
|
RESIDUAL_SUBKEYS = {
|
|
"residual.0": "norm1",
|
|
"residual.2": "conv1",
|
|
"residual.3": "norm2",
|
|
"residual.6": "conv2",
|
|
}
|
|
|
|
|
|
def _flat_encoder_layout(dim_mult, num_res_blocks, attn_scales):
|
|
"""Reproduce Encoder3d/QwenImageEncoder3d's flat down-path construction.
|
|
|
|
Returns a list of ('res'|'attn'|'down', flat_index) — identical ordering in both
|
|
comfy's `downsamples` Sequential and diffusers' `down_blocks` ModuleList, since
|
|
both are built by the same loop over the same config.
|
|
"""
|
|
layout = []
|
|
scale = 1.0
|
|
n_stages = len(dim_mult)
|
|
for i in range(n_stages):
|
|
for _ in range(num_res_blocks):
|
|
layout.append("res")
|
|
if scale in attn_scales:
|
|
layout.append("attn")
|
|
if i != n_stages - 1:
|
|
layout.append("down")
|
|
scale /= 2.0
|
|
return layout
|
|
|
|
|
|
def _flat_decoder_stage_layout(dim_mult, num_res_blocks, attn_scales):
|
|
"""Reproduce Decoder3d/QwenImageDecoder3d's up-path construction.
|
|
|
|
Returns a list of stages; each stage is a list of ('res'|'attn') entries plus a
|
|
trailing 'up' marker if that stage upsamples. comfy's `upsamples` Sequential is the
|
|
flat concatenation of all stages; diffusers' `up_blocks` is one QwenImageUpBlock per
|
|
stage (its own `.resnets` / `.upsamplers`).
|
|
"""
|
|
stages = []
|
|
scale = 1.0 / 2 ** (len(dim_mult) - 2)
|
|
n_stages = len(dim_mult)
|
|
for i in range(n_stages):
|
|
stage = []
|
|
for _ in range(num_res_blocks + 1):
|
|
stage.append("res")
|
|
if scale in attn_scales:
|
|
stage.append("attn")
|
|
if i != n_stages - 1:
|
|
stage.append("up")
|
|
scale *= 2.0
|
|
stages.append(stage)
|
|
return stages
|
|
|
|
|
|
def _map_residual_block(comfy_prefix, diffusers_prefix, mapping):
|
|
for comfy_sub, diffusers_sub in RESIDUAL_SUBKEYS.items():
|
|
mapping[f"{comfy_prefix}.{comfy_sub}"] = f"{diffusers_prefix}.{diffusers_sub}"
|
|
mapping[f"{comfy_prefix}.shortcut"] = f"{diffusers_prefix}.conv_shortcut"
|
|
|
|
|
|
def _map_attention_block(comfy_prefix, diffusers_prefix, mapping):
|
|
mapping[f"{comfy_prefix}.norm"] = f"{diffusers_prefix}.norm"
|
|
mapping[f"{comfy_prefix}.to_qkv"] = f"{diffusers_prefix}.to_qkv"
|
|
mapping[f"{comfy_prefix}.proj"] = f"{diffusers_prefix}.proj"
|
|
|
|
|
|
def _map_resample(comfy_prefix, diffusers_prefix, mapping):
|
|
mapping[f"{comfy_prefix}.resample"] = f"{diffusers_prefix}.resample"
|
|
mapping[f"{comfy_prefix}.time_conv"] = f"{diffusers_prefix}.time_conv"
|
|
|
|
|
|
def _map_middle(comfy_prefix, diffusers_prefix, mapping):
|
|
# comfy: Sequential[ResidualBlock, AttentionBlock, ResidualBlock]
|
|
# diffusers: QwenImageMidBlock(resnets=[ResidualBlock, ResidualBlock], attentions=[AttentionBlock])
|
|
_map_residual_block(f"{comfy_prefix}.0", f"{diffusers_prefix}.resnets.0", mapping)
|
|
_map_attention_block(f"{comfy_prefix}.1", f"{diffusers_prefix}.attentions.0", mapping)
|
|
_map_residual_block(f"{comfy_prefix}.2", f"{diffusers_prefix}.resnets.1", mapping)
|
|
|
|
|
|
def build_prefix_map(dim_mult=(1, 2, 4, 4), num_res_blocks=2, attn_scales=()):
|
|
"""Build a dict of module-path prefixes: comfy WanVAE -> diffusers AutoencoderKLQwenImage."""
|
|
mapping = {}
|
|
|
|
# top level
|
|
mapping["conv1"] = "quant_conv"
|
|
mapping["conv2"] = "post_quant_conv"
|
|
|
|
# --- encoder ---
|
|
mapping["encoder.conv1"] = "encoder.conv_in"
|
|
enc_layout = _flat_encoder_layout(dim_mult, num_res_blocks, attn_scales)
|
|
res_i = attn_i = 0
|
|
for flat_idx, kind in enumerate(enc_layout):
|
|
comfy_prefix = f"encoder.downsamples.{flat_idx}"
|
|
diffusers_prefix = f"encoder.down_blocks.{flat_idx}"
|
|
if kind == "res":
|
|
_map_residual_block(comfy_prefix, diffusers_prefix, mapping)
|
|
res_i += 1
|
|
elif kind == "attn":
|
|
_map_attention_block(comfy_prefix, diffusers_prefix, mapping)
|
|
attn_i += 1
|
|
elif kind == "down":
|
|
_map_resample(comfy_prefix, diffusers_prefix, mapping)
|
|
_map_middle("encoder.middle", "encoder.mid_block", mapping)
|
|
mapping["encoder.head.0"] = "encoder.norm_out"
|
|
mapping["encoder.head.2"] = "encoder.conv_out"
|
|
|
|
# --- decoder ---
|
|
mapping["decoder.conv1"] = "decoder.conv_in"
|
|
_map_middle("decoder.middle", "decoder.mid_block", mapping)
|
|
dec_stages = _flat_decoder_stage_layout(dim_mult, num_res_blocks, attn_scales)
|
|
flat_idx = 0
|
|
for stage_idx, stage in enumerate(dec_stages):
|
|
res_j = 0
|
|
for kind in stage:
|
|
comfy_prefix = f"decoder.upsamples.{flat_idx}"
|
|
if kind == "res":
|
|
diffusers_prefix = f"decoder.up_blocks.{stage_idx}.resnets.{res_j}"
|
|
_map_residual_block(comfy_prefix, diffusers_prefix, mapping)
|
|
res_j += 1
|
|
elif kind == "attn":
|
|
# diffusers QwenImageUpBlock has no attention support; only reachable if
|
|
# attn_scales is non-empty, which real Qwen-Image-Edit VAE configs don't use.
|
|
raise NotImplementedError(
|
|
"VAE config has attn_scales set in the decoder up-path; this converter "
|
|
"was only validated against the real Qwen-Image-Edit config (attn_scales=[])."
|
|
)
|
|
elif kind == "up":
|
|
diffusers_prefix = f"decoder.up_blocks.{stage_idx}.upsamplers.0"
|
|
_map_resample(comfy_prefix, diffusers_prefix, mapping)
|
|
flat_idx += 1
|
|
mapping["decoder.head.0"] = "decoder.norm_out"
|
|
mapping["decoder.head.2"] = "decoder.conv_out"
|
|
|
|
return mapping
|
|
|
|
|
|
def convert_state_dict(comfy_sd, dim_mult=(1, 2, 4, 4), num_res_blocks=2, attn_scales=()):
|
|
"""Convert a comfy WanVAE state dict into diffusers AutoencoderKLQwenImage key layout."""
|
|
prefix_map = build_prefix_map(dim_mult, num_res_blocks, attn_scales)
|
|
# longest-prefix-first so e.g. "encoder.downsamples.10.residual.0" doesn't get
|
|
# accidentally matched against a shorter unrelated prefix
|
|
sorted_prefixes = sorted(prefix_map.keys(), key=len, reverse=True)
|
|
|
|
out = {}
|
|
unmatched = []
|
|
for key, tensor in comfy_sd.items():
|
|
matched = False
|
|
for comfy_prefix in sorted_prefixes:
|
|
if key == comfy_prefix or key.startswith(comfy_prefix + "."):
|
|
suffix = key[len(comfy_prefix):]
|
|
out[prefix_map[comfy_prefix] + suffix] = tensor
|
|
matched = True
|
|
break
|
|
if not matched:
|
|
unmatched.append(key)
|
|
|
|
if unmatched:
|
|
raise RuntimeError(
|
|
f"vae_convert: {len(unmatched)} comfy VAE keys had no mapping, e.g. {unmatched[:5]}"
|
|
)
|
|
return out
|
|
|
|
|
|
def convert_and_validate(comfy_sd, target_model, dim_mult=(1, 2, 4, 4), num_res_blocks=2, attn_scales=()):
|
|
"""Convert and check the result exactly matches target_model's own state_dict keys/shapes."""
|
|
converted = convert_state_dict(comfy_sd, dim_mult, num_res_blocks, attn_scales)
|
|
target_sd = target_model.state_dict()
|
|
|
|
converted_keys = set(converted.keys())
|
|
target_keys = set(target_sd.keys())
|
|
missing = target_keys - converted_keys
|
|
extra = converted_keys - target_keys
|
|
if missing or extra:
|
|
raise RuntimeError(
|
|
"vae_convert: converted VAE state dict doesn't match diffusers model exactly. "
|
|
f"missing={sorted(missing)[:5]} extra={sorted(extra)[:5]}"
|
|
)
|
|
|
|
bad_shapes = [
|
|
(k, tuple(converted[k].shape), tuple(target_sd[k].shape))
|
|
for k in converted_keys
|
|
if tuple(converted[k].shape) != tuple(target_sd[k].shape)
|
|
]
|
|
if bad_shapes:
|
|
raise RuntimeError(f"vae_convert: shape mismatches after conversion, e.g. {bad_shapes[:5]}")
|
|
|
|
return converted
|