"""Explicit local construction and decoding for the Qwen Image VAE."""
from __future__ import annotations
from pathlib import Path
import re
import torch
from diffusion_cli.checkpoint import loadStateDict
from diffusion_cli.config import ModelSource
from diffusion_cli.errors import DiffusionCliError
QWEN_IMAGE_LATENT_MEAN = (
-0.7571,
-0.7089,
-0.9113,
0.1075,
-0.1745,
0.9653,
-0.1517,
1.5508,
0.4134,
-0.0715,
0.5517,
-0.3632,
-0.1922,
-0.9497,
0.2503,
-0.2921,
)
QWEN_IMAGE_LATENT_STD = (
2.8184,
1.4541,
2.3275,
2.6558,
1.2196,
1.7708,
2.6052,
2.0743,
3.2687,
2.1526,
2.8652,
1.5579,
1.6382,
1.1253,
2.8251,
1.9160,
)
def _residualKey(key: str) -> str:
"""Convert an original residual block's sequential module names."""
return (
key.replace(".residual.0.gamma", ".norm1.gamma")
.replace(".residual.2.bias", ".conv1.bias")
.replace(".residual.2.weight", ".conv1.weight")
.replace(".residual.3.gamma", ".norm2.gamma")
.replace(".residual.6.bias", ".conv2.bias")
.replace(".residual.6.weight", ".conv2.weight")
.replace(".shortcut.bias", ".conv_shortcut.bias")
.replace(".shortcut.weight", ".conv_shortcut.weight")
)
def convertWanVaeToDiffusers(
state_dict: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
"""Convert original Wan/Qwen VAE keys to Diffusers Qwen Image keys.
This mapping is attributed to Diffusers' ``convert_wan_vae_to_diffusers``
behavior and is maintained here so runtime loading never imports that
private helper.
"""
converted: dict[str, torch.Tensor] = {}
middle_mapping = {
"encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma",
"encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias",
"encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight",
"encoder.middle.0.residual.3.gamma": "encoder.mid_block.resnets.0.norm2.gamma",
"encoder.middle.0.residual.6.bias": "encoder.mid_block.resnets.0.conv2.bias",
"encoder.middle.0.residual.6.weight": "encoder.mid_block.resnets.0.conv2.weight",
"encoder.middle.2.residual.0.gamma": "encoder.mid_block.resnets.1.norm1.gamma",
"encoder.middle.2.residual.2.bias": "encoder.mid_block.resnets.1.conv1.bias",
"encoder.middle.2.residual.2.weight": "encoder.mid_block.resnets.1.conv1.weight",
"encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma",
"encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias",
"encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight",
"decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma",
"decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias",
"decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight",
"decoder.middle.0.residual.3.gamma": "decoder.mid_block.resnets.0.norm2.gamma",
"decoder.middle.0.residual.6.bias": "decoder.mid_block.resnets.0.conv2.bias",
"decoder.middle.0.residual.6.weight": "decoder.mid_block.resnets.0.conv2.weight",
"decoder.middle.2.residual.0.gamma": "decoder.mid_block.resnets.1.norm1.gamma",
"decoder.middle.2.residual.2.bias": "decoder.mid_block.resnets.1.conv1.bias",
"decoder.middle.2.residual.2.weight": "decoder.mid_block.resnets.1.conv1.weight",
"decoder.middle.2.residual.3.gamma": "decoder.mid_block.resnets.1.norm2.gamma",
"decoder.middle.2.residual.6.bias": "decoder.mid_block.resnets.1.conv2.bias",
"decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight",
}
attention_mapping = {
"encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma",
"encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight",
"encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias",
"encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight",
"encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias",
"decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma",
"decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight",
"decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias",
"decoder.middle.1.proj.weight": "decoder.mid_block.attentions.0.proj.weight",
"decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias",
}
head_mapping = {
"encoder.head.0.gamma": "encoder.norm_out.gamma",
"encoder.head.2.bias": "encoder.conv_out.bias",
"encoder.head.2.weight": "encoder.conv_out.weight",
"decoder.head.0.gamma": "decoder.norm_out.gamma",
"decoder.head.2.bias": "decoder.conv_out.bias",
"decoder.head.2.weight": "decoder.conv_out.weight",
"conv1.weight": "quant_conv.weight",
"conv1.bias": "quant_conv.bias",
"conv2.weight": "post_quant_conv.weight",
"conv2.bias": "post_quant_conv.bias",
}
for key, value in state_dict.items():
if key in middle_mapping:
converted[middle_mapping[key]] = value
elif key in attention_mapping:
converted[attention_mapping[key]] = value
elif key in head_mapping:
converted[head_mapping[key]] = value
elif key == "encoder.conv1.weight":
converted["encoder.conv_in.weight"] = value
elif key == "encoder.conv1.bias":
converted["encoder.conv_in.bias"] = value
elif key == "decoder.conv1.weight":
converted["decoder.conv_in.weight"] = value
elif key == "decoder.conv1.bias":
converted["decoder.conv_in.bias"] = value
elif key.startswith("encoder.downsamples."):
new_key = key.replace("encoder.downsamples.", "encoder.down_blocks.")
converted[_residualKey(new_key)] = value
elif key.startswith("decoder.upsamples."):
new_key = _convertDecoderUpsampleKey(key)
converted[new_key] = value
else:
converted[key] = value
return converted
def _convertDecoderUpsampleKey(key: str) -> str:
"""Convert one decoder upsample family key."""
match = re.match(r"decoder\.upsamples\.(\d+)\.(.*)", key)
if match is None:
return key
block_index = int(match.group(1))
suffix = match.group(2)
if suffix.startswith("residual.") or ".residual." in suffix:
groups = {0: (0, 0), 1: (0, 1), 2: (0, 2), 4: (1, 0), 5: (1, 1),
6: (1, 2), 8: (2, 0), 9: (2, 1), 10: (2, 2),
12: (3, 0), 13: (3, 1), 14: (3, 2)}
if block_index in groups:
up, residual = groups[block_index]
residual_suffix = suffix.removeprefix("residual.")
converted_suffix = _residualKey(
f"block.residual.{residual_suffix}"
).removeprefix("block.")
return f"decoder.up_blocks.{up}.resnets.{residual}.{converted_suffix}"
if suffix.startswith("shortcut.") or ".shortcut." in suffix:
if block_index == 4:
suffix = suffix.replace("shortcut.", "resnets.0.conv_shortcut.")
return f"decoder.up_blocks.1.{suffix}"
suffix = suffix.replace("shortcut.", "conv_shortcut.")
return f"decoder.up_blocks.{block_index}.{suffix}"
if suffix.startswith("resample.") or suffix.startswith("time_conv.") \
or ".resample." in suffix or ".time_conv." in suffix:
block = {3: 0, 7: 1, 11: 2}.get(block_index, block_index)
return f"decoder.up_blocks.{block}.upsamplers.0.{suffix}"
return f"decoder.up_blocks.{block_index}.{suffix}"
def normalizeQwenVaeStateDict(
state_dict: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
"""Remove common component prefixes before conversion."""
normalized = {}
for key, value in state_dict.items():
if key.startswith("vae."):
key = key.removeprefix("vae.")
normalized[key] = value
return normalized
class QwenImageVae:
"""Decoder-only wrapper around an explicitly constructed Qwen Image VAE."""
compression = 8
channels = 16
def __init__(
self,
model_path: ModelSource | Path,
device,
dtype=torch.bfloat16,
*,
model=None,
) -> None:
self.device = device
self.dtype = dtype
if model is not None:
self.model = model
else:
self.model = self._loadModel(model_path)
self.model.eval()
def decode(self, latent: torch.Tensor) -> torch.Tensor:
"""Decode an NCHW normalized latent into an RGB [0, 1] tensor."""
if latent.ndim != 4 or latent.shape[1] != self.channels:
raise DiffusionCliError("Qwen Image VAE expects NCHW 16-channel latent")
if not torch.isfinite(latent).all():
raise DiffusionCliError("Qwen Image VAE input contains NaN or Inf")
mean = torch.tensor(
QWEN_IMAGE_LATENT_MEAN,
device=latent.device,
dtype=latent.dtype,
).view(1, self.channels, 1, 1, 1)
std = torch.tensor(
QWEN_IMAGE_LATENT_STD,
device=latent.device,
dtype=latent.dtype,
).view(1, self.channels, 1, 1, 1)
value = (latent.unsqueeze(2) * std) + mean
with torch.inference_mode():
result = self.model.decode(value)
result = result.sample if hasattr(result, "sample") else result
if result.ndim != 5 or result.shape[1] != 3 or result.shape[2] != 1:
raise DiffusionCliError("Qwen Image VAE returned an invalid RGB shape")
result = result[:, :, 0]
result = result.add(1.0).div(2.0).clamp(0.0, 1.0)
if not torch.isfinite(result).all():
raise DiffusionCliError("Qwen Image VAE output contains NaN or Inf")
return result
def toDevice(self, device, dtype=None) -> None:
"""Move the VAE to the active device."""
self.device = device
if dtype is not None:
self.dtype = dtype
self.model.to(device=device, dtype=self.dtype)
self.model.eval()
def toCpu(self) -> None:
"""Move the VAE back to CPU."""
self.toDevice(torch.device("cpu"))
def _loadModel(self, model_path: ModelSource | Path):
"""Construct and strictly load the local VAE without Hub loaders."""
try:
from diffusers import AutoencoderKLQwenImage
except ImportError as exc:
raise DiffusionCliError(
"diffusers is required for the Krea Qwen Image VAE"
) from exc
source = model_path if isinstance(model_path, ModelSource) else ModelSource(
model_path,
"vae",
)
raw_state = normalizeQwenVaeStateDict(loadStateDict(source))
if any(
value.dtype in {torch.float8_e4m3fn, torch.float8_e5m2}
for value in raw_state.values()
):
raise DiffusionCliError(
"Unsupported FP8 convolution in Qwen Image VAE"
)
converted = convertWanVaeToDiffusers(raw_state)
model = AutoencoderKLQwenImage(
base_dim=96,
z_dim=16,
dim_mult=[1, 2, 4, 4],
num_res_blocks=2,
attn_scales=[],
temperal_downsample=[False, True, True],
dropout=0.0,
input_channels=3,
latents_mean=list(QWEN_IMAGE_LATENT_MEAN),
latents_std=list(QWEN_IMAGE_LATENT_STD),
)
try:
model.load_state_dict(converted, strict=True)
except RuntimeError as exc:
raise DiffusionCliError(
f"Qwen Image VAE conversion left unexpected keys: {exc}"
) from exc
model.to(device=self.device, dtype=self.dtype)
return model
loadQwenImageVae = QwenImageVae