Move z_image over to the new modeling class
This commit is contained in:
Jaret Burkett
parent
dcb3b329b2
commit
28f2c0acbe
|
|
@ -29,7 +29,10 @@ from diffusers import AutoencoderKL
|
|||
|
||||
try:
|
||||
from diffusers import ZImagePipeline
|
||||
from diffusers.models.transformers import ZImageTransformer2DModel
|
||||
|
||||
# our subclass of the diffusers transformer with the universal loading /
|
||||
# quantization mixin (see toolkit/models/classes/_mixin.py)
|
||||
from toolkit.models.v2.z_image import ZImageTransformer2DModel
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Diffusers is out of date. Update diffusers to the latest version by doing pip uninstall diffusers and then pip install -r requirements.txt"
|
||||
|
|
@ -47,73 +50,6 @@ scheduler_config = {
|
|||
SINGLE_FILE_EXTRAS_REPO = "Tongyi-MAI/Z-Image-Turbo"
|
||||
|
||||
|
||||
def convert_single_file_to_diffusers(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
|
||||
|
||||
|
||||
def convert_diffusers_to_single_file(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
|
||||
|
||||
|
||||
class ZImageModel(BaseModel):
|
||||
arch = "zimage"
|
||||
|
||||
|
|
@ -229,9 +165,10 @@ class ZImageModel(BaseModel):
|
|||
self.invert_assistant_lora = True
|
||||
|
||||
def load_transformer(self, model_path, base_model_path, dtype):
|
||||
"""Load the ZImage transformer from either a diffusers folder/repo or a
|
||||
single-file checkpoint. Returns (transformer, base_model_path) since the base
|
||||
path may be redirected to the hub repo for single-file checkpoints."""
|
||||
"""Load the ZImage transformer through the OstrisModelMixin universal loader
|
||||
(diffusers folder/repo, local or hub single-file checkpoint, pre-quantized
|
||||
checkpoint). Returns (transformer, base_model_path) since the base path may
|
||||
be redirected to the hub repo for single-file checkpoints."""
|
||||
if model_path.endswith(".safetensors"):
|
||||
# single-file checkpoint. Load the weights from the file and pull the
|
||||
# vae / text encoder / tokenizer / config from the base diffusers repo.
|
||||
|
|
@ -242,36 +179,32 @@ class ZImageModel(BaseModel):
|
|||
if base_model_path == model_path:
|
||||
# extras default to name_or_path which is the single file, fall back to the hub repo
|
||||
base_model_path = self.single_file_extras_repo
|
||||
elif os.path.exists(model_path):
|
||||
# check if the path is a full checkpoint.
|
||||
te_folder_path = os.path.join(model_path, "text_encoder")
|
||||
# if we have the te, this folder is a full checkpoint, use it as the base
|
||||
if os.path.exists(te_folder_path):
|
||||
base_model_path = model_path
|
||||
|
||||
state_dict = load_file(model_path)
|
||||
state_dict = convert_single_file_to_diffusers(state_dict)
|
||||
for key, value in state_dict.items():
|
||||
state_dict[key] = value.to(dtype=dtype)
|
||||
# quantization happens inside load_model unless an adapter has to be merged
|
||||
# into the full precision weights first (assistant lora / accuracy recovery
|
||||
# adapter); those paths quantize after the merge via quantize_model
|
||||
qtype = None
|
||||
if (
|
||||
self.model_config.quantize
|
||||
and self.model_config.assistant_lora_path is None
|
||||
and self.model_config.accuracy_recovery_adapter is None
|
||||
):
|
||||
qtype = self.model_config.qtype
|
||||
|
||||
config = ZImageTransformer2DModel.load_config(
|
||||
base_model_path, subfolder="transformer"
|
||||
)
|
||||
with torch.device("meta"):
|
||||
transformer = ZImageTransformer2DModel.from_config(config)
|
||||
transformer.load_state_dict(state_dict, assign=True)
|
||||
transformer.to(dtype=dtype)
|
||||
del state_dict
|
||||
flush()
|
||||
else:
|
||||
transformer_path = model_path
|
||||
transformer_subfolder = "transformer"
|
||||
if os.path.exists(transformer_path):
|
||||
transformer_subfolder = None
|
||||
transformer_path = os.path.join(transformer_path, "transformer")
|
||||
# check if the path is a full checkpoint.
|
||||
te_folder_path = os.path.join(model_path, "text_encoder")
|
||||
# if we have the te, this folder is a full checkpoint, use it as the base
|
||||
if os.path.exists(te_folder_path):
|
||||
base_model_path = model_path
|
||||
|
||||
transformer = ZImageTransformer2DModel.from_pretrained(
|
||||
transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype
|
||||
)
|
||||
transformer = ZImageTransformer2DModel.load_model(
|
||||
model_path,
|
||||
dtype=dtype,
|
||||
qtype=qtype,
|
||||
quantize_device=self.device_torch,
|
||||
config_path=base_model_path if self.is_single_file else None,
|
||||
)
|
||||
flush()
|
||||
|
||||
return transformer, base_model_path
|
||||
|
||||
|
|
@ -294,7 +227,9 @@ class ZImageModel(BaseModel):
|
|||
if self.model_config.qtype == "qfloat8":
|
||||
self.model_config.qtype = "float8"
|
||||
|
||||
if self.model_config.quantize:
|
||||
# already quantized inside load_transformer unless an adapter had to merge
|
||||
# into full precision weights first (or the checkpoint was pre-quantized)
|
||||
if self.model_config.quantize and not transformer.aitk_is_quantized:
|
||||
self.print_and_status_update("Quantizing Transformer")
|
||||
quantize_model(self, transformer)
|
||||
flush()
|
||||
|
|
@ -423,10 +358,10 @@ class ZImageModel(BaseModel):
|
|||
sc = self.get_bucket_divisibility()
|
||||
gen_config.width = int(gen_config.width // sc * sc)
|
||||
gen_config.height = int(gen_config.height // sc * sc)
|
||||
|
||||
|
||||
# CFG is 0 normalized for this model
|
||||
guidance = max(0.0, gen_config.guidance_scale - 1.0)
|
||||
|
||||
|
||||
img = pipeline(
|
||||
prompt_embeds=conditional_embeds.text_embeds,
|
||||
negative_prompt_embeds=unconditional_embeds.text_embeds,
|
||||
|
|
@ -486,6 +421,11 @@ class ZImageModel(BaseModel):
|
|||
def get_te_has_grad(self):
|
||||
return False
|
||||
|
||||
def get_quantization_exclude_modules(self):
|
||||
# the patterns live on the transformer class so the mixin quantization
|
||||
# uses them too (see toolkit/models/classes/z_image.py)
|
||||
return ZImageTransformer2DModel.get_quantization_exclude_modules()
|
||||
|
||||
def save_model(self, output_path, meta, save_dtype):
|
||||
transformer: ZImageTransformer2DModel = unwrap_model(self.model)
|
||||
if self.is_single_file:
|
||||
|
|
@ -497,7 +437,7 @@ class ZImageModel(BaseModel):
|
|||
save_dict[key] = (
|
||||
dequantize_if_quantized(value).clone().to("cpu", dtype=save_dtype)
|
||||
)
|
||||
save_dict = convert_diffusers_to_single_file(save_dict)
|
||||
save_dict = transformer.convert_state_dict_on_save(save_dict)
|
||||
|
||||
if not output_path.endswith(".safetensors"):
|
||||
output_path += ".safetensors"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,272 @@
|
|||
"""
|
||||
Universal model mixin for every model class the toolkit touches (transformers,
|
||||
unets, text encoders, VAEs, ...).
|
||||
|
||||
OstrisModelMixin gives a model class one loading entry point, `load_model`, that
|
||||
digests any of:
|
||||
- a local diffusers checkpoint directory (the model folder itself or a full
|
||||
checkpoint that contains it under `aitk_subfolder`)
|
||||
- a HuggingFace repo id ("org/repo")
|
||||
- a local single .safetensors file in the model's original key layout
|
||||
- a remote single .safetensors file ("org/repo/path/file.safetensors")
|
||||
plus optional automatic quantization of the loaded weights (`qtype=...`).
|
||||
|
||||
Model specific behavior lives in small overridable hooks (key conversion, config
|
||||
source, block names, backend loading), so subclassing for a new model usually means
|
||||
setting the `aitk_*` class attrs and overriding one or two hooks. The default
|
||||
backend hooks (`aitk_from_pretrained` / `aitk_load_config` / `aitk_from_config`)
|
||||
speak the diffusers ModelMixin API; transformers-lib models (text encoders) override
|
||||
those three to speak PreTrainedModel/AutoConfig instead.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import torch
|
||||
from huggingface_hub import hf_hub_download
|
||||
from safetensors.torch import load_file
|
||||
from tqdm import tqdm
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin
|
||||
from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
|
||||
from toolkit.basic import flush
|
||||
|
||||
|
||||
class BasicModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
|
||||
pass
|
||||
|
||||
|
||||
class OstrisModelMixin:
|
||||
# ---- per-model configuration, override in subclasses ----
|
||||
# subfolder that holds this model inside a diffusers style checkpoint
|
||||
# (e.g. "transformer", "text_encoder", "vae")
|
||||
aitk_subfolder: Optional[str] = None
|
||||
# hub repo (or local path) to pull the model config from when loading a single
|
||||
# .safetensors file and no config_path is given
|
||||
aitk_config_repo: Optional[str] = None
|
||||
|
||||
# ---- state set by the loader / quantizer ----
|
||||
aitk_is_quantized: bool = False
|
||||
aitk_qtype: Optional[str] = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# extendable hooks
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@classmethod
|
||||
def convert_state_dict_on_load(cls, state_dict: Dict[str, torch.Tensor]):
|
||||
"""Convert a single-file state dict from the model's original key layout to
|
||||
this class's layout. Default is a passthrough."""
|
||||
return state_dict
|
||||
|
||||
@classmethod
|
||||
def convert_state_dict_on_save(cls, state_dict: Dict[str, torch.Tensor]):
|
||||
"""Convert this class's state dict back to the original single-file layout.
|
||||
Default is a passthrough."""
|
||||
return state_dict
|
||||
|
||||
@classmethod
|
||||
def get_quantization_block_names(cls) -> Optional[List[str]]:
|
||||
"""Names (dotted paths allowed) of the repeated block lists to quantize one
|
||||
block at a time so the whole model never has to sit on the gpu at once."""
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def get_quantization_exclude_modules(cls) -> Optional[List[str]]:
|
||||
"""fnmatch patterns of sensitive modules to keep in full precision."""
|
||||
return None
|
||||
|
||||
# ---- backend hooks: default to the diffusers ModelMixin API. transformers-lib
|
||||
# models (text encoders) override these three to use PreTrainedModel/AutoConfig.
|
||||
|
||||
@classmethod
|
||||
def aitk_from_pretrained(cls, path, subfolder=None, dtype=None, **kwargs):
|
||||
return cls.from_pretrained(
|
||||
path, subfolder=subfolder, torch_dtype=dtype, **kwargs
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def aitk_load_config(cls, path, subfolder=None):
|
||||
return cls.load_config(path, subfolder=subfolder)
|
||||
|
||||
@classmethod
|
||||
def aitk_from_config(cls, config):
|
||||
with torch.device("meta"):
|
||||
return cls.from_config(config)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# loading
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@classmethod
|
||||
def load_model(
|
||||
cls,
|
||||
name_or_path: str,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
device: Optional[torch.device] = None,
|
||||
qtype: Optional[str] = None,
|
||||
quantize_device: Optional[torch.device] = None,
|
||||
exclude_quant_modules: Optional[List[str]] = None,
|
||||
config_path: Optional[str] = None,
|
||||
subfolder: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Load a model universally from a given name or path.
|
||||
|
||||
name_or_path can be a local diffusers directory, a hub repo id, a local
|
||||
single .safetensors file, or a remote single file
|
||||
("org/repo/file.safetensors").
|
||||
|
||||
qtype: quantize the weights after loading. quantize_device: where to run the
|
||||
quantization math; blocks are moved there one at a time and returned to where
|
||||
they were.
|
||||
config_path: config source for single-file loads, overriding aitk_config_repo.
|
||||
device: move the finished model there before returning.
|
||||
"""
|
||||
if subfolder is None:
|
||||
subfolder = cls.aitk_subfolder
|
||||
|
||||
if name_or_path.endswith(".safetensors"):
|
||||
file_path = cls._resolve_single_file(name_or_path)
|
||||
model = cls._load_single_file(
|
||||
file_path, dtype=dtype, config_path=config_path, subfolder=subfolder
|
||||
)
|
||||
else:
|
||||
if os.path.isdir(name_or_path):
|
||||
# a local dir may be the model folder itself or a full checkpoint
|
||||
# that nests it under the subfolder
|
||||
if subfolder is not None and not os.path.isdir(
|
||||
os.path.join(name_or_path, subfolder)
|
||||
):
|
||||
subfolder = None
|
||||
model = cls.aitk_from_pretrained(
|
||||
name_or_path, subfolder=subfolder, dtype=dtype, **kwargs
|
||||
)
|
||||
|
||||
if qtype is not None:
|
||||
model.quantize_(
|
||||
qtype, device=quantize_device, exclude=exclude_quant_modules
|
||||
)
|
||||
|
||||
if device is not None:
|
||||
model.to(device)
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def _resolve_single_file(name_or_path: str) -> str:
|
||||
"""Resolve a .safetensors reference to a local file, downloading
|
||||
'org/repo/path/file.safetensors' hub references as needed."""
|
||||
if os.path.isfile(name_or_path):
|
||||
return name_or_path
|
||||
parts = name_or_path.split("/")
|
||||
if len(parts) < 3:
|
||||
raise ValueError(
|
||||
f"'{name_or_path}' is not a local file or a hub file reference "
|
||||
"('org/repo/filename.safetensors')."
|
||||
)
|
||||
return hf_hub_download(
|
||||
repo_id="/".join(parts[:2]), filename="/".join(parts[2:])
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _load_single_file_config(
|
||||
cls, config_path: Optional[str], subfolder: Optional[str]
|
||||
):
|
||||
config_source = config_path if config_path is not None else cls.aitk_config_repo
|
||||
if config_source is None:
|
||||
raise ValueError(
|
||||
f"{cls.__name__} cannot load a single-file checkpoint without a "
|
||||
"config source; pass config_path or set aitk_config_repo."
|
||||
)
|
||||
if (
|
||||
subfolder is not None
|
||||
and os.path.isdir(config_source)
|
||||
and not os.path.isdir(os.path.join(config_source, subfolder))
|
||||
):
|
||||
subfolder = None
|
||||
return cls.aitk_load_config(config_source, subfolder=subfolder)
|
||||
|
||||
@classmethod
|
||||
def _load_single_file(
|
||||
cls,
|
||||
file_path: str,
|
||||
dtype: torch.dtype,
|
||||
config_path: Optional[str] = None,
|
||||
subfolder: Optional[str] = None,
|
||||
):
|
||||
config = cls._load_single_file_config(config_path, subfolder)
|
||||
|
||||
state_dict = load_file(file_path)
|
||||
state_dict = cls.convert_state_dict_on_load(state_dict)
|
||||
for key, value in state_dict.items():
|
||||
state_dict[key] = value.to(dtype=dtype)
|
||||
|
||||
model = cls.aitk_from_config(config)
|
||||
model.load_state_dict(state_dict, assign=True)
|
||||
model.to(dtype=dtype)
|
||||
del state_dict
|
||||
flush()
|
||||
return model
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# quantization
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@torch.no_grad()
|
||||
def quantize_(
|
||||
self,
|
||||
qtype: str,
|
||||
device: Optional[torch.device] = None,
|
||||
exclude: Optional[List[str]] = None,
|
||||
):
|
||||
"""Quantize the model weights in place. When device is given, the repeated
|
||||
blocks (get_quantization_block_names) are moved there one at a time for the
|
||||
quantization math and returned to their original device, so the whole model
|
||||
never has to fit on the gpu in full precision."""
|
||||
from optimum.quanto import freeze
|
||||
from toolkit.dequantize import patch_dequantization_on_save
|
||||
from toolkit.util.quantize import get_qtype, quantize
|
||||
|
||||
# make full-model saves emit plain full precision weights
|
||||
patch_dequantization_on_save(self)
|
||||
|
||||
quantization_type = get_qtype(qtype)
|
||||
exclude = list(exclude or []) + list(
|
||||
self.get_quantization_exclude_modules() or []
|
||||
)
|
||||
|
||||
blocks: List[torch.nn.Module] = []
|
||||
for name in self.get_quantization_block_names() or []:
|
||||
# name may be a dotted path for models that nest their blocks
|
||||
block_list = self
|
||||
for part in name.split("."):
|
||||
block_list = getattr(block_list, part, None)
|
||||
if block_list is None:
|
||||
break
|
||||
if block_list is not None:
|
||||
blocks += list(block_list)
|
||||
|
||||
for block in tqdm(blocks, desc=f"Quantizing blocks ({qtype})"):
|
||||
first_param = next(block.parameters(), None)
|
||||
orig_device = first_param.device if first_param is not None else None
|
||||
if device is not None and orig_device is not None:
|
||||
block.to(device, non_blocking=True)
|
||||
quantize(block, weights=quantization_type)
|
||||
freeze(block)
|
||||
if device is not None and orig_device is not None:
|
||||
# NOT non_blocking: an async D2H allocates the cpu destination in pinned
|
||||
# memory, which the caching host allocator keeps forever (with power-of-2
|
||||
# bucket rounding on top) — that silently retained a model-sized chunk of
|
||||
# host ram after the weights moved back to the gpu for training
|
||||
block.to(orig_device)
|
||||
|
||||
# everything the block pass did not cover (embedders, norms, projections, ...)
|
||||
quantize(self, weights=quantization_type, exclude=exclude)
|
||||
freeze(self)
|
||||
|
||||
self.aitk_is_quantized = True
|
||||
self.aitk_qtype = qtype
|
||||
flush()
|
||||
return self
|
||||
|
|
@ -0,0 +1,101 @@
|
|||
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
|
||||
Loading…
Reference in New Issue