484 lines
19 KiB
Python
484 lines
19 KiB
Python
"""
|
|
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
|
|
|
|
|
|
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 (Phase 2 flips the loading default to these) ----
|
|
# hub repo the comfy-format weight files are published in
|
|
aitk_comfy_repo: Optional[str] = None
|
|
# standard name_or_path (hub repo id) -> repo-relative comfy weight file
|
|
# (ComfyUI folder layout: diffusion_models/, text_encoders/, vae/, ...)
|
|
aitk_comfy_weight_names: Dict[str, 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
|
|
|
|
# ---- 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
|
|
|
|
# ---- 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.
|
|
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 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,
|
|
):
|
|
state_dict = load_file(file_path)
|
|
return cls.load_from_state_dict(
|
|
state_dict, dtype, config_path=config_path, subfolder=subfolder
|
|
)
|
|
|
|
@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,
|
|
subfolder: Optional[str] = None,
|
|
):
|
|
"""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)."""
|
|
state_dict = cls.convert_state_dict_on_load(state_dict)
|
|
has_quant_markers = any(k.endswith(".comfy_quant") for k in state_dict)
|
|
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
|
|
)
|
|
cls._load_state_dict_with_quantized(model, state_dict)
|
|
model.aitk_is_quantized = True
|
|
else:
|
|
for key, value in state_dict.items():
|
|
state_dict[key] = value.to(dtype=dtype)
|
|
model.load_state_dict(state_dict, assign=True)
|
|
model.to(dtype=dtype)
|
|
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 find_comfy_weights(cls, name_or_path: str) -> Optional[str]:
|
|
"""Local comfy-format weight file registered for a standard
|
|
``name_or_path``, or None. Never downloads — Phase 2 flips the loading
|
|
default to comfy sources; until then callers opt in explicitly."""
|
|
from toolkit.models.v2.resolver import resolve_comfy_file
|
|
|
|
rel_path = cls.aitk_comfy_weight_names.get(name_or_path)
|
|
if rel_path is None:
|
|
return None
|
|
return resolve_comfy_file(
|
|
rel_path, repo_id=cls.aitk_comfy_repo, local_only=True
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# 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 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 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,
|
|
):
|
|
"""Save as a single-file .safetensors in the model's original (comfy)
|
|
key layout, via convert_state_dict_on_save. Quantized weights are
|
|
dequantized to full precision; dtype, when given, casts the floating
|
|
point tensors. (Quantized-storage saves — comfy_quant markers — come
|
|
with Phase 2.)"""
|
|
from safetensors.torch import save_file
|
|
from toolkit.util.quantize import dequantize_if_quantized
|
|
|
|
state_dict = {}
|
|
for key, value in self.state_dict().items():
|
|
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 = 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.from_pretrained(
|
|
path, subfolder=subfolder or "", torch_dtype=dtype, **kwargs
|
|
)
|
|
|
|
@classmethod
|
|
def aitk_load_config(cls, path, subfolder=None):
|
|
from transformers import AutoConfig
|
|
|
|
return AutoConfig.from_pretrained(path, subfolder=subfolder or "")
|
|
|
|
@classmethod
|
|
def aitk_from_config(cls, config):
|
|
with torch.device("meta"):
|
|
return cls(config)
|