diff --git a/extensions_built_in/diffusion_models/z_image/z_image.py b/extensions_built_in/diffusion_models/z_image/z_image.py index 6f7ba2c1..bf13983e 100644 --- a/extensions_built_in/diffusion_models/z_image/z_image.py +++ b/extensions_built_in/diffusion_models/z_image/z_image.py @@ -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" diff --git a/toolkit/models/v2/__init__.py b/toolkit/models/v2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py new file mode 100644 index 00000000..6fb71ecf --- /dev/null +++ b/toolkit/models/v2/_mixin.py @@ -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 diff --git a/toolkit/models/v2/z_image.py b/toolkit/models/v2/z_image.py new file mode 100644 index 00000000..90df66d7 --- /dev/null +++ b/toolkit/models/v2/z_image.py @@ -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