""" 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=...`). A single-file checkpoint carrying ComfyUI `comfy_quant` markers loads with its pre-quantized layers attached to the toolkit's quantization backends (convrot8 / nvfp4) automatically. `save_model` writes a single-file .safetensors back in the original (comfy) key layout. 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 def adopt_component(module: torch.nn.Module, wrapper_cls: type) -> torch.nn.Module: """Rebind an already-loaded component instance onto its v2 wrapper class (in-place class swap, like the OstrisLinear conversion). Used when a component arrives from machinery that builds the base class directly (e.g. a diffusers pipeline load), so every resident component is a mixin instance the inference engine can pool and manage.""" if isinstance(module, wrapper_cls): return module base = wrapper_cls.__mro__[1] if not isinstance(module, base): raise TypeError( f"cannot adopt {type(module).__name__} into {wrapper_cls.__name__} " f"(expected a {base.__name__})" ) module.__class__ = wrapper_cls return module 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 # ---- comfy weight sources ---- # hub repo the comfy-format weight files are published in (Comfy-Org/...) aitk_comfy_repo: Optional[str] = None # standard name_or_path (hub repo id) -> repo-relative comfy weight file # candidates for this module: every precision variant the class can digest. # Selection ranks them convrot8 > float8 mixed > float8 > bf16 > fp16 # (resolver.comfy_precision_rank); a locally-present candidate always wins # over a download. aitk_comfy_weight_names: Dict[str, List[str]] = {} # ---- tokenizer/processor source, for text-encoder modules ---- aitk_tokenizer_repo: Optional[str] = None aitk_tokenizer_subfolder: Optional[str] = None aitk_processor_repo: Optional[str] = None aitk_processor_subfolder: Optional[str] = None # single-file loads without quant markers cast tensors to the requested # dtype; classes whose checkpoints carry a deliberate precision mix (e.g. # fp32 norms next to bf16 weights) set this False to load at stored # precision instead aitk_cast_on_load: bool = True # pre-quantized (comfy marker) checkpoints normally load their # NON-quantized tensors at stored precision (the mix can be deliberate, # e.g. ltx2.5's fp32 tables). Classes whose forward assumes one uniform # dtype (wan: fp16 conv biases next to fp32 tables in the fp8 files) set # this True to cast those float tensors to the load dtype instead aitk_cast_quantized_load: bool = False # ---- 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_transformer_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 def get_offload_ignore_modules(self) -> Optional[List]: """Module/parameter objects layer offloading must keep resident (pad tokens, modulation tables, ...). Default: none.""" 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._local_first( cls.from_pretrained, path, subfolder=subfolder, torch_dtype=dtype, **kwargs ) @classmethod def aitk_load_config(cls, path, subfolder=None): return cls._local_first(cls.load_config, path, subfolder=subfolder) @staticmethod def _local_first(loader, *args, **kwargs): """Serve fully-cached repos without touching the hub. Online from_pretrained does an etag HEAD round-trip per file (seconds of pure network latency on every warm load); a local_files_only attempt is instant when the cache can't satisfy it, so failing over to the online path costs nothing.""" import os as _os if _os.environ.get("HF_HUB_OFFLINE") == "1": return loader(*args, **kwargs) try: return loader(*args, local_files_only=True, **kwargs) except Exception: return loader(*args, **kwargs) @classmethod def aitk_from_config(cls, config): # params on meta (materialized by the state-dict assign), buffers real: # non-persistent buffers (e.g. rope tables) are computed at init and # never appear in checkpoints from accelerate import init_empty_weights with init_empty_weights(include_buffers=False): if hasattr(cls, "from_config"): return cls.from_config(config) # plain nn.Module classes take their params/config object directly return cls(config) # ------------------------------------------------------------------ # loading # ------------------------------------------------------------------ @classmethod def load( cls, name_or_path: str, qtype: Optional[str] = None, offload: float = 0.0, dtype: torch.dtype = torch.bfloat16, device: Optional[torch.device] = None, base_model=None, quantize_device: Optional[torch.device] = None, exclude_quant_modules: Optional[List[str]] = None, config_path: Optional[str] = None, config=None, subfolder: Optional[str] = None, use_comfy_weights: bool = True, **kwargs, ): """One-call component load. Resolves/downloads the source (comfy candidates included), builds the module, quantizes it, attaches layer offloading, and places it — everything a holder used to hand-roll. qtype: quantization type; "type|path.safetensors" also attaches an accuracy recovery adapter (requires base_model for the network wiring). Pre-quantized checkpoints skip quantization automatically. offload: layer-offload fraction (0 disables); the class's get_offload_ignore_modules hook pins sensitive pieces. device: final placement. With offload active the memory manager owns staging, so no blanket device move happens. base_model: the holder — supplies status prints, quantize_kwargs, and ARA wiring when given. """ model = cls.load_model( name_or_path, dtype=dtype, config_path=config_path, config=config, subfolder=subfolder, use_comfy_weights=use_comfy_weights, # comfy candidate ranking prefers the file whose shipped # quantization matches the request (strip any ARA suffix); # quantization itself happens in aitk_post_load below qtype=(qtype or "").split("|", 1)[0] or None, quantize_on_load=False, **kwargs, ) return model.aitk_post_load( qtype=qtype, offload=offload, dtype=dtype, device=device, base_model=base_model, quantize_device=quantize_device, exclude_quant_modules=exclude_quant_modules, ) def aitk_post_load( self, qtype: Optional[str] = None, offload: float = 0.0, dtype: torch.dtype = torch.bfloat16, device: Optional[torch.device] = None, base_model=None, quantize_device: Optional[torch.device] = None, exclude_quant_modules: Optional[List[str]] = None, **_load_only_kwargs, # tolerate component_load_kwargs() extras ): """The post-load half of `load` — quantize (incl. "qtype|ara"), attach layer offloading, place on device. Callable directly by holders whose checkpoint sourcing is custom but want the standard policy.""" from toolkit.print import print_acc status_fn = ( base_model.print_and_status_update if base_model is not None else print_acc ) ara_path = None if qtype is not None and "|" in qtype: qtype, ara_path = qtype.split("|", 1) # pre-quantized checkpoint: keep the shipped quantization IFF it # exactly matches the request. Anything else — no quantization # requested (full finetuning), a different backend, or an ARA — # restores/requantizes to what was asked for. if getattr(self, "aitk_is_quantized", False): from toolkit.util.ostris_quant import OstrisLinear from toolkit.util.quantize import ( dequantize_ostris_to_linear, get_qtype, ostristype, ) shipped = sorted( { getattr(m.ostris_quantizer, "qtype", None) for m in self.modules() if isinstance(m, OstrisLinear) } - {None} ) # a checkpoint may deliberately mix backends (minimax's TE: nvfp4 # LM linears + int8 embeddings). The request matches when the # requested backend is among the shipped ones — the whole shipped # quantization is then kept exactly as-is. matches = ( qtype is not None and ara_path is None and qtype in shipped ) if matches: status_fn( f"Checkpoint is pre-quantized ({'/'.join(shipped)}); keeping it" ) else: target_is_ostris = qtype is not None and isinstance( get_qtype(qtype), ostristype ) if qtype is None or ara_path is not None or not target_is_ostris: # full precision requested, an ARA (quantizes fresh from # full weights), or a quanto/torchao backend that cannot # re-quantize an OstrisLinear: dequantize first status_fn( f"Dequantizing shipped {'/'.join(shipped)} weights " f"-> {qtype or 'full precision'}" ) dequantize_ostris_to_linear(self) else: # ostris -> ostris: quantize_module re-quantizes each # layer in place (dequant -> requant, one layer at a time) status_fn(f"Requantizing {'/'.join(shipped)} -> {qtype}") self.aitk_is_quantized = False self.aitk_qtype = None if qtype and not getattr(self, "aitk_is_quantized", False): from toolkit.util.quantize import ( attach_ara_and_quantize, quantize_module, ) exclude = list(exclude_quant_modules or []) + list( self.get_quantization_exclude_modules() or [] ) q_device = quantize_device if quantize_device is not None else device # final home is the quantize gpu (no offloading, no cpu # parking): quantized weights stay put — one bus crossing # instead of three, and no model-sized host-ram copy keep_on_device = ( not offload and device is not None and q_device is not None and torch.device(device) == torch.device(q_device) and torch.device(device).type != "cpu" ) if ara_path is not None: if base_model is None: raise ValueError( "an accuracy recovery adapter needs base_model= for the " "network wiring" ) status_fn("Quantizing with accuracy recovery adapter") attach_ara_and_quantize( base_model, self, ara_path, exclude=exclude, device=q_device, keep_on_device=keep_on_device, ) else: status_fn(f"Quantizing ({qtype})") quantize_module( self, qtype, device=q_device, dtype=dtype, block_names=self.get_transformer_block_names(), exclude=exclude, quantize_kwargs=( base_model.model_config.quantize_kwargs if base_model is not None else None ), status_fn=status_fn, keep_on_device=keep_on_device, ) self.aitk_is_quantized = True self.aitk_qtype = qtype if offload and offload > 0: from toolkit.memory_management import MemoryManager # the manager stages layers to the COMPUTE device. `device` is the # component's parking device — "cpu" under low_vram — which must # never become the manager's process device (bounces would stage # to cpu against cuda activations). quantize_device carries the # actual gpu in the holder flow. compute_device = quantize_device if quantize_device is not None else device MemoryManager.attach( self, torch.device(compute_device), offload_percent=offload, ignore_modules=list(self.get_offload_ignore_modules() or []), ) elif device is not None: self.to(device) return self @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, config=None, subfolder: Optional[str] = None, use_comfy_weights: bool = True, quantize_on_load: bool = True, **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"). When name_or_path is a standard hub repo this class has comfy weight candidates registered for (aitk_comfy_weight_names), the comfy file is the preferred source: the best-ranked local candidate, or the top-preference candidate downloaded into the comfy layout under MODELS_PATH. The standard repo still supplies the config. Local dirs and unregistered repos load as-is; use_comfy_weights=False opts out. 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. subfolder: overrides the class's aitk_subfolder; pass "" to explicitly load from the checkpoint root (e.g. a raw hub repo). """ if subfolder is None: subfolder = cls.aitk_subfolder elif subfolder == "": subfolder = None if ( use_comfy_weights and not name_or_path.endswith(".safetensors") and not os.path.exists(name_or_path) ): comfy_path = cls.resolve_comfy_weights( name_or_path, subfolder=subfolder, qtype=qtype ) if comfy_path is not None: if config_path is None: # the standard repo supplies the config for the comfy file config_path = name_or_path name_or_path = comfy_path 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, config=config, subfolder=subfolder, **kwargs, ) 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 ) # quantize_on_load=False: qtype was only a comfy-candidate ranking # hint (the .load() path quantizes in aitk_post_load instead) if qtype is not None and quantize_on_load: 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) @staticmethod def _readahead(file_path: str): """Queue kernel readahead for the whole file. safetensors tensors are mmap-backed; without this a cold cache faults pages in per-tensor cudaMemcpy order at a fraction of the drive's sequential bandwidth.""" if not hasattr(os, "posix_fadvise"): return # non-POSIX: mmap readahead heuristics only try: fd = os.open(file_path, os.O_RDONLY) try: os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_WILLNEED) finally: os.close(fd) except OSError: pass @classmethod def _load_single_file( cls, file_path: str, dtype: torch.dtype, config_path: Optional[str] = None, config=None, subfolder: Optional[str] = None, **kwargs, ): cls._readahead(file_path) state_dict = load_file(file_path) return cls.load_from_state_dict( state_dict, dtype, config_path=config_path, config=config, subfolder=subfolder, **kwargs, ) @classmethod def aitk_config_from_state_dict(cls, state_dict: Dict[str, torch.Tensor]): """Derive the model config from the checkpoint itself (e.g. sniffing block counts from key indices). Return None (the default) to load the config from config_path / aitk_config_repo instead.""" return None @classmethod def load_from_state_dict( cls, state_dict: Dict[str, torch.Tensor], dtype: torch.dtype, config_path: Optional[str] = None, config=None, subfolder: Optional[str] = None, **kwargs, ): """Build the model and load an already-read single-file state dict (the tail of the single-file path; also callable directly for checkpoints read from non-safetensors sources). ``config``, when given, is used directly (for models whose config comes from the holder, e.g. model_kwargs-driven archs).""" state_dict = cls.convert_state_dict_on_load(state_dict) has_quant_markers = "scaled_fp8" in state_dict or any( k.endswith(".comfy_quant") for k in state_dict ) if config is None: config = cls.aitk_config_from_state_dict(state_dict) if config is None: config = cls._load_single_file_config(config_path, subfolder) model = cls.aitk_from_config(config) if has_quant_markers: # pre-quantized comfy checkpoint: attach the quantized linears onto # the toolkit's backends, load the rest at its stored precision # (the checkpoint's bf16/fp16/fp32 mix is deliberate) from toolkit.util.comfy_quant_import import import_comfy_quantized_layers state_dict, num_quantized = import_comfy_quantized_layers( model, state_dict, orig_dtype=dtype ) if cls.aitk_cast_quantized_load: # uniform-dtype models: the leftover (non-quantized) float # tensors follow the load dtype, like the non-marker path for key, value in state_dict.items(): if value.is_floating_point(): state_dict[key] = value.to(dtype=dtype) cls._load_state_dict_with_quantized(model, state_dict) if cls.aitk_cast_quantized_load: # the importer assigns quantized-layer biases directly at # stored dtype; they must follow too (a stray fp16 bias also # makes ModelMixin.dtype — the pipelines' cast target — lie) from toolkit.util.ostris_quant import OstrisLinear for m in model.modules(): if isinstance(m, OstrisLinear) and m.bias is not None: m.bias.data = m.bias.data.to(dtype=dtype) model.aitk_is_quantized = True elif cls.aitk_cast_on_load: for key, value in state_dict.items(): if value.is_floating_point(): state_dict[key] = value.to(dtype=dtype) model.load_state_dict(state_dict, assign=True) model.to(dtype=dtype) else: # stored-precision load (the checkpoint's dtype mix is deliberate) model.load_state_dict(state_dict, assign=True) del state_dict flush() return model @staticmethod def _load_state_dict_with_quantized(model, state_dict): """Load a state dict onto a (meta-built) model whose quantized linears were already attached by import_comfy_quantized_layers: their weight (and importer-assigned bias) legitimately report as missing keys.""" from toolkit.util.ostris_quant import OstrisLinear result = model.load_state_dict(state_dict, assign=True, strict=False) quantized_param_keys = set() for name, m in model.named_modules(): if isinstance(m, OstrisLinear): quantized_param_keys.add(f"{name}.weight") if m.bias is not None: quantized_param_keys.add(f"{name}.bias") bad_missing = [k for k in result.missing_keys if k not in quantized_param_keys] if bad_missing or result.unexpected_keys: raise ValueError( f"{type(model).__name__} load mismatch: missing {bad_missing[:8]}, " f"unexpected {result.unexpected_keys[:8]}" ) leftover_meta = [n for n, p in model.named_parameters() if p.is_meta] if leftover_meta: raise ValueError( f"{type(model).__name__} load left meta parameters: " f"{leftover_meta[:8]}" ) # ------------------------------------------------------------------ # comfy weight sources # ------------------------------------------------------------------ @classmethod def resolve_comfy_weights( cls, name_or_path: str, subfolder: Optional[str] = None, local_only: bool = False, hf_token: Optional[str] = None, status_fn: Optional[callable] = None, qtype: Optional[str] = None, ) -> Optional[str]: """The comfy-format weight file replacing a standard ``name_or_path``, or None when this class has none registered for it. Best-ranked local candidate wins; otherwise the top-preference candidate is downloaded to the comfy layout under MODELS_PATH (unless local_only). Candidate keys may be plain repo ids or ``(repo_id, subfolder)`` tuples for checkpoints holding several of this component (e.g. wan2.2's transformer / transformer_2).""" from toolkit.models.v2.resolver import resolve_comfy_candidates candidates = None if subfolder is not None: candidates = cls.aitk_comfy_weight_names.get((name_or_path, subfolder)) if candidates is None: candidates = cls.aitk_comfy_weight_names.get(name_or_path) repo_id = cls.aitk_comfy_repo if isinstance(candidates, dict): # entries may carry their own comfy repo ({"repo": ..., "files": [...]}) repo_id = candidates.get("repo", repo_id) candidates = candidates.get("files") if not candidates or repo_id is None: return None return resolve_comfy_candidates( candidates, repo_id=repo_id, hf_token=hf_token, status_fn=status_fn, local_only=local_only, qtype=qtype, ) # ------------------------------------------------------------------ # tokenizer / processor (text-encoder modules) # ------------------------------------------------------------------ @classmethod def load_tokenizer( cls, name_or_path: Optional[str] = None, subfolder: Optional[str] = None, **kwargs, ): """Load the tokenizer from ``name_or_path`` (a checkpoint dir or repo holding it at aitk_tokenizer_subfolder), falling back to the class's aitk_tokenizer_repo. subfolder overrides the class default; "" loads from the checkpoint root.""" from transformers import AutoTokenizer source = name_or_path if name_or_path is not None else cls.aitk_tokenizer_repo if source is None: raise ValueError(f"{cls.__name__} does not declare aitk_tokenizer_repo") if subfolder is None: subfolder = cls.aitk_tokenizer_subfolder if subfolder and os.path.isdir(source) and not os.path.isdir( os.path.join(source, subfolder) ): subfolder = None return cls._local_first( AutoTokenizer.from_pretrained, source, subfolder=subfolder or "", **kwargs ) @classmethod def load_processor( cls, name_or_path: Optional[str] = None, subfolder: Optional[str] = None, **kwargs, ): from transformers import AutoProcessor source = name_or_path if name_or_path is not None else cls.aitk_processor_repo if source is None: raise ValueError(f"{cls.__name__} does not declare aitk_processor_repo") if subfolder is None: subfolder = cls.aitk_processor_subfolder if subfolder and os.path.isdir(source) and not os.path.isdir( os.path.join(source, subfolder) ): subfolder = None return cls._local_first( AutoProcessor.from_pretrained, source, subfolder=subfolder or "", **kwargs ) # ------------------------------------------------------------------ # saving # ------------------------------------------------------------------ @torch.no_grad() def save_model( self, output_path: str, dtype: Optional[torch.dtype] = None, metadata: Optional[Dict[str, str]] = None, quantized: Optional[bool] = None, ): """Save as a single-file .safetensors in the model's original (comfy) key layout, via convert_state_dict_on_save. quantized: None (auto) keeps quantized storage — comfy_quant markers, loadable by ComfyUI and by this class — when the model holds quantized layers whose backend has a comfy format (convrot8 / nvfp4 / convrotcomfyw4a4), and saves dequantized full precision otherwise. True forces quantized storage (raises if any quantized layer's backend has no comfy format); False forces a dequantized save. dtype, when given, casts the floating point (non-quantized) tensors.""" from safetensors.torch import save_file from toolkit.util.quantize import dequantize_if_quantized q_entries: Dict[str, torch.Tensor] = {} exported: List[str] = [] if quantized is None or quantized: from toolkit.util.comfy_quant_export import export_comfy_quantized_layers q_entries, exported, unexportable = export_comfy_quantized_layers(self) if unexportable: if quantized: raise ValueError( f"{type(self).__name__}: quantized layers without a comfy " f"storage format: {unexportable[:8]}" ) # auto mode: a partially-quantized file would be inconsistent — # fall back to a fully dequantized save q_entries, exported = {}, [] skip_keys = {f"{name}.weight" for name in exported} state_dict = {} for key, value in self.state_dict().items(): if key in skip_keys: continue value = dequantize_if_quantized(value) if dtype is not None and value.is_floating_point(): value = value.to(dtype=dtype) state_dict[key] = value.detach().to("cpu").contiguous() state_dict.update(q_entries) state_dict = self.convert_state_dict_on_save(state_dict) parent = os.path.dirname(output_path) if parent: os.makedirs(parent, exist_ok=True) save_file(state_dict, output_path, metadata=metadata) del state_dict flush() # ------------------------------------------------------------------ # 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_transformer_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_transformer_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 class OstrisTransformersMixin(OstrisModelMixin): """OstrisModelMixin with the backend hooks speaking the transformers-lib API (PreTrainedModel / AutoConfig) instead of diffusers ModelMixin. Base for text-encoder and vision-encoder modules.""" @classmethod def aitk_from_pretrained(cls, path, subfolder=None, dtype=None, **kwargs): return cls._local_first( cls.from_pretrained, path, subfolder=subfolder or "", torch_dtype=dtype, **kwargs ) @classmethod def aitk_load_config(cls, path, subfolder=None): from transformers import AutoConfig return cls._local_first( AutoConfig.from_pretrained, path, subfolder=subfolder or "" ) @classmethod def aitk_from_config(cls, config): from accelerate import init_empty_weights with init_empty_weights(include_buffers=False): return cls(config)