102 lines
4.4 KiB
Python
102 lines
4.4 KiB
Python
import torch
|
|
from diffusers.models.transformers import (
|
|
ZImageTransformer2DModel as DiffusersZImageTransformer2DModel,
|
|
)
|
|
|
|
from ._mixin import OstrisModelMixin
|
|
|
|
|
|
class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMixin):
|
|
aitk_subfolder = "transformer"
|
|
# repo to pull the config from when loading a single-file checkpoint
|
|
aitk_config_repo = "Tongyi-MAI/Z-Image-Turbo"
|
|
|
|
@classmethod
|
|
def get_quantization_block_names(cls):
|
|
return ["layers"]
|
|
|
|
@classmethod
|
|
def get_quantization_exclude_modules(cls):
|
|
# sensitive modules kept in full precision (fnmatch patterns on module
|
|
# names within ZImageTransformer2DModel):
|
|
# t_embedder* - timestep embedder; feeds every block's
|
|
# adaLN_modulation and the final layers
|
|
# cap_embedder* - caption feature -> model width projection
|
|
# all_x_embedder* - patchified latent input projections
|
|
# all_final_layer* - final adaLN-modulated output projections
|
|
# siglip_embedder* - siglip feature projection (edit models only)
|
|
return [
|
|
"t_embedder*",
|
|
"cap_embedder*",
|
|
"all_x_embedder*",
|
|
"all_final_layer*",
|
|
"siglip_embedder*",
|
|
]
|
|
|
|
@classmethod
|
|
def convert_state_dict_on_load(cls, state_dict):
|
|
"""Convert a single-file Z-Image checkpoint to diffusers transformer keys."""
|
|
new_sd = {}
|
|
for key, value in state_dict.items():
|
|
k = key
|
|
if k.endswith(".attention.qkv.weight"):
|
|
# the single file fuses q,k,v into one tensor (in that order); diffusers keeps them split
|
|
prefix = k[: -len(".attention.qkv.weight")]
|
|
q, k_proj, v = torch.chunk(value, 3, dim=0)
|
|
new_sd[prefix + ".attention.to_q.weight"] = q
|
|
new_sd[prefix + ".attention.to_k.weight"] = k_proj
|
|
new_sd[prefix + ".attention.to_v.weight"] = v
|
|
continue
|
|
k = k.replace(".attention.out.weight", ".attention.to_out.0.weight")
|
|
k = k.replace(".attention.q_norm.weight", ".attention.norm_q.weight")
|
|
k = k.replace(".attention.k_norm.weight", ".attention.norm_k.weight")
|
|
if k.startswith("x_embedder."):
|
|
k = "all_x_embedder.2-1." + k[len("x_embedder.") :]
|
|
elif k.startswith("final_layer."):
|
|
k = "all_final_layer.2-1." + k[len("final_layer.") :]
|
|
new_sd[k] = value
|
|
return new_sd
|
|
|
|
@classmethod
|
|
def convert_state_dict_on_save(cls, state_dict):
|
|
"""Convert a diffusers transformer state dict back to the single-file layout."""
|
|
new_sd = {}
|
|
qkv_cache = {}
|
|
for key, value in state_dict.items():
|
|
k = key
|
|
matched = False
|
|
for suffix in (
|
|
".attention.to_q.weight",
|
|
".attention.to_k.weight",
|
|
".attention.to_v.weight",
|
|
):
|
|
if k.endswith(suffix):
|
|
prefix = k[: -len(suffix)]
|
|
cache = qkv_cache.setdefault(prefix, {})
|
|
cache[suffix] = value
|
|
if len(cache) == 3:
|
|
# the single file expects q,k,v fused in that order
|
|
qkv = torch.cat(
|
|
[
|
|
cache[".attention.to_q.weight"],
|
|
cache[".attention.to_k.weight"],
|
|
cache[".attention.to_v.weight"],
|
|
],
|
|
dim=0,
|
|
)
|
|
new_sd[prefix + ".attention.qkv.weight"] = qkv
|
|
del qkv_cache[prefix]
|
|
matched = True
|
|
break
|
|
if matched:
|
|
continue
|
|
k = k.replace(".attention.to_out.0.weight", ".attention.out.weight")
|
|
k = k.replace(".attention.norm_q.weight", ".attention.q_norm.weight")
|
|
k = k.replace(".attention.norm_k.weight", ".attention.k_norm.weight")
|
|
if k.startswith("all_x_embedder.2-1."):
|
|
k = "x_embedder." + k[len("all_x_embedder.2-1.") :]
|
|
elif k.startswith("all_final_layer.2-1."):
|
|
k = "final_layer." + k[len("all_final_layer.2-1.") :]
|
|
new_sd[k] = value
|
|
return new_sd
|