From e8d9cf6d35976d5905c0578c1b0825a6eaa32305 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 27 Aug 2026 10:51:43 -0600 Subject: [PATCH 01/14] Models v2 - phase 0 --- .../audio_models/base_audio_model.py | 16 +- .../boogu_image/boogu_image.py | 11 +- .../diffusion_models/chroma/chroma_model.py | 15 +- .../chroma/chroma_radiance_model.py | 15 +- .../ernie_image/ernie_image.py | 13 +- .../example_model/example_model.py | 19 +- .../diffusion_models/f_light/f_light.py | 15 +- .../diffusion_models/flux2/flux2_model.py | 14 +- .../diffusion_models/hidream/hidream_model.py | 15 +- .../diffusion_models/ideogram4/ideogram4.py | 13 +- .../diffusion_models/krea2/krea2.py | 11 +- .../diffusion_models/ltx2/ltx2.py | 110 ++------- .../diffusion_models/mageflow/mageflow.py | 13 +- .../diffusion_models/minimax_h3/minimax_h3.py | 93 ++----- .../nucleus_image/nucleus_image_model.py | 13 +- .../diffusion_models/omnigen2/__init__.py | 16 +- .../prx_pixel_t2i/prx_pixel_t2i.py | 11 +- .../diffusion_models/qwen_image/qwen_image.py | 14 +- .../diffusion_models/z_image/z_image.py | 17 +- .../zeta_chroma/zeta_chroma_model.py | 14 +- toolkit/models/base_model.py | 15 +- toolkit/models/v2/PLANNING.md | 227 ++++++++++++++++++ toolkit/models/v2/_mixin.py | 151 +++++++++++- .../models/v2/diffusion_models/__init__.py | 0 .../v2/{ => diffusion_models}/z_image.py | 4 +- toolkit/models/v2/resolver.py | 129 ++++++++++ toolkit/models/v2/text_encoders/__init__.py | 0 toolkit/models/v2/vae/__init__.py | 0 toolkit/models/v2/vision_encoders/__init__.py | 0 29 files changed, 583 insertions(+), 401 deletions(-) create mode 100644 toolkit/models/v2/PLANNING.md create mode 100644 toolkit/models/v2/diffusion_models/__init__.py rename toolkit/models/v2/{ => diffusion_models}/z_image.py (98%) create mode 100644 toolkit/models/v2/resolver.py create mode 100644 toolkit/models/v2/text_encoders/__init__.py create mode 100644 toolkit/models/v2/vae/__init__.py create mode 100644 toolkit/models/v2/vision_encoders/__init__.py diff --git a/extensions_built_in/audio_models/base_audio_model.py b/extensions_built_in/audio_models/base_audio_model.py index 6860a1a..f47b8f6 100644 --- a/extensions_built_in/audio_models/base_audio_model.py +++ b/extensions_built_in/audio_models/base_audio_model.py @@ -69,21 +69,7 @@ class BaseAudioModel(BaseModel): "save_model is not implemented for this model. Use the pipeline directly instead." ) - def convert_lora_weights_before_save(self, state_dict): - # currently starte with transformer. but needs to start with diffusion_model. for comfyui - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd - - def convert_lora_weights_before_load(self, state_dict): - # saved as diffusion_model. but needs to be transformer. for ai-toolkit - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True def encode_images(self, image_list: torch.Tensor, device=None, dtype=None): # make it more obvious for audio models diff --git a/extensions_built_in/diffusion_models/boogu_image/boogu_image.py b/extensions_built_in/diffusion_models/boogu_image/boogu_image.py index fe5b75b..8241561 100644 --- a/extensions_built_in/diffusion_models/boogu_image/boogu_image.py +++ b/extensions_built_in/diffusion_models/boogu_image/boogu_image.py @@ -446,14 +446,5 @@ class BooguImageModel(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ["double_stream_layers", "single_stream_layers"] - def convert_lora_weights_before_save(self, state_dict): - return { - k.replace("transformer.", "diffusion_model."): v - for k, v in state_dict.items() - } + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - return { - k.replace("diffusion_model.", "transformer."): v - for k, v in state_dict.items() - } diff --git a/extensions_built_in/diffusion_models/chroma/chroma_model.py b/extensions_built_in/diffusion_models/chroma/chroma_model.py index 236d950..15da572 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_model.py @@ -443,21 +443,8 @@ class ChromaModel(BaseModel): batch = kwargs.get('batch') return (noise - batch.latents).detach() - def convert_lora_weights_before_save(self, state_dict): - # currently starte with transformer. but needs to start with diffusion_model. for comfyui - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - # saved as diffusion_model. but needs to be transformer. for ai-toolkit - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd def get_base_model_version(self): return "chroma" diff --git a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py index 333600e..e5a79fe 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py @@ -425,21 +425,8 @@ class ChromaRadianceModel(BaseModel): batch = kwargs.get('batch') return (noise - batch.latents).detach() - def convert_lora_weights_before_save(self, state_dict): - # currently starte with transformer. but needs to start with diffusion_model. for comfyui - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - # saved as diffusion_model. but needs to be transformer. for ai-toolkit - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd def get_base_model_version(self): return "chroma_radiance" diff --git a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py index dba1204..bb296f6 100644 --- a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py +++ b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py @@ -377,16 +377,5 @@ class ErnieImageModel(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ["layers"] - def convert_lora_weights_before_save(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd diff --git a/extensions_built_in/diffusion_models/example_model/example_model.py b/extensions_built_in/diffusion_models/example_model/example_model.py index a813f39..51b973e 100644 --- a/extensions_built_in/diffusion_models/example_model/example_model.py +++ b/extensions_built_in/diffusion_models/example_model/example_model.py @@ -504,19 +504,8 @@ class ExampleModel(BaseModel): attribute in src/model.py.""" return ["blocks"] - def convert_lora_weights_before_save(self, state_dict): - """Map internal LoRA keys to the ecosystem-standard naming right before - the .safetensors is written. Most modern models ship LoRAs with a - ``diffusion_model.`` prefix (ComfyUI convention); internally ai-toolkit - uses ``transformer.``.""" - return { - k.replace("transformer.", "diffusion_model."): v - for k, v in state_dict.items() - } + # LoRA keys save with the ecosystem-standard ``diffusion_model.`` prefix + # (ComfyUI convention) and load back to the internal ``transformer.`` + # prefix; see BaseModel.convert_lora_weights_before_save/load + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - """Inverse of the above, applied when resuming from a saved LoRA.""" - return { - k.replace("diffusion_model.", "transformer."): v - for k, v in state_dict.items() - } diff --git a/extensions_built_in/diffusion_models/f_light/f_light.py b/extensions_built_in/diffusion_models/f_light/f_light.py index 2fc4f5a..813e6c3 100644 --- a/extensions_built_in/diffusion_models/f_light/f_light.py +++ b/extensions_built_in/diffusion_models/f_light/f_light.py @@ -270,21 +270,8 @@ class FLiteModel(BaseModel): # return (noise - batch.latents).detach() return (batch.latents - noise).detach() - def convert_lora_weights_before_save(self, state_dict): - # currently starte with transformer. but needs to start with diffusion_model. for comfyui - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - # saved as diffusion_model. but needs to be transformer. for ai-toolkit - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd def get_base_model_version(self): return "f-lite" diff --git a/extensions_built_in/diffusion_models/flux2/flux2_model.py b/extensions_built_in/diffusion_models/flux2/flux2_model.py index e22a1cd..e39e468 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_model.py @@ -505,19 +505,7 @@ class Flux2Model(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ["double_blocks", "single_blocks"] - def convert_lora_weights_before_save(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd - - def convert_lora_weights_before_load(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None): if device is None: diff --git a/extensions_built_in/diffusion_models/hidream/hidream_model.py b/extensions_built_in/diffusion_models/hidream/hidream_model.py index 7bba831..71921aa 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_model.py @@ -432,21 +432,8 @@ class HidreamModel(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ['double_stream_blocks', 'single_stream_blocks'] - def convert_lora_weights_before_save(self, state_dict): - # currently starte with transformer. but needs to start with diffusion_model. for comfyui - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - # saved as diffusion_model. but needs to be transformer. for ai-toolkit - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd def get_base_model_version(self): return "hidream_i1" diff --git a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py index e42a4a2..390f7a4 100644 --- a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py +++ b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py @@ -620,16 +620,5 @@ class Ideogram4Model(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ["layers"] - def convert_lora_weights_before_save(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd diff --git a/extensions_built_in/diffusion_models/krea2/krea2.py b/extensions_built_in/diffusion_models/krea2/krea2.py index 128fdb0..22d75a8 100644 --- a/extensions_built_in/diffusion_models/krea2/krea2.py +++ b/extensions_built_in/diffusion_models/krea2/krea2.py @@ -867,14 +867,5 @@ class Krea2Model(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ["blocks"] - def convert_lora_weights_before_save(self, state_dict): - return { - k.replace("transformer.", "diffusion_model."): v - for k, v in state_dict.items() - } + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - return { - k.replace("diffusion_model.", "transformer."): v - for k, v in state_dict.items() - } diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index bbf5e57..2821e23 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -1191,23 +1191,17 @@ class LTX2Model(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ["transformer_blocks"] + lora_keys_use_comfy_prefix = True + def convert_lora_weights_before_save(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - new_sd = convert_lora_diffusers_to_original(new_sd, version=self.ltx_version) - return new_sd + state_dict = super().convert_lora_weights_before_save(state_dict) + return convert_lora_diffusers_to_original(state_dict, version=self.ltx_version) def convert_lora_weights_before_load(self, state_dict): state_dict = convert_lora_original_to_diffusers( state_dict, version=self.ltx_version ) - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd + return super().convert_lora_weights_before_load(state_dict) class LTX23Model(LTX2Model): @@ -1239,88 +1233,30 @@ class LTX25Model(LTX2Model): ltx_te_path = None # ------------------------------------------------------------------ - # ComfyUI-style file resolution (mirrors MinimaxH3Model) + # ComfyUI-style file resolution (toolkit/models/v2/resolver.py) # ------------------------------------------------------------------ - @staticmethod - def _find_file_recursive(root_dir: str, filename: str) -> Optional[str]: - if not os.path.isdir(root_dir): - return None - for dirpath, dirnames, filenames in os.walk(root_dir): - dirnames.sort() - if filename in filenames: - return os.path.join(dirpath, filename) - return None - def _resolve_comfy_file(self, component: str) -> str: - """Find a weight file at its local location, or download it there - when (and only when) it is missing. - - Search order: model_kwargs override, the repo-relative path under - MODELS_PATH (diffusion_models/, text_encoders/, vae/), the bare - filename at the root, any subfolder of the component's category - folder (recursive), then the hub — downloaded to the repo-relative - path under MODELS_PATH. - """ - override = self.model_config.model_kwargs.get(f"{component}_path", None) - if override is not None: - if not os.path.exists(override): - raise FileNotFoundError( - f"model_kwargs.{component}_path does not exist: {override}" - ) - return override - - rel_path = COMFY_LTX25_FILES[component] - filename = os.path.basename(rel_path) - category = os.path.dirname(rel_path) - for rel in (rel_path, filename): - candidate = os.path.join(MODELS_PATH, rel) - if os.path.exists(candidate): - return candidate - found = self._find_file_recursive(os.path.join(MODELS_PATH, category), filename) - if found is not None: - return found - - repo_id = COMFY_LTX25_REPO - name_or_path = self.model_config.name_or_path - if ( - name_or_path - and not os.path.exists(name_or_path) - and not name_or_path.endswith(".safetensors") - and "/" in name_or_path - ): - repo_id = name_or_path - self.print_and_status_update( - f"Downloading {rel_path} from {repo_id} into {MODELS_PATH}" + from toolkit.models.v2.resolver import ( + repo_id_from_name_or_path, + resolve_comfy_file, ) - return huggingface_hub.hf_hub_download( - repo_id=repo_id, filename=rel_path, token=HF_TOKEN, local_dir=MODELS_PATH + + return resolve_comfy_file( + COMFY_LTX25_FILES[component], + repo_id=repo_id_from_name_or_path( + self.model_config.name_or_path, COMFY_LTX25_REPO + ), + override_path=self.model_config.model_kwargs.get( + f"{component}_path", None + ), + hf_token=HF_TOKEN, + status_fn=self.print_and_status_update, ) def _resolve_named_file(self, path: str, component: str) -> str: - """Resolve an explicit .safetensors path: local file, models-folder - file, or an 'org/repo/path/file.safetensors' hub path (downloaded - into the models folder).""" - if os.path.exists(path): - return path - splits = path.split("/") - if len(splits) < 3: - raise ValueError( - f"Invalid {component} path: {path}. Must be a local file or " - "'repo_id/repo/filename.safetensors' to download from the Hugging Face Hub." - ) - rel_path = "/".join(splits[2:]) - for candidate in ( - os.path.join(MODELS_PATH, rel_path), - os.path.join(MODELS_PATH, splits[-1]), - ): - if os.path.exists(candidate): - return candidate - return huggingface_hub.hf_hub_download( - repo_id="/".join(splits[:2]), - filename=rel_path, - token=HF_TOKEN, - local_dir=MODELS_PATH, - ) + from toolkit.models.v2.resolver import resolve_named_file + + return resolve_named_file(path, component=component, hf_token=HF_TOKEN) def _resolve_dit_path(self) -> str: name_or_path = self.model_config.name_or_path diff --git a/extensions_built_in/diffusion_models/mageflow/mageflow.py b/extensions_built_in/diffusion_models/mageflow/mageflow.py index 0d15ffd..589f9f2 100644 --- a/extensions_built_in/diffusion_models/mageflow/mageflow.py +++ b/extensions_built_in/diffusion_models/mageflow/mageflow.py @@ -629,18 +629,7 @@ class MageFlowModel(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ["transformer_blocks"] - def convert_lora_weights_before_save(self, state_dict): - return { - k.replace("transformer.", "diffusion_model."): v - for k, v in state_dict.items() - } - - def convert_lora_weights_before_load(self, state_dict): - return { - k.replace("diffusion_model.", "transformer."): v - for k, v in state_dict.items() - } - + lora_keys_use_comfy_prefix = True class MageFlowEditModel(MageFlowModel): arch = "mageflow_edit" diff --git a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py index 1ec8619..ab0d7fa 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py @@ -52,6 +52,11 @@ from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.memory_management import MemoryManager from toolkit.metadata import get_meta_for_safetensors from toolkit.models.base_model import BaseModel +from toolkit.models.v2.resolver import ( + find_file_recursive, + repo_id_from_name_or_path, + resolve_comfy_file, +) from toolkit.paths import MODELS_PATH from toolkit.util.comfy_quant_import import import_comfy_quantized_layers from toolkit.util.ostris_quant import OstrisLinear @@ -233,64 +238,23 @@ class MinimaxH3Model(BaseModel): # ------------------------------------------------------------------ # Loading # ------------------------------------------------------------------ - @staticmethod - def _find_file_recursive(root_dir: str, filename: str) -> Optional[str]: - """First (breadth-stable, sorted) match of ``filename`` anywhere under - ``root_dir``.""" - if not os.path.isdir(root_dir): - return None - for dirpath, dirnames, filenames in os.walk(root_dir): - dirnames.sort() - if filename in filenames: - return os.path.join(dirpath, filename) - return None - def _resolve_comfy_file(self, component: str) -> str: - """Find a weight file at its local location, or download it there - when (and only when) it is missing. - - Search order: model_kwargs override, the repo-relative path under - MODELS_PATH (diffusion_models/, text_encoders/, vae/), the bare - filename at the root, any subfolder of the component's category - folder (recursive — e.g. diffusion_models/my_custom_sub/), the same - spots under name_or_path when it is a local folder, then the hub — - downloaded to the repo-relative path under MODELS_PATH. - """ - override = self.model_config.model_kwargs.get(f"{component}_path", None) - if override is not None: - if not os.path.exists(override): - raise FileNotFoundError( - f"model_kwargs.{component}_path does not exist: {override}" - ) - return override - - rel_path = COMFY_FILES[component] - filename = os.path.basename(rel_path) - category = os.path.dirname(rel_path) - roots = [MODELS_PATH] + """Find a weight file at its local location (model_kwargs override, + comfy layout under MODELS_PATH or a local name_or_path dir), or + download it there when (and only when) it is missing — see + toolkit/models/v2/resolver.py for the search order.""" name_or_path = self.model_config.name_or_path - if name_or_path and os.path.isdir(name_or_path): - roots.append(name_or_path) - for root in roots: - for rel in (rel_path, filename): - candidate = os.path.join(root, rel) - if os.path.exists(candidate): - return candidate - for root in roots: - found = self._find_file_recursive(os.path.join(root, category), filename) - if found is not None: - return found - - import huggingface_hub - - repo_id = COMFY_REPO - if name_or_path and not os.path.exists(name_or_path) and "/" in name_or_path: - repo_id = name_or_path - self.print_and_status_update( - f"Downloading {rel_path} from {repo_id} into {MODELS_PATH}" + extra_roots = ( + [name_or_path] if name_or_path and os.path.isdir(name_or_path) else [] ) - return huggingface_hub.hf_hub_download( - repo_id=repo_id, filename=rel_path, local_dir=MODELS_PATH + return resolve_comfy_file( + COMFY_FILES[component], + repo_id=repo_id_from_name_or_path(name_or_path, COMFY_REPO), + override_path=self.model_config.model_kwargs.get( + f"{component}_path", None + ), + extra_roots=extra_roots, + status_fn=self.print_and_status_update, ) def _dit_component(self) -> str: @@ -322,7 +286,7 @@ class MinimaxH3Model(BaseModel): lora_path = self.model_config.assistant_lora_path if not os.path.exists(lora_path): filename = os.path.basename(lora_path) - found = self._find_file_recursive( + found = find_file_recursive( os.path.join(MODELS_PATH, "loras"), filename ) if found is not None: @@ -1226,20 +1190,9 @@ class MinimaxH3Model(BaseModel): "*adaln_proj*", ] - def convert_lora_weights_before_save(self, state_dict): - # ComfyUI's MiniMax-H3 keys are the original checkpoint keys, so the - # standard diffusion_model prefix maps directly - return { - k.replace("transformer.", "diffusion_model."): v - for k, v in state_dict.items() - } - - def convert_lora_weights_before_load(self, state_dict): - return { - k.replace("diffusion_model.", "transformer."): v - for k, v in state_dict.items() - } - + # ComfyUI's MiniMax-H3 keys are the original checkpoint keys, so the + # standard diffusion_model prefix maps directly + lora_keys_use_comfy_prefix = True class MinimaxH3Ref2VAModel(MinimaxH3Model): """Reference-to-video (ref2va): the control images ride along as reference diff --git a/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py b/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py index 572fda8..54ec7ff 100644 --- a/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py +++ b/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py @@ -405,16 +405,5 @@ class NucleusImageModel(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ["transformer_blocks"] - def convert_lora_weights_before_save(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd diff --git a/extensions_built_in/diffusion_models/omnigen2/__init__.py b/extensions_built_in/diffusion_models/omnigen2/__init__.py index 77ce910..edb10bd 100644 --- a/extensions_built_in/diffusion_models/omnigen2/__init__.py +++ b/extensions_built_in/diffusion_models/omnigen2/__init__.py @@ -343,21 +343,7 @@ class OmniGen2Model(BaseModel): return ["noise_refiner", "context_refiner", "ref_image_refiner", "layers"] return ["noise_refiner", "context_refiner", "layers"] - def convert_lora_weights_before_save(self, state_dict): - # currently starte with transformer. but needs to start with diffusion_model. for comfyui - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd - - def convert_lora_weights_before_load(self, state_dict): - # saved as diffusion_model. but needs to be transformer. for ai-toolkit - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True def get_base_model_version(self): return "omnigen2" diff --git a/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py b/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py index fb44ded..b6f2137 100644 --- a/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py +++ b/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py @@ -332,14 +332,5 @@ class PRXPixelT2IModel(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ["blocks"] - def convert_lora_weights_before_save(self, state_dict): - return { - k.replace("transformer.", "diffusion_model."): v - for k, v in state_dict.items() - } + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - return { - k.replace("diffusion_model.", "transformer."): v - for k, v in state_dict.items() - } diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py index c03af90..5199321 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py @@ -415,19 +415,7 @@ class QwenImageModel(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ["transformer_blocks"] - def convert_lora_weights_before_save(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd - - def convert_lora_weights_before_load(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None): if device is None: 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 bf13983..324b76a 100644 --- a/extensions_built_in/diffusion_models/z_image/z_image.py +++ b/extensions_built_in/diffusion_models/z_image/z_image.py @@ -31,8 +31,8 @@ try: from diffusers import ZImagePipeline # 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 + # quantization mixin (see toolkit/models/v2/_mixin.py) + from toolkit.models.v2.diffusion_models.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" @@ -464,16 +464,5 @@ class ZImageModel(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ["layers"] - def convert_lora_weights_before_save(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py index ca8b491..a4f3c0c 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py @@ -27,7 +27,6 @@ from .zeta_chroma_transformer import ZImageDCT, ZImageDCTParams, vae_flatten, va from .zeta_chroma_pipeline import ZetaChromaPipeline - scheduler_config = { "num_train_timesteps": 1000, "use_dynamic_shifting": False, @@ -368,16 +367,5 @@ class ZetaChromaModel(BaseModel): def get_transformer_block_names(self) -> Optional[List[str]]: return ["layers"] - def convert_lora_weights_before_save(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("transformer.", "diffusion_model.") - new_sd[new_key] = value - return new_sd + lora_keys_use_comfy_prefix = True - def convert_lora_weights_before_load(self, state_dict): - new_sd = {} - for key, value in state_dict.items(): - new_key = key.replace("diffusion_model.", "transformer.") - new_sd[new_key] = value - return new_sd diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index 9aadeea..ba4a1d7 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -98,6 +98,9 @@ UNET_IN_CHANNELS = 4 # Stable Diffusion の in_channels は 4 で固定。XLも class BaseModel: # override these in child classes arch = None + # rename LoRA keys transformer. <-> diffusion_model. (the ComfyUI-standard + # prefix) on save/load + lora_keys_use_comfy_prefix = False def __init__( self, @@ -1627,10 +1630,20 @@ class BaseModel: def convert_lora_weights_before_save(self, state_dict): # can be overridden in child classes to convert weights before saving + if self.lora_keys_use_comfy_prefix: + return { + k.replace("transformer.", "diffusion_model."): v + for k, v in state_dict.items() + } return state_dict - + def convert_lora_weights_before_load(self, state_dict): # can be overridden in child classes to convert weights before loading + if self.lora_keys_use_comfy_prefix: + return { + k.replace("diffusion_model.", "transformer."): v + for k, v in state_dict.items() + } return state_dict def condition_noisy_latents(self, latents: torch.Tensor, batch:'DataLoaderBatchDTO'): diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md new file mode 100644 index 0000000..796723f --- /dev/null +++ b/toolkit/models/v2/PLANNING.md @@ -0,0 +1,227 @@ +# v2 Model Module Restructure — Planning + +## Goal + +Every component the toolkit loads (DiTs/transformers, unets, text encoders, vision +encoders, VAEs, audio VAEs) becomes a class extending one base module in +`toolkit/models/v2`. `BaseModel` (`toolkit/models/base_model.py`) stays as the +multimodal holder that each arch in `extensions_built_in/diffusion_models` extends — +that layer is good. The layer below it is what gets unified: one loading entry point, +one quantization path, one save path, shared component definitions instead of +per-model-folder copies. + +End state this enables: + +- **Live server with model hot-swap**: a resident process where, when a generation or + training run requests a different model, the unused components are dropped and the + new ones loaded. Shared component classes (same TE/VAE reused across archs) make + component-level reuse possible instead of full teardown/reload. +- **Model loading test suite**: a test that loads each registered arch one at a time + and runs an inference pass. Every model type gets added to this suite as it is + migrated (see Testing below). +- **Comfy-aligned weights**: weights live in the ComfyUI folder layout under + `MODELS_PATH` (shareable with a ComfyUI install), download there when missing, and + saves are comfy-format. Eventually defaults move to comfy / our own prequantized + releases for everything. + +## Current state (survey 2026-08-27) + +Three generations of loading conventions coexist: + +1. Legacy monolith `toolkit/stable_diffusion_model.py` (`is_flux` / `is_v3` branches); + still the silent fallback in `toolkit/util/get_model.py` when an arch string + doesn't match. +2. `BaseModel` subclass per arch (35 registered classes), each with a hand-written + `load_model()` / `save_model()`. +3. `toolkit/models/v2/_mixin.py` (`OstrisModelMixin`) — the intended fix, currently + used by one model (`v2/z_image.py` → z_image extension). + +### Duplication highlights + +- BFL KL autoencoder: full copies in `flux2/src/autoencoder.py` and + `ideogram4/src/vae.py` (header says "Flux2 KL autoencoder"). +- Qwen3-VL text encoder loaded independently in qwen_image, nucleus_image, krea2, + ideogram4, mageflow, minimax_h3 — from several different repo sources; krea2 + hand-patches the vision tower locally. +- `Qwen3ForCausalLM` TE load: 3 verbatim line-for-line copies + (`z_image/z_image.py:257`, `z_image/z_image_l2p_model.py:471`, + `zeta_chroma/zeta_chroma_model.py:146`). +- Flux1 VAE + T5 + CLIP trio loaded 4x from 2 different repos (chroma ×2, + flux_kontext, legacy SD path). +- Comfy-file resolver copy-pasted: `minimax_h3/minimax_h3.py:248` + (`_resolve_comfy_file`) → `ltx2/ltx2.py:1254` ("mirrors MinimaxH3Model"). +- `AutoencoderKLQwenImage` latents mean/std handling triplicated (qwen_image, + nucleus_image, krea2). +- `transformer.` ↔ `diffusion_model.` LoRA key rename copy-pasted ~15x in + `convert_lora_weights_before_save/load` overrides. +- Fake CLIP/TE/config stubs redefined in ~5 places + (canonical: `toolkit/models/FakeVAE.py`, `toolkit/unloader.py`). + +### Inconsistency highlights + +- **Quantization, 5 paths**: `quantize_model()` (block-streaming, ARA-aware — ~21 + users), raw `quantize()` (~10 users, no block streaming/excludes, but the only path + honoring `quantize_kwargs`), hidream's hand-rolled block loop, the v2 mixin's own + `quantize_`, and bare TE quantization everywhere. +- **Known bugs**: ~9 sites quantize the TE with `qtype` instead of `qtype_te` + (chroma ×2, flux2, flux_kontext, cogview4, wan21, legacy SD, ...); + `toolkit/models/loaders/umt5.py` accepts a `comfy_files` param it never uses, so + wan21's comfy-TE path is a silent no-op. +- **Saving, 4 incompatible styles**: diffusers `save_pretrained` folders, flat + safetensors, safetensors-inside-diffusers-folder hybrids, and z_image's + loaded-format-dependent branch. Dequant-on-save done 3 ways; the + `isinstance(v, QTensor)` variant (chroma, flux2, boogu_image, ideogram4) misses + torchao and Ostris weights entirely. Every `save_pretrained` override ignores its + `save_dtype` argument. +- Registry: linear scan in `toolkit/util/get_model.py`, silent SD1 fallback on a + typo'd arch, eager import of every model file at startup. Second unsynchronized + registry in `ui/src/app/jobs/new/options.tsx`. + +## Decisions (locked in) + +1. **Save format = ComfyUI format.** Single-file safetensors in comfy key layout. + Must support saving quantized — primarily convrot8 and nvfp4 (comfy_quant marker + format, see `toolkit/util/comfy_quant_import.py`) — and plain bf16, all in comfy + format. Loading stays backwards compatible: diffusers dirs, transformers repos, + and single files all still digest through `load_model`; only saving standardizes + on comfy. +2. **v2 folder layout mirrors the comfy save path structure**: + + ``` + toolkit/models/v2/ + _mixin.py # base module (OstrisModelMixin, evolving) + resolver.py # comfy-layout weight resolution (lift from minimax_h3) + diffusion_models/ # one file per DiT/unet family + text_encoders/ # qwen3_vl.py, qwen3.py, t5.py, clip.py, gemma.py, ... + vae/ # flux_kl.py, qwen_image.py, wan.py, audio VAEs, ... + vision_encoders/ + ``` + +3. **Method names win from `BaseModel`**: `get_transformer_block_names` and + `get_quantization_exclude_modules`. The mixin's `get_quantization_block_names` + gets renamed to match; resolve the classmethod-vs-instance-method mismatch while + doing so. +4. **Loading policy**: + - Per-model special handling is allowed via the hook methods. + - If `name_or_path` is a diffusers/transformers source, load it with + diffusers/transformers for now. **Step 1 is migrating every model to the v2 + module format and loader without breaking anything** — same weights, same + sources, same results. + - Each model declares a `comfy_weight_names` dict keyed per standard + `name_or_path`. If the user points at a local folder or a non-standard repo, + load it as-is. If `name_or_path` is the standard repo and we have matching + comfy weight names, load those instead when any of them exist (locally under + `MODELS_PATH` in comfy layout, or downloadable to there). + - Eventually the default flips to comfy weights / our own prequantized releases + for everything. + +## Base module: what `OstrisModelMixin` still needs + +The mixin already handles: diffusers dir / hub repo / local single file / +`org/repo/file.safetensors`, key-conversion hooks on load and save, overridable +backend hooks for transformers-lib models, block-wise quantize. + +To add: + +- [x] **Comfy weight spec + resolver.** `aitk_comfy_repo` / `aitk_comfy_weight_names` + class attrs + `find_comfy_weights` (local-only until Phase 2); resolution chain + generalized from `minimax_h3._resolve_comfy_file` into `v2/resolver.py`: + explicit override → `MODELS_PATH` at the repo-relative comfy path → flat at + root → recursive walk of the category folder → hub download **to the + repo-relative path** (folder stays shareable with ComfyUI, no duplicate + downloads). +- [x] **Automatic prequantized import.** Single-file path sniffs `comfy_quant` + markers and routes through `import_comfy_quantized_layers` before + `load_state_dict`, including the OstrisLinear missing-key whitelist that + minimax_h3 and ltx2 each hand-rolled. +- [x] **One save path.** `save_model(path, dtype)`: dequantize via + `dequantize_if_quantized` (honors dtype), run `convert_state_dict_on_save`, + write single-file comfy-layout safetensors. (Quantized-storage saves — + convrot8 / nvfp4 with comfy_quant markers — land with Phase 2; diffusers-folder + save as an explicit flag still to add.) +- [x] **Tokenizer/processor declaration** for text encoders + (`aitk_tokenizer_repo`/`aitk_processor_repo` + `load_tokenizer`/`load_processor`). +- [x] Rename quantization hooks to the `BaseModel` spellings (decision 3): + `get_transformer_block_names` (classmethod on the module). + +## Migration steps + +Track progress here; check items off as they land. + +### Phase 0 — foundation (done 2026-08-27) +- [x] Evolve `_mixin.py` per the list above (comfy spec local-only until Phase 2) +- [x] Create `v2/diffusion_models/`, `v2/text_encoders/`, `v2/vae/`, + `v2/vision_encoders/`; move `v2/z_image.py` → `v2/diffusion_models/z_image.py` +- [x] Lift the comfy resolver out of minimax_h3 into `v2/resolver.py`; point + minimax_h3 and ltx2 at it (delete their copies) +- [x] `BaseModel` default `convert_lora_weights_before_save/load` doing the + `transformer.` ↔ `diffusion_model.` rename, gated on the class attr + `lora_keys_use_comfy_prefix` (default False, so passthrough models keep + their behavior); the ~18 identical overrides replaced with the flag. + Custom conversions (anima, hidream_o1, ltx2, wan21) keep their overrides; + ltx2's now composes with the flag via super(). + +### Phase 1 — migrate all models to v2 modules, no behavior change +Every arch's components become v2 classes; if `name_or_path` is diffusers, it still +loads via diffusers. Nothing about sources or outputs changes yet. Suggested order +(worst duplication first), each including its loading test (see Testing): + +- [ ] Shared TEs: `text_encoders/qwen3_vl.py`, `text_encoders/qwen3.py`, + `text_encoders/t5.py`, `text_encoders/clip.py` +- [ ] Shared VAEs: `vae/flux_kl.py` (kills flux2 + ideogram4 copies and the 4 + scattered `AutoencoderKL` loads), `vae/qwen_image.py` (mean/std handling + built in) +- [ ] z_image / z_image_l2p (already half-migrated) +- [ ] qwen_image family (qwen_image, qwen_image_edit, qwen_image_edit_plus) +- [ ] nucleus_image, krea2, ideogram4, mageflow +- [ ] chroma, chroma_radiance, zeta_chroma +- [ ] flux2, flux_kontext +- [ ] minimax_h3 (+ ref2va), ltx2 family +- [ ] wan21 / wan22 family +- [ ] hidream family, omnigen2 +- [ ] anima, boogu_image, ernie_image, f_light, prx_pixel_t2i +- [ ] audio_models (ace_step) +- [ ] Per-model fixes folded in as each migrates: `qtype_te` bug, dequant-on-save + (`dequantize_if_quantized` everywhere), raw-`quantize()` → `quantize_model()` + +### Phase 2 — comfy weights become the preferred source +- [ ] Wire `comfy_weight_names` per model; standard-repo `name_or_path` + existing + comfy weights → load comfy +- [ ] Comfy-format save (bf16 + convrot8/nvfp4 quantized) as the default + full-weight save +- [ ] Publish/verify comfy repacks per model as they flip + +### Phase 3 — live server +- [ ] Component-level identity (which TE/VAE instances are shared between archs) so + a model switch drops only what the next run doesn't need +- [ ] Resident process: request comes in → diff requested components vs loaded → + unload/load the difference +- [ ] Legacy `stable_diffusion_model.py` archs: grandfather or port last + +## Testing + +- [ ] `testing/` (or `tests/`) harness: for each migrated arch, load the model via + its v2 modules and run one small inference pass (single low-step sample; video + models at minimum frame count). One arch at a time, full unload between archs. +- [ ] Weights resolved through the normal resolver against `MODELS_PATH` + (GPU + local-weights test, not CI-portable at first; skip archs whose weights + are absent rather than failing). +- [ ] Round-trip test per model: load → save comfy format → reload from the save → + outputs match (bf16) / load cleanly (quantized saves). +- [ ] Each newly migrated model adds its test in the same PR as its migration. + +## TODO / look at later + +- [ ] Quantize-path consolidation quirks: `quantize_kwargs` is honored only by the + raw `quantize()` call sites and silently dropped by `quantize_model()`; the + ARA path inside `quantize_model` hardcodes `uint8`. Decide the unified + behavior when consolidating. +- [ ] `toolkit/models/loaders/umt5.py` dead `comfy_files` param (wan21 comfy-TE + no-op) — fix when wan migrates. +- [ ] Registry hardening: error (don't fall back to SD1) on unknown arch; lazy + per-arch imports; single source of truth shared with the UI's + `options.tsx` model list. +- [ ] Fake/stub components: consolidate on `toolkit/models/FakeVAE.py` / + `toolkit/unloader.py`, delete local copies. +- [ ] Vendored upstream code (hidream/src, omnigen2/src, ltx2 converter's private + comfy-quant parser): dedupe against toolkit utils where practical. diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index 6fb71ec..4e82640 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -9,7 +9,11 @@ digests any of: - 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=...`). +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 @@ -47,6 +51,19 @@ class OstrisModelMixin: # .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 @@ -68,7 +85,7 @@ class OstrisModelMixin: return state_dict @classmethod - def get_quantization_block_names(cls) -> Optional[List[str]]: + 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 @@ -200,16 +217,132 @@ class OstrisModelMixin: 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) - + has_quant_markers = any(k.endswith(".comfy_quant") for k in state_dict) model = cls.aitk_from_config(config) - model.load_state_dict(state_dict, assign=True) - model.to(dtype=dtype) + + 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, **kwargs): + from transformers import AutoTokenizer + + if cls.aitk_tokenizer_repo is None: + raise ValueError(f"{cls.__name__} does not declare aitk_tokenizer_repo") + return AutoTokenizer.from_pretrained( + cls.aitk_tokenizer_repo, subfolder=cls.aitk_tokenizer_subfolder, **kwargs + ) + + @classmethod + def load_processor(cls, **kwargs): + from transformers import AutoProcessor + + if cls.aitk_processor_repo is None: + raise ValueError(f"{cls.__name__} does not declare aitk_processor_repo") + return AutoProcessor.from_pretrained( + cls.aitk_processor_repo, subfolder=cls.aitk_processor_subfolder, **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 # ------------------------------------------------------------------ @@ -222,7 +355,7 @@ class OstrisModelMixin: 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 + 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 @@ -238,7 +371,7 @@ class OstrisModelMixin: ) blocks: List[torch.nn.Module] = [] - for name in self.get_quantization_block_names() or []: + 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("."): diff --git a/toolkit/models/v2/diffusion_models/__init__.py b/toolkit/models/v2/diffusion_models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/toolkit/models/v2/z_image.py b/toolkit/models/v2/diffusion_models/z_image.py similarity index 98% rename from toolkit/models/v2/z_image.py rename to toolkit/models/v2/diffusion_models/z_image.py index 90df66d..5ee32c4 100644 --- a/toolkit/models/v2/z_image.py +++ b/toolkit/models/v2/diffusion_models/z_image.py @@ -3,7 +3,7 @@ from diffusers.models.transformers import ( ZImageTransformer2DModel as DiffusersZImageTransformer2DModel, ) -from ._mixin import OstrisModelMixin +from .._mixin import OstrisModelMixin class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMixin): @@ -12,7 +12,7 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix aitk_config_repo = "Tongyi-MAI/Z-Image-Turbo" @classmethod - def get_quantization_block_names(cls): + def get_transformer_block_names(cls): return ["layers"] @classmethod diff --git a/toolkit/models/v2/resolver.py b/toolkit/models/v2/resolver.py new file mode 100644 index 0000000..fda8ff0 --- /dev/null +++ b/toolkit/models/v2/resolver.py @@ -0,0 +1,129 @@ +"""ComfyUI-layout weight file resolution. + +Weight files live under MODELS_PATH in ComfyUI's folder layout +(diffusion_models/, text_encoders/, vae/, ...) so the folder is shareable with +a ComfyUI install. Files are used in place when present and downloaded to +exactly their repo-relative location only when missing, so nothing is ever +duplicated on re-run. + +Lifted from the minimax_h3 / ltx2.5 model implementations; those now call +into here. +""" + +import os +from typing import Callable, Iterable, Optional + +from toolkit.paths import MODELS_PATH + + +def find_file_recursive(root_dir: str, filename: str) -> Optional[str]: + """First (breadth-stable, sorted) match of ``filename`` anywhere under + ``root_dir``.""" + if not os.path.isdir(root_dir): + return None + for dirpath, dirnames, filenames in os.walk(root_dir): + dirnames.sort() + if filename in filenames: + return os.path.join(dirpath, filename) + return None + + +def repo_id_from_name_or_path( + name_or_path: Optional[str], default: str +) -> str: + """Treat a hub-style ``name_or_path`` ("org/repo") as a replacement comfy + repo; anything local (or an explicit .safetensors file) keeps the + default.""" + if ( + name_or_path + and not os.path.exists(name_or_path) + and not name_or_path.endswith(".safetensors") + and "/" in name_or_path + ): + return name_or_path + return default + + +def resolve_comfy_file( + rel_path: str, + repo_id: str, + override_path: Optional[str] = None, + extra_roots: Optional[Iterable[str]] = None, + hf_token: Optional[str] = None, + status_fn: Optional[Callable[[str], None]] = None, + local_only: bool = False, +) -> Optional[str]: + """Find a weight file at its local location, or download it there when + (and only when) it is missing. + + Search order: ``override_path`` (must exist), the repo-relative path under + MODELS_PATH (and each of ``extra_roots``), the bare filename at each root, + any subfolder of the category folder (recursive — e.g. + diffusion_models/my_custom_sub/), then the hub — downloaded to the + repo-relative path under MODELS_PATH. With ``local_only`` the hub is never + touched and a miss returns None. + """ + if override_path is not None: + if not os.path.exists(override_path): + raise FileNotFoundError( + f"Override path for {rel_path} does not exist: {override_path}" + ) + return override_path + + filename = os.path.basename(rel_path) + category = os.path.dirname(rel_path) + roots = [MODELS_PATH] + [r for r in (extra_roots or []) if os.path.isdir(r)] + for root in roots: + for rel in (rel_path, filename): + candidate = os.path.join(root, rel) + if os.path.exists(candidate): + return candidate + for root in roots: + found = find_file_recursive(os.path.join(root, category), filename) + if found is not None: + return found + + if local_only: + return None + + import huggingface_hub + + if status_fn is not None: + status_fn(f"Downloading {rel_path} from {repo_id} into {MODELS_PATH}") + return huggingface_hub.hf_hub_download( + repo_id=repo_id, filename=rel_path, token=hf_token, local_dir=MODELS_PATH + ) + + +def resolve_named_file( + path: str, + component: str = "model", + hf_token: Optional[str] = None, +) -> str: + """Resolve an explicit .safetensors reference: a local file, a file already + under MODELS_PATH, or an 'org/repo/path/file.safetensors' hub path + (downloaded into the models folder at its repo-relative path).""" + if os.path.exists(path): + return path + splits = path.split("/") + if len(splits) < 3: + raise ValueError( + f"Invalid {component} path: {path}. Must be a local file or " + "'org/repo/filename.safetensors' to download from the Hugging Face Hub." + ) + rel_path = "/".join(splits[2:]) + for candidate in ( + os.path.join(MODELS_PATH, rel_path), + os.path.join(MODELS_PATH, splits[-1]), + ): + if os.path.exists(candidate): + return candidate + + import huggingface_hub + + return huggingface_hub.hf_hub_download( + repo_id="/".join(splits[:2]), + filename=rel_path, + token=hf_token, + local_dir=MODELS_PATH, + ) diff --git a/toolkit/models/v2/text_encoders/__init__.py b/toolkit/models/v2/text_encoders/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/toolkit/models/v2/vae/__init__.py b/toolkit/models/v2/vae/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/toolkit/models/v2/vision_encoders/__init__.py b/toolkit/models/v2/vision_encoders/__init__.py new file mode 100644 index 0000000..e69de29 From 8db198ec0a20ca430c49048fb3cb546b989ac5d8 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 27 Aug 2026 11:53:08 -0600 Subject: [PATCH 02/14] Phase 1 --- .../captioner/Qwen3VLCaptioner.py | 25 +- .../boogu_image/boogu_image.py | 32 +- .../boogu_image/src/transformer.py | 11 +- .../diffusion_models/chroma/chroma_model.py | 63 +-- .../chroma/chroma_radiance_model.py | 53 +-- .../diffusion_models/chroma/src/model.py | 35 +- .../diffusion_models/chroma/src/radiance.py | 35 +- .../ernie_image/ernie_image.py | 10 +- .../ernie_image/transformer.py | 10 +- .../diffusion_models/f_light/f_light.py | 35 +- .../diffusion_models/f_light/src/model.py | 10 +- .../diffusion_models/flux2/flux2_model.py | 6 +- .../diffusion_models/flux2/src/autoencoder.py | 435 ------------------ .../diffusion_models/flux2/src/pipeline.py | 2 +- .../flux_kontext/flux_kontext.py | 49 +- .../hidream/hidream_e1_model.py | 2 +- .../diffusion_models/hidream/hidream_model.py | 68 +-- .../transformers/transformer_hidream_image.py | 10 +- .../diffusion_models/ideogram4/ideogram4.py | 6 +- .../diffusion_models/krea2/krea2.py | 115 +---- .../diffusion_models/mageflow/mageflow.py | 13 +- .../mageflow/src/text_encoder.py | 27 +- .../nucleus_image/nucleus_image_model.py | 89 +--- .../diffusion_models/omnigen2/__init__.py | 4 +- .../transformers/transformer_omnigen2.py | 12 +- .../prx_pixel_t2i/prx_pixel_t2i.py | 4 +- .../prx_pixel_t2i/src/transformer_prx.py | 10 +- .../diffusion_models/qwen_image/qwen_image.py | 145 +----- .../diffusion_models/wan22/wan22_14b_model.py | 18 +- .../diffusion_models/z_image/z_image.py | 39 +- .../z_image/z_image_l2p_model.py | 33 +- .../zeta_chroma/zeta_chroma_model.py | 33 +- toolkit/models/base_model.py | 34 +- toolkit/models/loaders/umt5.py | 53 +-- toolkit/models/v2/PLANNING.md | 97 +++- toolkit/models/v2/_mixin.py | 94 +++- toolkit/models/v2/diffusion_models/flux.py | 14 + toolkit/models/v2/diffusion_models/hidream.py | 18 + .../v2/diffusion_models/nucleus_image.py | 15 + .../models/v2/diffusion_models/qwen_image.py | 29 ++ toolkit/models/v2/diffusion_models/wan.py | 13 + toolkit/models/v2/text_encoders/clip.py | 27 ++ toolkit/models/v2/text_encoders/qwen25_vl.py | 22 + toolkit/models/v2/text_encoders/qwen3.py | 17 + toolkit/models/v2/text_encoders/qwen3_vl.py | 53 +++ toolkit/models/v2/text_encoders/t5.py | 16 + toolkit/models/v2/text_encoders/umt5.py | 49 ++ toolkit/models/v2/vae/autoencoder_kl.py | 10 + .../models/v2/vae/flux2_kl.py | 131 ++++-- toolkit/models/v2/vae/qwen_image.py | 86 ++++ toolkit/models/wan21/wan21.py | 41 +- 51 files changed, 985 insertions(+), 1273 deletions(-) delete mode 100644 extensions_built_in/diffusion_models/flux2/src/autoencoder.py create mode 100644 toolkit/models/v2/diffusion_models/flux.py create mode 100644 toolkit/models/v2/diffusion_models/hidream.py create mode 100644 toolkit/models/v2/diffusion_models/nucleus_image.py create mode 100644 toolkit/models/v2/diffusion_models/qwen_image.py create mode 100644 toolkit/models/v2/diffusion_models/wan.py create mode 100644 toolkit/models/v2/text_encoders/clip.py create mode 100644 toolkit/models/v2/text_encoders/qwen25_vl.py create mode 100644 toolkit/models/v2/text_encoders/qwen3.py create mode 100644 toolkit/models/v2/text_encoders/qwen3_vl.py create mode 100644 toolkit/models/v2/text_encoders/t5.py create mode 100644 toolkit/models/v2/text_encoders/umt5.py create mode 100644 toolkit/models/v2/vae/autoencoder_kl.py rename extensions_built_in/diffusion_models/ideogram4/src/vae.py => toolkit/models/v2/vae/flux2_kl.py (83%) create mode 100644 toolkit/models/v2/vae/qwen_image.py diff --git a/extensions_built_in/captioner/Qwen3VLCaptioner.py b/extensions_built_in/captioner/Qwen3VLCaptioner.py index 3c8a5c4..b4ec37e 100644 --- a/extensions_built_in/captioner/Qwen3VLCaptioner.py +++ b/extensions_built_in/captioner/Qwen3VLCaptioner.py @@ -12,6 +12,8 @@ from optimum.quanto import freeze from toolkit.basic import flush from toolkit.util.quantize import quantize, get_qtype +from toolkit.models.v2.text_encoders.qwen3_vl import patch_qwen_vl_patch_embed + from .BaseCaptioner import BaseCaptioner import transformers import logging @@ -19,29 +21,6 @@ import traceback import warnings -def patch_qwen_vl_patch_embed(model): - """Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain - linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and - falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear - (a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves. - Returns the number of patch_embed modules patched.""" - patched = 0 - for module in model.modules(): - proj = getattr(module, "proj", None) - if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple( - proj.stride - ): - - def fast_forward(hidden_states, _proj=proj): - w = _proj.weight.reshape(_proj.weight.shape[0], -1) - x = hidden_states.view(-1, w.shape[1]).to(w.dtype) - return F.linear(x, w, _proj.bias) - - module.forward = fast_forward - patched += 1 - return patched - - # transformers.logging.set_verbosity_error() warnings.filterwarnings("ignore") logging.disable(logging.WARNING) diff --git a/extensions_built_in/diffusion_models/boogu_image/boogu_image.py b/extensions_built_in/diffusion_models/boogu_image/boogu_image.py index 8241561..5f70298 100644 --- a/extensions_built_in/diffusion_models/boogu_image/boogu_image.py +++ b/extensions_built_in/diffusion_models/boogu_image/boogu_image.py @@ -30,6 +30,7 @@ from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds from toolkit.basic import flush from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.models.base_model import BaseModel +from toolkit.models.v2.text_encoders.qwen3_vl import patch_qwen_vl_patch_embed from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) @@ -71,33 +72,6 @@ SYSTEM_PROMPT_T2I = ( HF_TOKEN = os.getenv("HF_TOKEN", None) -def patch_qwen_vl_patch_embed(model) -> int: - """Swap Qwen-VL's vision ``patch_embed`` Conv3d for the equivalent ``F.linear``. - - Qwen-VL's patch_embed is a Conv3d whose kernel == stride, i.e. just a linear - projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and - falls back to a slow path that effectively locks up image caching for the edit - (TI2I) model. The weight is read lazily so this survives later ``.to()`` moves. - Returns the number of patch_embed modules patched. (Vendored from - extensions_built_in/captioner/Qwen3VLCaptioner.py.) - """ - patched = 0 - for module in model.modules(): - proj = getattr(module, "proj", None) - if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple( - proj.stride - ): - - def fast_forward(hidden_states, _proj=proj): - w = _proj.weight.reshape(_proj.weight.shape[0], -1) - x = hidden_states.view(-1, w.shape[1]).to(w.dtype) - return F.linear(x, w, _proj.bias) - - module.forward = fast_forward - patched += 1 - return patched - - class BooguImageModel(BaseModel): arch = "boogu_image" # Default HF repo when model.name_or_path is unset (overridden by the edit model). @@ -167,8 +141,8 @@ class BooguImageModel(BaseModel): # deserialize -- use the bf16 repo and let ai-toolkit quantize if wanted. self.print_and_status_update("Loading transformer") try: - transformer = BooguImageTransformer2DModel.from_pretrained( - base, subfolder="transformer", torch_dtype=dtype, token=HF_TOKEN + transformer = BooguImageTransformer2DModel.load_model( + base, dtype=dtype, token=HF_TOKEN ) except OSError as e: raise OSError( diff --git a/extensions_built_in/diffusion_models/boogu_image/src/transformer.py b/extensions_built_in/diffusion_models/boogu_image/src/transformer.py index ce93406..090a089 100644 --- a/extensions_built_in/diffusion_models/boogu_image/src/transformer.py +++ b/extensions_built_in/diffusion_models/boogu_image/src/transformer.py @@ -21,6 +21,8 @@ from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.nn as nn from diffusers.configuration_utils import ConfigMixin, register_to_config + +from toolkit.models.v2._mixin import OstrisModelMixin from diffusers.loaders import PeftAdapterMixin from diffusers.loaders.single_file_model import FromOriginalModelMixin from diffusers.models.attention_processor import Attention @@ -490,10 +492,17 @@ class BooguImageDoubleStreamTransformerBlock(nn.Module): class BooguImageTransformer2DModel( - ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin + ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, OstrisModelMixin ): """Boogu-Image transformer with mixed double-stream -> single-stream topology.""" + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["double_stream_layers", "single_stream_layers"] + + _supports_gradient_checkpointing = True _no_split_modules = [ "BooguImageTransformerBlock", diff --git a/extensions_built_in/diffusion_models/chroma/chroma_model.py b/extensions_built_in/diffusion_models/chroma/chroma_model.py index 15da572..0cd0166 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_model.py @@ -5,8 +5,9 @@ import torch from toolkit.config_modules import GenerateImageConfig, ModelConfig from PIL import Image from toolkit.models.base_model import BaseModel +from toolkit.models.v2.text_encoders.t5 import T5TextEncoder +from toolkit.models.v2.vae.autoencoder_kl import KLVAE from toolkit.basic import flush -from diffusers import AutoencoderKL # from toolkit.pixel_shuffle_encoder import AutoencoderPixelMixer from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler @@ -14,7 +15,6 @@ from toolkit.dequantize import patch_dequantization_on_save from toolkit.accelerator import unwrap_model from optimum.quanto import freeze, QTensor from toolkit.util.quantize import quantize, get_qtype -from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer from .pipeline import ChromaPipeline, prepare_latent_image_ids from einops import rearrange, repeat import random @@ -148,37 +148,17 @@ class ChromaModel(BaseModel): self.print_and_status_update("Loading transformer") - chroma_state_dict = load_file(model_path, 'cpu') - - # determine number of double and single blocks - double_blocks = 0 - single_blocks = 0 - for key in chroma_state_dict.keys(): - if "double_blocks" in key: - block_num = int(key.split(".")[1]) + 1 - if block_num > double_blocks: - double_blocks = block_num - elif "single_blocks" in key: - block_num = int(key.split(".")[1]) + 1 - if block_num > single_blocks: - single_blocks = block_num - print(f"Double Blocks: {double_blocks}") - print(f"Single Blocks: {single_blocks}") - - chroma_params.depth = double_blocks - chroma_params.depth_single_blocks = single_blocks - transformer = Chroma(chroma_params) - + if model_path.endswith(".safetensors"): + transformer = Chroma.load_model(model_path, dtype=dtype) + else: + transformer = Chroma.load_from_state_dict(load_file(model_path, "cpu"), dtype) # add dtype, not sure why it doesnt have it transformer.dtype = dtype - # load the state dict into the model - transformer.load_state_dict(chroma_state_dict) - transformer.to(self.quantize_device, dtype=dtype) - + transformer.config = FakeConfig() - transformer.config.num_layers = double_blocks - transformer.config.num_single_layers = single_blocks + transformer.config.num_layers = transformer.params.depth + transformer.config.num_single_layers = transformer.params.depth_single_blocks if self.model_config.quantize: # patch the state dict method @@ -195,21 +175,9 @@ class ChromaModel(BaseModel): flush() self.print_and_status_update("Loading T5") - tokenizer_2 = T5TokenizerFast.from_pretrained( - extras_path, subfolder="tokenizer_2", torch_dtype=dtype - ) - text_encoder_2 = T5EncoderModel.from_pretrained( - extras_path, subfolder="text_encoder_2", torch_dtype=dtype - ) - text_encoder_2.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing T5") - quantize(text_encoder_2, weights=get_qtype( - self.model_config.qtype)) - freeze(text_encoder_2) - flush() + tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path) + text_encoder_2 = T5TextEncoder.load_model(extras_path, dtype=dtype) + self.prepare_text_encoder(text_encoder_2, dtype=dtype) # self.print_and_status_update("Loading CLIP") text_encoder = FakeCLIP() @@ -219,12 +187,7 @@ class ChromaModel(BaseModel): self.noise_scheduler = ChromaModel.get_train_scheduler() self.print_and_status_update("Loading VAE") - vae = AutoencoderKL.from_pretrained( - extras_path, - subfolder="vae", - torch_dtype=dtype - ) - vae = vae.to(self.device_torch, dtype=dtype) + vae = KLVAE.load_model(extras_path, dtype=dtype, device=self.device_torch) self.print_and_status_update("Making pipe") diff --git a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py index e5a79fe..0a6133b 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py @@ -5,8 +5,8 @@ import torch from toolkit.config_modules import GenerateImageConfig, ModelConfig from PIL import Image from toolkit.models.base_model import BaseModel +from toolkit.models.v2.text_encoders.t5 import T5TextEncoder from toolkit.basic import flush -from diffusers import AutoencoderKL # from toolkit.pixel_shuffle_encoder import AutoencoderPixelMixer from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler @@ -14,7 +14,6 @@ from toolkit.dequantize import patch_dequantization_on_save from toolkit.accelerator import unwrap_model from optimum.quanto import freeze, QTensor from toolkit.util.quantize import quantize, get_qtype -from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer from .pipeline import ChromaPipeline, prepare_latent_image_ids from einops import rearrange, repeat import random @@ -152,38 +151,16 @@ class ChromaRadianceModel(BaseModel): if model_path.endswith('.pth') or model_path.endswith('.pt'): chroma_state_dict = torch.load(model_path, map_location='cpu', weights_only=True) + transformer = Chroma.load_from_state_dict(chroma_state_dict, dtype) else: - chroma_state_dict = load_file(model_path, 'cpu') - - # determine number of double and single blocks - double_blocks = 0 - single_blocks = 0 - for key in chroma_state_dict.keys(): - if "double_blocks" in key: - block_num = int(key.split(".")[1]) + 1 - if block_num > double_blocks: - double_blocks = block_num - elif "single_blocks" in key: - block_num = int(key.split(".")[1]) + 1 - if block_num > single_blocks: - single_blocks = block_num - print(f"Double Blocks: {double_blocks}") - print(f"Single Blocks: {single_blocks}") - - chroma_params.depth = double_blocks - chroma_params.depth_single_blocks = single_blocks - transformer = Chroma(chroma_params) - + transformer = Chroma.load_model(model_path, dtype=dtype) # add dtype, not sure why it doesnt have it transformer.dtype = dtype - # load the state dict into the model - transformer.load_state_dict(chroma_state_dict) - transformer.to(self.quantize_device, dtype=dtype) - + transformer.config = FakeConfig() - transformer.config.num_layers = double_blocks - transformer.config.num_single_layers = single_blocks + transformer.config.num_layers = transformer.params.depth + transformer.config.num_single_layers = transformer.params.depth_single_blocks if self.model_config.quantize: # patch the state dict method @@ -200,21 +177,9 @@ class ChromaRadianceModel(BaseModel): flush() self.print_and_status_update("Loading T5") - tokenizer_2 = T5TokenizerFast.from_pretrained( - extras_path, subfolder="tokenizer_2", torch_dtype=dtype - ) - text_encoder_2 = T5EncoderModel.from_pretrained( - extras_path, subfolder="text_encoder_2", torch_dtype=dtype - ) - text_encoder_2.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing T5") - quantize(text_encoder_2, weights=get_qtype( - self.model_config.qtype)) - freeze(text_encoder_2) - flush() + tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path) + text_encoder_2 = T5TextEncoder.load_model(extras_path, dtype=dtype) + self.prepare_text_encoder(text_encoder_2, dtype=dtype) # self.print_and_status_update("Loading CLIP") text_encoder = FakeCLIP() diff --git a/extensions_built_in/diffusion_models/chroma/src/model.py b/extensions_built_in/diffusion_models/chroma/src/model.py index ebdf69d..5ef4dcf 100644 --- a/extensions_built_in/diffusion_models/chroma/src/model.py +++ b/extensions_built_in/diffusion_models/chroma/src/model.py @@ -1,4 +1,6 @@ -from dataclasses import dataclass +from dataclasses import dataclass, replace + +from toolkit.models.v2._mixin import OstrisModelMixin import torch from torch import Tensor, nn @@ -86,11 +88,40 @@ def modify_mask_to_attend_padding(mask, max_seq_length, num_extra_padding=8): return modified_mask -class Chroma(nn.Module): +class Chroma(nn.Module, OstrisModelMixin): """ Transformer model for flow matching on sequences. """ + @classmethod + def aitk_config_from_state_dict(cls, state_dict): + # block counts come from the checkpoint's key indices + double_blocks = 0 + single_blocks = 0 + for key in state_dict.keys(): + if "double_blocks" in key: + block_num = int(key.split(".")[1]) + 1 + if block_num > double_blocks: + double_blocks = block_num + elif "single_blocks" in key: + block_num = int(key.split(".")[1]) + 1 + if block_num > single_blocks: + single_blocks = block_num + print(f"Double Blocks: {double_blocks}") + print(f"Single Blocks: {single_blocks}") + return replace( + chroma_params, depth=double_blocks, depth_single_blocks=single_blocks + ) + + @classmethod + def aitk_from_config(cls, config): + with torch.device("meta"): + return cls(config) + + @classmethod + def get_transformer_block_names(cls): + return ["double_blocks", "single_blocks"] + def __init__(self, params: ChromaParams): super().__init__() self.params = params diff --git a/extensions_built_in/diffusion_models/chroma/src/radiance.py b/extensions_built_in/diffusion_models/chroma/src/radiance.py index d328f26..e6f7be7 100644 --- a/extensions_built_in/diffusion_models/chroma/src/radiance.py +++ b/extensions_built_in/diffusion_models/chroma/src/radiance.py @@ -1,4 +1,6 @@ -from dataclasses import dataclass +from dataclasses import dataclass, replace + +from toolkit.models.v2._mixin import OstrisModelMixin import torch from torch import Tensor, nn @@ -100,11 +102,40 @@ def modify_mask_to_attend_padding(mask, max_seq_length, num_extra_padding=8): return modified_mask -class Chroma(nn.Module): +class Chroma(nn.Module, OstrisModelMixin): """ Transformer model for flow matching on sequences. """ + @classmethod + def aitk_config_from_state_dict(cls, state_dict): + # block counts come from the checkpoint's key indices + double_blocks = 0 + single_blocks = 0 + for key in state_dict.keys(): + if "double_blocks" in key: + block_num = int(key.split(".")[1]) + 1 + if block_num > double_blocks: + double_blocks = block_num + elif "single_blocks" in key: + block_num = int(key.split(".")[1]) + 1 + if block_num > single_blocks: + single_blocks = block_num + print(f"Double Blocks: {double_blocks}") + print(f"Single Blocks: {single_blocks}") + return replace( + chroma_params, depth=double_blocks, depth_single_blocks=single_blocks + ) + + @classmethod + def aitk_from_config(cls, config): + with torch.device("meta"): + return cls(config) + + @classmethod + def get_transformer_block_names(cls): + return ["double_blocks", "single_blocks"] + def __init__(self, params: ChromaParams): super().__init__() self.params = params diff --git a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py index bb296f6..c9d76a6 100644 --- a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py +++ b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py @@ -79,20 +79,14 @@ class ErnieImageModel(BaseModel): self.print_and_status_update("Loading transformer") - transformer_path = model_path - transformer_subfolder = "transformer" - if os.path.exists(transformer_path): - transformer_subfolder = None - transformer_path = os.path.join(transformer_path, "transformer") + if 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 - transformer = ErnieImageTransformer2DModel.from_pretrained( - transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype - ) + transformer = ErnieImageTransformer2DModel.load_model(model_path, dtype=dtype) if self.model_config.quantize: self.print_and_status_update("Quantizing Transformer") diff --git a/extensions_built_in/diffusion_models/ernie_image/transformer.py b/extensions_built_in/diffusion_models/ernie_image/transformer.py index 3d27efe..ca3a2ae 100644 --- a/extensions_built_in/diffusion_models/ernie_image/transformer.py +++ b/extensions_built_in/diffusion_models/ernie_image/transformer.py @@ -33,6 +33,8 @@ from diffusers.models.attention_dispatch import dispatch_attention_fn from diffusers.models.attention_processor import Attention from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.modeling_utils import ModelMixin + +from toolkit.models.v2._mixin import OstrisModelMixin from diffusers.models.normalization import RMSNorm @@ -288,8 +290,14 @@ class ErnieImageAdaLNContinuous(nn.Module): return x -class ErnieImageTransformer2DModel(ModelMixin, ConfigMixin): +class ErnieImageTransformer2DModel(ModelMixin, ConfigMixin, OstrisModelMixin): _supports_gradient_checkpointing = True + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["layers"] + @register_to_config def __init__( diff --git a/extensions_built_in/diffusion_models/f_light/f_light.py b/extensions_built_in/diffusion_models/f_light/f_light.py index 813e6c3..2f3795c 100644 --- a/extensions_built_in/diffusion_models/f_light/f_light.py +++ b/extensions_built_in/diffusion_models/f_light/f_light.py @@ -6,15 +6,15 @@ import yaml from toolkit.config_modules import GenerateImageConfig, ModelConfig from PIL import Image from toolkit.models.base_model import BaseModel +from toolkit.models.v2.text_encoders.t5 import T5TextEncoder +from toolkit.models.v2.vae.autoencoder_kl import KLVAE from toolkit.basic import flush -from diffusers import AutoencoderKL from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler from toolkit.dequantize import patch_dequantization_on_save from toolkit.accelerator import unwrap_model from optimum.quanto import freeze, QTensor from toolkit.util.quantize import quantize, get_qtype -from transformers import T5TokenizerFast, T5EncoderModel from .src import FLitePipeline, DiT if TYPE_CHECKING: @@ -74,11 +74,7 @@ class FLiteModel(BaseModel): self.print_and_status_update("Loading transformer") - transformer = DiT.from_pretrained( - model_path, - subfolder="dit_model", - torch_dtype=dtype, - ) + transformer = DiT.load_model(model_path, dtype=dtype) transformer.to(self.quantize_device, dtype=dtype) @@ -97,31 +93,16 @@ class FLiteModel(BaseModel): flush() self.print_and_status_update("Loading T5") - tokenizer = T5TokenizerFast.from_pretrained( - extras_path, subfolder="tokenizer", torch_dtype=dtype + tokenizer = T5TextEncoder.load_tokenizer(extras_path, subfolder="tokenizer") + text_encoder = T5TextEncoder.load_model( + extras_path, dtype=dtype, subfolder="text_encoder" ) - text_encoder = T5EncoderModel.from_pretrained( - extras_path, subfolder="text_encoder", torch_dtype=dtype - ) - text_encoder.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing T5") - quantize(text_encoder, weights=get_qtype( - self.model_config.qtype)) - freeze(text_encoder) - flush() + self.prepare_text_encoder(text_encoder, dtype=dtype) self.noise_scheduler = FLiteModel.get_train_scheduler() self.print_and_status_update("Loading VAE") - vae = AutoencoderKL.from_pretrained( - extras_path, - subfolder="vae", - torch_dtype=dtype - ) - vae = vae.to(self.device_torch, dtype=dtype) + vae = KLVAE.load_model(extras_path, dtype=dtype, device=self.device_torch) self.print_and_status_update("Making pipe") diff --git a/extensions_built_in/diffusion_models/f_light/src/model.py b/extensions_built_in/diffusion_models/f_light/src/model.py index 903d492..26b3a9c 100644 --- a/extensions_built_in/diffusion_models/f_light/src/model.py +++ b/extensions_built_in/diffusion_models/f_light/src/model.py @@ -7,6 +7,8 @@ import torch.nn.functional as F from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin from diffusers.models.modeling_utils import ModelMixin + +from toolkit.models.v2._mixin import OstrisModelMixin from diffusers.utils.accelerate_utils import apply_forward_hook from einops import rearrange from peft import get_peft_model_state_dict, set_peft_model_state_dict @@ -302,7 +304,13 @@ def apply_rotary_emb(x, cos, sin): return torch.cat([y1, y2], 3).to(dtype=orig_dtype) -class DiT(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin): # type: ignore[misc] +class DiT(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin, OstrisModelMixin): # type: ignore[misc] + aitk_subfolder = "dit_model" + + @classmethod + def get_transformer_block_names(cls): + return ["blocks"] + _supports_gradient_checkpointing = True @register_to_config diff --git a/extensions_built_in/diffusion_models/flux2/flux2_model.py b/extensions_built_in/diffusion_models/flux2/flux2_model.py index e39e468..20aedf0 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_model.py @@ -21,7 +21,11 @@ from toolkit.util.quantize import quantize, get_qtype, quantize_model from transformers import AutoProcessor, Mistral3ForConditionalGeneration from .src.model import Flux2, Flux2Params from .src.pipeline import Flux2Pipeline -from .src.autoencoder import AutoEncoder, AutoEncoderParams, AutoEncoderSmallDecoderParams +from toolkit.models.v2.vae.flux2_kl import ( + AutoEncoder, + AutoEncoderParams, + AutoEncoderSmallDecoderParams, +) from safetensors.torch import load_file, save_file from PIL import Image import torch.nn.functional as F diff --git a/extensions_built_in/diffusion_models/flux2/src/autoencoder.py b/extensions_built_in/diffusion_models/flux2/src/autoencoder.py deleted file mode 100644 index 3d8230e..0000000 --- a/extensions_built_in/diffusion_models/flux2/src/autoencoder.py +++ /dev/null @@ -1,435 +0,0 @@ -from dataclasses import dataclass, field - -import torch -from einops import rearrange -from torch import Tensor, nn -import math -import torch.utils.checkpoint as ckpt - - -@dataclass -class AutoEncoderParams: - resolution: int = 256 - in_channels: int = 3 - ch: int = 128 - out_ch: int = 3 - ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4]) - num_res_blocks: int = 2 - z_channels: int = 32 - -@dataclass -class AutoEncoderSmallDecoderParams: - resolution: int = 256 - in_channels: int = 3 - ch: int = 128 - ch_encoder: int = 96 - out_ch: int = 3 - ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4]) - num_res_blocks: int = 2 - z_channels: int = 32 - - -def swish(x: Tensor) -> Tensor: - return x * torch.sigmoid(x) - - -class AttnBlock(nn.Module): - def __init__(self, in_channels: int): - super().__init__() - self.in_channels = in_channels - - self.norm = nn.GroupNorm( - num_groups=32, num_channels=in_channels, eps=1e-6, affine=True - ) - - self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1) - self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1) - self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1) - self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1) - - def attention(self, h_: Tensor) -> Tensor: - h_ = self.norm(h_) - q = self.q(h_) - k = self.k(h_) - v = self.v(h_) - - b, c, h, w = q.shape - q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous() - k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous() - v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous() - h_ = nn.functional.scaled_dot_product_attention(q, k, v) - - return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b) - - def forward(self, x: Tensor) -> Tensor: - return x + self.proj_out(self.attention(x)) - - -class ResnetBlock(nn.Module): - def __init__(self, in_channels: int, out_channels: int): - super().__init__() - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - - self.norm1 = nn.GroupNorm( - num_groups=32, num_channels=in_channels, eps=1e-6, affine=True - ) - self.conv1 = nn.Conv2d( - in_channels, out_channels, kernel_size=3, stride=1, padding=1 - ) - self.norm2 = nn.GroupNorm( - num_groups=32, num_channels=out_channels, eps=1e-6, affine=True - ) - self.conv2 = nn.Conv2d( - out_channels, out_channels, kernel_size=3, stride=1, padding=1 - ) - if self.in_channels != self.out_channels: - self.nin_shortcut = nn.Conv2d( - in_channels, out_channels, kernel_size=1, stride=1, padding=0 - ) - - def forward(self, x): - h = x - h = self.norm1(h) - h = swish(h) - h = self.conv1(h) - - h = self.norm2(h) - h = swish(h) - h = self.conv2(h) - - if self.in_channels != self.out_channels: - x = self.nin_shortcut(x) - - return x + h - - -class Downsample(nn.Module): - def __init__(self, in_channels: int): - super().__init__() - # no asymmetric padding in torch conv, must do it ourselves - self.conv = nn.Conv2d( - in_channels, in_channels, kernel_size=3, stride=2, padding=0 - ) - - def forward(self, x: Tensor): - pad = (0, 1, 0, 1) - x = nn.functional.pad(x, pad, mode="constant", value=0) - x = self.conv(x) - return x - - -class Upsample(nn.Module): - def __init__(self, in_channels: int): - super().__init__() - self.conv = nn.Conv2d( - in_channels, in_channels, kernel_size=3, stride=1, padding=1 - ) - - def forward(self, x: Tensor): - x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") - x = self.conv(x) - return x - - -class Encoder(nn.Module): - def __init__( - self, - resolution: int, - in_channels: int, - ch: int, - ch_mult: list[int], - num_res_blocks: int, - z_channels: int, - ): - super().__init__() - self.quant_conv = torch.nn.Conv2d(2 * z_channels, 2 * z_channels, 1) - self.ch = ch - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.in_channels = in_channels - # downsampling - self.conv_in = nn.Conv2d( - in_channels, self.ch, kernel_size=3, stride=1, padding=1 - ) - - curr_res = resolution - in_ch_mult = (1,) + tuple(ch_mult) - self.in_ch_mult = in_ch_mult - self.down = nn.ModuleList() - block_in = self.ch - for i_level in range(self.num_resolutions): - block = nn.ModuleList() - attn = nn.ModuleList() - block_in = ch * in_ch_mult[i_level] - block_out = ch * ch_mult[i_level] - for _ in range(self.num_res_blocks): - block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) - block_in = block_out - down = nn.Module() - down.block = block - down.attn = attn - if i_level != self.num_resolutions - 1: - down.downsample = Downsample(block_in) - curr_res = curr_res // 2 - self.down.append(down) - - # middle - self.mid = nn.Module() - self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) - self.mid.attn_1 = AttnBlock(block_in) - self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) - - # end - self.norm_out = nn.GroupNorm( - num_groups=32, num_channels=block_in, eps=1e-6, affine=True - ) - self.conv_out = nn.Conv2d( - block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1 - ) - self.gradient_checkpointing = False - - def enable_gradient_checkpointing(self): - self.gradient_checkpointing = True - - def forward(self, x: Tensor) -> Tensor: - # downsampling - hs = [self.conv_in(x)] - for i_level in range(self.num_resolutions): - for i_block in range(self.num_res_blocks): - if torch.is_grad_enabled() and self.gradient_checkpointing: - h = ckpt.checkpoint(self.down[i_level].block[i_block], hs[-1]) - if len(self.down[i_level].attn) > 0: - h = ckpt.checkpoint(self.down[i_level].attn[i_block], h) - else: - h = self.down[i_level].block[i_block](hs[-1]) - if len(self.down[i_level].attn) > 0: - h = self.down[i_level].attn[i_block](h) - hs.append(h) - if i_level != self.num_resolutions - 1: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hs.append(ckpt.checkpoint(self.down[i_level].downsample, hs[-1])) - else: - hs.append(self.down[i_level].downsample(hs[-1])) - - # middle - h = hs[-1] - if torch.is_grad_enabled() and self.gradient_checkpointing: - h = ckpt.checkpoint(self.mid.block_1, h) - h = ckpt.checkpoint(self.mid.attn_1, h) - h = ckpt.checkpoint(self.mid.block_2, h) - else: - h = self.mid.block_1(h) - h = self.mid.attn_1(h) - h = self.mid.block_2(h) - # end - h = self.norm_out(h) - h = swish(h) - h = self.conv_out(h) - h = self.quant_conv(h) - return h - - -class Decoder(nn.Module): - def __init__( - self, - ch: int, - out_ch: int, - ch_mult: list[int], - num_res_blocks: int, - in_channels: int, - resolution: int, - z_channels: int, - ): - super().__init__() - self.post_quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1) - self.ch = ch - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.in_channels = in_channels - self.ffactor = 2 ** (self.num_resolutions - 1) - - # compute in_ch_mult, block_in and curr_res at lowest res - block_in = ch * ch_mult[self.num_resolutions - 1] - curr_res = resolution // 2 ** (self.num_resolutions - 1) - self.z_shape = (1, z_channels, curr_res, curr_res) - - # z to block_in - self.conv_in = nn.Conv2d( - z_channels, block_in, kernel_size=3, stride=1, padding=1 - ) - - # middle - self.mid = nn.Module() - self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) - self.mid.attn_1 = AttnBlock(block_in) - self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) - - # upsampling - self.up = nn.ModuleList() - for i_level in reversed(range(self.num_resolutions)): - block = nn.ModuleList() - attn = nn.ModuleList() - block_out = ch * ch_mult[i_level] - for _ in range(self.num_res_blocks + 1): - block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) - block_in = block_out - up = nn.Module() - up.block = block - up.attn = attn - if i_level != 0: - up.upsample = Upsample(block_in) - curr_res = curr_res * 2 - self.up.insert(0, up) # prepend to get consistent order - - # end - self.norm_out = nn.GroupNorm( - num_groups=32, num_channels=block_in, eps=1e-6, affine=True - ) - self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1) - self.gradient_checkpointing = False - - def enable_gradient_checkpointing(self): - self.gradient_checkpointing = True - - def forward(self, z: Tensor) -> Tensor: - z = self.post_quant_conv(z) - - # get dtype for proper tracing - upscale_dtype = next(self.up.parameters()).dtype - - # z to block_in - h = self.conv_in(z) - - # middle - if torch.is_grad_enabled() and self.gradient_checkpointing: - h = ckpt.checkpoint(self.mid.block_1, h) - h = ckpt.checkpoint(self.mid.attn_1, h) - h = ckpt.checkpoint(self.mid.block_2, h) - else: - h = self.mid.block_1(h) - h = self.mid.attn_1(h) - h = self.mid.block_2(h) - - # cast to proper dtype - h = h.to(upscale_dtype) - # upsampling - for i_level in reversed(range(self.num_resolutions)): - for i_block in range(self.num_res_blocks + 1): - if torch.is_grad_enabled() and self.gradient_checkpointing: - h = ckpt.checkpoint(self.up[i_level].block[i_block], h) - if len(self.up[i_level].attn) > 0: - h = ckpt.checkpoint(self.up[i_level].attn[i_block], h) - else: - h = self.up[i_level].block[i_block](h) - if len(self.up[i_level].attn) > 0: - h = self.up[i_level].attn[i_block](h) - if i_level != 0: - if torch.is_grad_enabled() and self.gradient_checkpointing: - h = ckpt.checkpoint(self.up[i_level].upsample, h) - else: - h = self.up[i_level].upsample(h) - - # end - h = self.norm_out(h) - h = swish(h) - h = self.conv_out(h) - return h - - -class AutoEncoder(nn.Module): - def __init__(self, params: AutoEncoderParams): - super().__init__() - self.params = params - self.encoder = Encoder( - resolution=params.resolution, - in_channels=params.in_channels, - ch=params.ch, - ch_mult=params.ch_mult, - num_res_blocks=params.num_res_blocks, - z_channels=params.z_channels, - ) - decoder_ch = params.ch - if hasattr(params, "ch_encoder"): - decoder_ch = params.ch_encoder - self.decoder = Decoder( - resolution=params.resolution, - in_channels=params.in_channels, - ch=decoder_ch, - out_ch=params.out_ch, - ch_mult=params.ch_mult, - num_res_blocks=params.num_res_blocks, - z_channels=params.z_channels, - ) - - self.bn_eps = 1e-4 - self.bn_momentum = 0.1 - self.ps = [2, 2] - self.bn = torch.nn.BatchNorm2d( - math.prod(self.ps) * params.z_channels, - eps=self.bn_eps, - momentum=self.bn_momentum, - affine=False, - track_running_stats=True, - ) - self._gradient_checkpointing = False - - @property - def gradient_checkpointing(self): - return self._gradient_checkpointing - - @gradient_checkpointing.setter - def gradient_checkpointing(self, value: bool): - self._gradient_checkpointing = value - self.encoder.gradient_checkpointing = value - self.decoder.gradient_checkpointing = value - - @property - def device(self): - return next(self.parameters()).device - - @property - def dtype(self): - return next(self.parameters()).dtype - - def enable_gradient_checkpointing(self): - self.gradient_checkpointing = True - self.encoder.enable_gradient_checkpointing() - self.decoder.enable_gradient_checkpointing() - - def normalize(self, z): - self.bn.eval() - return self.bn(z) - - def inv_normalize(self, z): - self.bn.eval() - s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.bn_eps) - m = self.bn.running_mean.view(1, -1, 1, 1) - return z * s + m - - def encode(self, x: Tensor) -> Tensor: - moments = self.encoder(x) - mean = torch.chunk(moments, 2, dim=1)[0] - - z = rearrange( - mean, - "... c (i pi) (j pj) -> ... (c pi pj) i j", - pi=self.ps[0], - pj=self.ps[1], - ) - z = self.normalize(z) - return z - - def decode(self, z: Tensor) -> Tensor: - z = self.inv_normalize(z) - z = rearrange( - z, - "... (c pi pj) i j -> ... c (i pi) (j pj)", - pi=self.ps[0], - pj=self.ps[1], - ) - dec = self.decoder(z) - return dec diff --git a/extensions_built_in/diffusion_models/flux2/src/pipeline.py b/extensions_built_in/diffusion_models/flux2/src/pipeline.py index 8e64639..aec4dd0 100644 --- a/extensions_built_in/diffusion_models/flux2/src/pipeline.py +++ b/extensions_built_in/diffusion_models/flux2/src/pipeline.py @@ -12,7 +12,7 @@ from diffusers.utils import ( from diffusers.utils.torch_utils import randn_tensor from diffusers.pipelines.pipeline_utils import DiffusionPipeline from diffusers.utils import BaseOutput -from .autoencoder import AutoEncoder +from toolkit.models.v2.vae.flux2_kl import AutoEncoder from .model import Flux2 from einops import rearrange from transformers import AutoProcessor, Mistral3ForConditionalGeneration diff --git a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py index e9eeee5..927477f 100644 --- a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py +++ b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py @@ -8,7 +8,11 @@ from toolkit import train_tools from toolkit.config_modules import GenerateImageConfig, ModelConfig from PIL import Image from toolkit.models.base_model import BaseModel -from diffusers import FluxTransformer2DModel, AutoencoderKL, FluxKontextPipeline +from toolkit.models.v2.text_encoders.t5 import T5TextEncoder +from toolkit.models.v2.text_encoders.clip import CLIPTextEncoder +from toolkit.models.v2.vae.autoencoder_kl import KLVAE +from diffusers import FluxKontextPipeline +from toolkit.models.v2.diffusion_models.flux import FluxTransformer2DModel from toolkit.basic import flush from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler @@ -18,7 +22,6 @@ from toolkit.accelerator import get_accelerator, unwrap_model from optimum.quanto import freeze, QTensor from toolkit.util.mask import generate_random_mask, random_dialate_mask from toolkit.util.quantize import quantize, get_qtype -from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer from einops import rearrange, repeat import random import torch.nn.functional as F @@ -80,11 +83,7 @@ class FluxKontextModel(BaseModel): # so we need this for the VAE, te, etc base_model_path = self.model_config.extras_name_or_path - transformer_path = model_path - transformer_subfolder = 'transformer' - if os.path.exists(transformer_path): - transformer_subfolder = None - transformer_path = os.path.join(transformer_path, 'transformer') + if 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 @@ -92,11 +91,7 @@ class FluxKontextModel(BaseModel): base_model_path = model_path self.print_and_status_update("Loading transformer") - transformer = FluxTransformer2DModel.from_pretrained( - transformer_path, - subfolder=transformer_subfolder, - torch_dtype=dtype - ) + transformer = FluxTransformer2DModel.load_model(model_path, dtype=dtype) transformer.to(self.quantize_device, dtype=dtype) if self.model_config.quantize: @@ -114,32 +109,18 @@ class FluxKontextModel(BaseModel): flush() self.print_and_status_update("Loading T5") - tokenizer_2 = T5TokenizerFast.from_pretrained( - base_model_path, subfolder="tokenizer_2", torch_dtype=dtype - ) - text_encoder_2 = T5EncoderModel.from_pretrained( - base_model_path, subfolder="text_encoder_2", torch_dtype=dtype - ) - text_encoder_2.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing T5") - quantize(text_encoder_2, weights=get_qtype( - self.model_config.qtype)) - freeze(text_encoder_2) - flush() + tokenizer_2 = T5TextEncoder.load_tokenizer(base_model_path) + text_encoder_2 = T5TextEncoder.load_model(base_model_path, dtype=dtype) + self.prepare_text_encoder(text_encoder_2, dtype=dtype) self.print_and_status_update("Loading CLIP") - text_encoder = CLIPTextModel.from_pretrained( - base_model_path, subfolder="text_encoder", torch_dtype=dtype) - tokenizer = CLIPTokenizer.from_pretrained( - base_model_path, subfolder="tokenizer", torch_dtype=dtype) - text_encoder.to(self.device_torch, dtype=dtype) + text_encoder = CLIPTextEncoder.load_model( + base_model_path, dtype=dtype, device=self.device_torch + ) + tokenizer = CLIPTextEncoder.load_tokenizer(base_model_path, use_fast=False) self.print_and_status_update("Loading VAE") - vae = AutoencoderKL.from_pretrained( - base_model_path, subfolder="vae", torch_dtype=dtype) + vae = KLVAE.load_model(base_model_path, dtype=dtype) self.noise_scheduler = FluxKontextModel.get_train_scheduler() diff --git a/extensions_built_in/diffusion_models/hidream/hidream_e1_model.py b/extensions_built_in/diffusion_models/hidream/hidream_e1_model.py index 5306ad5..2981807 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_e1_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_e1_model.py @@ -7,7 +7,7 @@ from toolkit.accelerator import unwrap_model import torch from toolkit.prompt_utils import PromptEmbeds from toolkit.config_modules import GenerateImageConfig -from diffusers.models import HiDreamImageTransformer2DModel +from toolkit.models.v2.diffusion_models.hidream import HiDreamImageTransformer2DModel import torch.nn.functional as F from PIL import Image diff --git a/extensions_built_in/diffusion_models/hidream/hidream_model.py b/extensions_built_in/diffusion_models/hidream/hidream_model.py index 71921aa..3d72f04 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_model.py @@ -9,7 +9,10 @@ from toolkit import train_tools from toolkit.config_modules import GenerateImageConfig, ModelConfig from PIL import Image from toolkit.models.base_model import BaseModel -from diffusers import AutoencoderKL, TorchAoConfig +from toolkit.models.v2.text_encoders.t5 import T5TextEncoder +from toolkit.models.v2.text_encoders.clip import CLIPTextEncoderWithProjection +from toolkit.models.v2.vae.autoencoder_kl import KLVAE +from diffusers import TorchAoConfig from toolkit.basic import flush from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler @@ -19,7 +22,6 @@ from toolkit.accelerator import get_accelerator, unwrap_model from optimum.quanto import freeze, QTensor from toolkit.util.mask import generate_random_mask, random_dialate_mask from toolkit.util.quantize import quantize, get_qtype -from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer, TorchAoConfig as TorchAoConfigTransformers from .src.pipelines.hidream_image.pipeline_hidream_image import HiDreamImagePipeline from .src.models.transformers.transformer_hidream_image import HiDreamImageTransformer2DModel from .src.schedulers.fm_solvers_unipc import FlowUniPCMultistepScheduler @@ -28,15 +30,6 @@ from einops import rearrange, repeat import random import torch.nn.functional as F from tqdm import tqdm -from transformers import ( - CLIPTextModelWithProjection, - CLIPTokenizer, - T5EncoderModel, - T5Tokenizer, - LlamaForCausalLM, - PreTrainedTokenizerFast -) - if TYPE_CHECKING: from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO @@ -125,10 +118,8 @@ class HidreamModel(BaseModel): self.print_and_status_update("Loading transformer") - transformer = self.hidream_transformer_class.from_pretrained( - model_path, - subfolder="transformer", - torch_dtype=torch.bfloat16 + transformer = self.hidream_transformer_class.load_model( + model_path, dtype=torch.bfloat16 ) if not self.low_vram: @@ -164,44 +155,34 @@ class HidreamModel(BaseModel): self.print_and_status_update("Loading vae") - vae = AutoencoderKL.from_pretrained( - extras_path, - subfolder="vae", - torch_dtype=torch.bfloat16 - ).to(self.device_torch, dtype=dtype) + vae = KLVAE.load_model(extras_path, dtype=torch.bfloat16).to( + self.device_torch, dtype=dtype + ) self.print_and_status_update("Loading clip encoders") - text_encoder = CLIPTextModelWithProjection.from_pretrained( - extras_path, - subfolder="text_encoder", - torch_dtype=torch.bfloat16 + text_encoder = CLIPTextEncoderWithProjection.load_model( + extras_path, dtype=torch.bfloat16 ).to(self.device_torch, dtype=dtype) - - tokenizer = CLIPTokenizer.from_pretrained( - extras_path, - subfolder="tokenizer" + + tokenizer = CLIPTextEncoderWithProjection.load_tokenizer( + extras_path, use_fast=False ) - - text_encoder_2 = CLIPTextModelWithProjection.from_pretrained( - extras_path, - subfolder="text_encoder_2", - torch_dtype=torch.bfloat16 + + text_encoder_2 = CLIPTextEncoderWithProjection.load_model( + extras_path, dtype=torch.bfloat16, subfolder="text_encoder_2" ).to(self.device_torch, dtype=dtype) - - tokenizer_2 = CLIPTokenizer.from_pretrained( - extras_path, - subfolder="tokenizer_2" + + tokenizer_2 = CLIPTextEncoderWithProjection.load_tokenizer( + extras_path, subfolder="tokenizer_2", use_fast=False ) flush() self.print_and_status_update("Loading T5 encoders") - text_encoder_3 = T5EncoderModel.from_pretrained( - extras_path, - subfolder="text_encoder_3", - torch_dtype=torch.bfloat16 + text_encoder_3 = T5TextEncoder.load_model( + extras_path, dtype=torch.bfloat16, subfolder="text_encoder_3" ).to(self.device_torch, dtype=dtype) if self.model_config.quantize_te: @@ -211,9 +192,8 @@ class HidreamModel(BaseModel): freeze(text_encoder_3) flush() - tokenizer_3 = T5Tokenizer.from_pretrained( - extras_path, - subfolder="tokenizer_3" + tokenizer_3 = T5TextEncoder.load_tokenizer( + extras_path, subfolder="tokenizer_3", use_fast=False ) flush() diff --git a/extensions_built_in/diffusion_models/hidream/src/models/transformers/transformer_hidream_image.py b/extensions_built_in/diffusion_models/hidream/src/models/transformers/transformer_hidream_image.py index f7eb104..1c23bf7 100644 --- a/extensions_built_in/diffusion_models/hidream/src/models/transformers/transformer_hidream_image.py +++ b/extensions_built_in/diffusion_models/hidream/src/models/transformers/transformer_hidream_image.py @@ -8,6 +8,8 @@ from einops import repeat from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin from diffusers.models.modeling_utils import ModelMixin + +from toolkit.models.v2._mixin import OstrisModelMixin from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers from diffusers.utils.torch_utils import maybe_allow_in_graph from diffusers.models.modeling_outputs import Transformer2DModelOutput @@ -228,9 +230,15 @@ class HiDreamImageBlock(nn.Module): ) class HiDreamImageTransformer2DModel( - ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin + ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, OstrisModelMixin ): _supports_gradient_checkpointing = True + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["double_stream_blocks", "single_stream_blocks"] + _no_split_modules = ["HiDreamImageBlock"] @register_to_config diff --git a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py index 390f7a4..2048bdd 100644 --- a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py +++ b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py @@ -26,7 +26,11 @@ from huggingface_hub.errors import EntryNotFoundError from transformers import AutoModel, AutoTokenizer from .src.transformer import Ideogram4Config, Ideogram4Transformer2DModel -from .src.vae import AutoEncoder, AutoEncoderParams, convert_diffusers_state_dict +from toolkit.models.v2.vae.flux2_kl import ( + AutoEncoder, + AutoEncoderParams, + convert_diffusers_state_dict, +) from .src.latent_norm import get_latent_norm from .src.pipeline import ( Ideogram4Pipeline, diff --git a/extensions_built_in/diffusion_models/krea2/krea2.py b/extensions_built_in/diffusion_models/krea2/krea2.py index 22d75a8..b0fa0a6 100644 --- a/extensions_built_in/diffusion_models/krea2/krea2.py +++ b/extensions_built_in/diffusion_models/krea2/krea2.py @@ -25,18 +25,21 @@ from safetensors.torch import load_file, save_file import huggingface_hub from huggingface_hub.errors import EntryNotFoundError -from diffusers import AutoencoderKLQwenImage from transformers import ( AutoProcessor, AutoTokenizer, Qwen2TokenizerFast, - Qwen3VLForConditionalGeneration, ) from optimum.quanto import freeze from toolkit.config_modules import GenerateImageConfig, ModelConfig, NetworkConfig from toolkit.lora_special import LoRASpecialNetwork from toolkit.models.base_model import BaseModel +from toolkit.models.v2.vae.qwen_image import QwenImageVAE, QwenImageVAEHolderMixin +from toolkit.models.v2.text_encoders.qwen3_vl import ( + Qwen3VLTextEncoder, + patch_qwen_vl_patch_embed, +) from toolkit.basic import flush from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import ( @@ -100,30 +103,6 @@ QWEN_IMAGE_VAE_PATH = "Qwen/Qwen-Image" HF_TOKEN = os.getenv("HF_TOKEN", None) -def patch_qwen_vl_patch_embed(model): - """Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain - linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and - falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear - (a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves. - Returns the number of patch_embed modules patched. (Same patch as the - Qwen3VLCaptioner extension.)""" - patched = 0 - for module in model.modules(): - proj = getattr(module, "proj", None) - if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple( - proj.stride - ): - - def fast_forward(hidden_states, _proj=proj): - w = _proj.weight.reshape(_proj.weight.shape[0], -1) - x = hidden_states.view(-1, w.shape[1]).to(w.dtype) - return F.linear(x, w, _proj.bias) - - module.forward = fast_forward - patched += 1 - return patched - - def _load_mmdit_state_dict(name_or_path: str, filename: Optional[str]) -> dict: """Load the MMDiT weights from a local safetensors file/dir or the HF hub. @@ -163,7 +142,7 @@ def _load_mmdit_state_dict(name_or_path: str, filename: Optional[str]) -> dict: return load_file(path) -class Krea2Model(BaseModel): +class Krea2Model(QwenImageVAEHolderMixin, BaseModel): arch = "krea2" def __init__( @@ -268,8 +247,8 @@ class Krea2Model(BaseModel): processor = Qwen2TokenizerFast.from_pretrained( te_path, max_length=self.max_text_length, token=HF_TOKEN ) - text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( - te_path, torch_dtype=dtype, token=HF_TOKEN + text_encoder = Qwen3VLTextEncoder.load_model( + te_path, dtype=dtype, subfolder="", token=HF_TOKEN ) vl_processor = None if self.is_edit: @@ -281,8 +260,7 @@ class Krea2Model(BaseModel): else: # We only ever encode text, so the vision tower is dead weight -- drop it to # free VRAM and skip loading its (bf16-slow) Conv3d patch_embed onto the GPU. - if getattr(text_encoder.model, "visual", None) is not None: - text_encoder.model.visual = None + text_encoder.drop_vision_tower() text_encoder.eval() text_encoder.requires_grad_(False) flush() @@ -291,8 +269,8 @@ class Krea2Model(BaseModel): def _load_vae(self): vae_path = self.model_config.model_kwargs.get("vae_path", QWEN_IMAGE_VAE_PATH) self.print_and_status_update(f"Loading Qwen-Image VAE from {vae_path}") - vae = AutoencoderKLQwenImage.from_pretrained( - vae_path, subfolder="vae", torch_dtype=self.vae_torch_dtype, token=HF_TOKEN + vae = QwenImageVAE.load_model( + vae_path, dtype=self.vae_torch_dtype, token=HF_TOKEN ) vae.eval() vae.requires_grad_(False) @@ -772,76 +750,9 @@ class Krea2Model(BaseModel): return False # ------------------------------------------------------------------ - # VAE (Qwen-Image AutoencoderKLQwenImage -- same handling as qwen_image arch) + # VAE (Qwen-Image AutoencoderKLQwenImage -- shared QwenImageVAEHolderMixin) # ------------------------------------------------------------------ - def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None): - if device is None: - device = self.vae_device_torch - if dtype is None: - dtype = self.vae_torch_dtype - - if self.vae.device == torch.device("cpu"): - self.vae.to(device) - self.vae.eval() - self.vae.requires_grad_(False) - - image_list = [image.to(device, dtype=dtype) for image in image_list] - images = torch.stack(image_list).to(device, dtype=dtype) - - # AutoencoderKLQwenImage is a video VAE: add a frame dim. - images = images.unsqueeze(2) - latents = self.vae.encode(images).latent_dist.sample() - - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latents.device, latents.dtype) - ) - latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view( - 1, self.vae.config.z_dim, 1, 1, 1 - ).to(latents.device, latents.dtype) - - latents = (latents - latents_mean) * latents_std - latents = latents.squeeze(2) # drop frame dim - return latents.to(device, dtype=dtype) - - def decode_latents(self, latents: torch.Tensor, device=None, dtype=None): - if device is None: - device = self.vae_device_torch - if dtype is None: - dtype = self.vae_torch_dtype - - if self.vae.device == torch.device("cpu"): - self.vae.to(device) - - latents = latents.to(device, dtype=dtype) - latents = latents.unsqueeze(2) # add frame dim - - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latents.device, latents.dtype) - ) - latents = latents * latents_std + latents_mean - - # Full-resolution decode spikes VRAM; tile it when low on VRAM (decode - # only -- encode stays untiled). - tiled = self.model_config.low_vram - if tiled: - self.vae.enable_tiling() - try: - images = self.vae.decode(latents).sample - finally: - if tiled: - self.vae.disable_tiling() - images = images.squeeze(2) # drop frame dim - return images.to(device, dtype=dtype) - + vae_decode_tiled_on_low_vram = True # ------------------------------------------------------------------ # Saving / bookkeeping # ------------------------------------------------------------------ diff --git a/extensions_built_in/diffusion_models/mageflow/mageflow.py b/extensions_built_in/diffusion_models/mageflow/mageflow.py index 589f9f2..49b3530 100644 --- a/extensions_built_in/diffusion_models/mageflow/mageflow.py +++ b/extensions_built_in/diffusion_models/mageflow/mageflow.py @@ -34,7 +34,8 @@ from torchvision.transforms.functional import to_tensor from safetensors.torch import load_file, save_file import huggingface_hub -from transformers import AutoProcessor, AutoTokenizer, Qwen3VLForConditionalGeneration +from transformers import AutoProcessor, AutoTokenizer +from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder from optimum.quanto import freeze from toolkit.config_modules import GenerateImageConfig, ModelConfig @@ -209,8 +210,11 @@ class MageFlowModel(BaseModel): self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}") tokenizer = AutoTokenizer.from_pretrained(te_path, token=HF_TOKEN, **te_kwargs) - text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( - te_path, torch_dtype=dtype, token=HF_TOKEN, **te_kwargs + text_encoder = Qwen3VLTextEncoder.load_model( + te_path, + dtype=dtype, + subfolder=te_kwargs.get("subfolder", ""), + token=HF_TOKEN, ) vl_processor = None if self.is_edit: @@ -224,8 +228,7 @@ class MageFlowModel(BaseModel): else: # We only ever encode text, so the vision tower is dead weight -- # drop it to free VRAM. - if getattr(text_encoder.model, "visual", None) is not None: - text_encoder.model.visual = None + text_encoder.drop_vision_tower() text_encoder.eval() text_encoder.requires_grad_(False) flush() diff --git a/extensions_built_in/diffusion_models/mageflow/src/text_encoder.py b/extensions_built_in/diffusion_models/mageflow/src/text_encoder.py index 4f3508a..0f05ed6 100644 --- a/extensions_built_in/diffusion_models/mageflow/src/text_encoder.py +++ b/extensions_built_in/diffusion_models/mageflow/src/text_encoder.py @@ -18,7 +18,8 @@ import math from typing import List, Optional import torch -import torch.nn.functional as F + +from toolkit.models.v2.text_encoders.qwen3_vl import patch_qwen_vl_patch_embed # Prompt templates from the reference (mage_flow/models/utils.py). ``start_idx`` @@ -57,30 +58,6 @@ def edit_prompt_body(instruction: str, num_refs: int) -> str: return prefix + instruction -def patch_qwen_vl_patch_embed(model): - """Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain - linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and - falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear - (a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves. - Returns the number of patch_embed modules patched. (Same patch as the krea2 - extension / Qwen3VLCaptioner.)""" - patched = 0 - for module in model.modules(): - proj = getattr(module, "proj", None) - if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple( - proj.stride - ): - - def fast_forward(hidden_states, _proj=proj): - w = _proj.weight.reshape(_proj.weight.shape[0], -1) - x = hidden_states.view(-1, w.shape[1]).to(w.dtype) - return F.linear(x, w, _proj.bias) - - module.forward = fast_forward - patched += 1 - return patched - - def resize_vl_images( images: List[torch.Tensor], max_long_edge: int = 384 ) -> List["PIL.Image.Image"]: diff --git a/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py b/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py index 54ec7ff..ee80c39 100644 --- a/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py +++ b/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py @@ -6,21 +6,25 @@ import torch import yaml from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.models.base_model import BaseModel +from toolkit.models.v2.vae.qwen_image import QwenImageVAE, QwenImageVAEHolderMixin +from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder from toolkit.basic import flush from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.accelerator import unwrap_model -from optimum.quanto import freeze -from toolkit.util.quantize import quantize, get_qtype, quantize_model +from toolkit.util.quantize import quantize_model from toolkit.memory_management import MemoryManager -from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor +from transformers import Qwen3VLProcessor import torch.nn.functional as F try: - from diffusers import NucleusMoEImagePipeline, NucleusMoEImageTransformer2DModel, AutoencoderKLQwenImage + from diffusers import NucleusMoEImagePipeline, AutoencoderKLQwenImage + from toolkit.models.v2.diffusion_models.nucleus_image import ( + NucleusMoEImageTransformer2DModel, + ) from diffusers.models.transformers.transformer_nucleusmoe_image import SwiGLUExperts except ImportError: raise ImportError( @@ -46,7 +50,7 @@ scheduler_config = { } -class NucleusImageModel(BaseModel): +class NucleusImageModel(QwenImageVAEHolderMixin, BaseModel): arch = "nucleus_image" def __init__( @@ -81,20 +85,14 @@ class NucleusImageModel(BaseModel): self.print_and_status_update("Loading transformer") - transformer_path = model_path - transformer_subfolder = "transformer" - if os.path.exists(transformer_path): - transformer_subfolder = None - transformer_path = os.path.join(transformer_path, "transformer") + if 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 - transformer = NucleusMoEImageTransformer2DModel.from_pretrained( - transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype - ) + transformer = NucleusMoEImageTransformer2DModel.load_model(model_path, dtype=dtype) # handle versions of pytorch that don't have grouped mm, by disabling it in the SwiGLUExperts if not hasattr(torch.nn.functional, "grouped_mm"): @@ -129,33 +127,13 @@ class NucleusImageModel(BaseModel): tokenizer = Qwen3VLProcessor.from_pretrained( base_model_path, subfolder="processor", torch_dtype=dtype ) - text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( - base_model_path, subfolder="text_encoder", torch_dtype=dtype - ) - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - - text_encoder.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Text Encoder") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() + text_encoder = Qwen3VLTextEncoder.load_model(base_model_path, dtype=dtype) + self.prepare_text_encoder(text_encoder, dtype=dtype) self.print_and_status_update("Loading VAE") - vae = AutoencoderKLQwenImage.from_pretrained( - base_model_path, subfolder="vae", torch_dtype=dtype - ).to(self.device_torch, dtype=dtype) + vae = QwenImageVAE.load_model( + base_model_path, dtype=dtype, device=self.device_torch + ) self.noise_scheduler = NucleusImageModel.get_train_scheduler() @@ -214,41 +192,6 @@ class NucleusImageModel(BaseModel): return pipeline - def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None): - if device is None: - device = self.vae_device_torch - if dtype is None: - dtype = self.vae_torch_dtype - - # Move to vae to device if on cpu - if self.vae.device == torch.device("cpu"): - self.vae.to(device) - self.vae.eval() - self.vae.requires_grad_(False) - # move to device and dtype - image_list = [image.to(device, dtype=dtype) for image in image_list] - images = torch.stack(image_list).to(device, dtype=dtype) - # it uses wan vae, so add dim for frame count - - images = images.unsqueeze(2) - latents = self.vae.encode(images).latent_dist.sample() - - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latents.device, latents.dtype) - ) - latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view( - 1, self.vae.config.z_dim, 1, 1, 1 - ).to(latents.device, latents.dtype) - - latents = (latents - latents_mean) * latents_std - latents = latents.to(device, dtype=dtype) - - latents = latents.squeeze(2) # remove the frame count dimension - - return latents - def generate_single_image( self, pipeline: NucleusMoEImagePipeline, diff --git a/extensions_built_in/diffusion_models/omnigen2/__init__.py b/extensions_built_in/diffusion_models/omnigen2/__init__.py index edb10bd..77741be 100644 --- a/extensions_built_in/diffusion_models/omnigen2/__init__.py +++ b/extensions_built_in/diffusion_models/omnigen2/__init__.py @@ -96,8 +96,8 @@ class OmniGen2Model(BaseModel): self.print_and_status_update("Loading transformer") - transformer = OmniGen2Transformer2DModel.from_pretrained( - model_path, subfolder="transformer", torch_dtype=torch.bfloat16 + transformer = OmniGen2Transformer2DModel.load_model( + model_path, dtype=torch.bfloat16 ) if not self.low_vram: diff --git a/extensions_built_in/diffusion_models/omnigen2/src/models/transformers/transformer_omnigen2.py b/extensions_built_in/diffusion_models/omnigen2/src/models/transformers/transformer_omnigen2.py index 8e7fef6..abb3e0c 100644 --- a/extensions_built_in/diffusion_models/omnigen2/src/models/transformers/transformer_omnigen2.py +++ b/extensions_built_in/diffusion_models/omnigen2/src/models/transformers/transformer_omnigen2.py @@ -15,6 +15,8 @@ from diffusers.models.attention_processor import Attention from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.modeling_utils import ModelMixin +from toolkit.models.v2._mixin import OstrisModelMixin + from ..attention_processor import OmniGen2AttnProcessorFlash2Varlen, OmniGen2AttnProcessor from .repo import OmniGen2RotaryPosEmbed from .block_lumina2 import LuminaLayerNormContinuous, LuminaRMSNormZero, LuminaFeedForward, Lumina2CombinedTimestepCaptionEmbedding @@ -177,7 +179,15 @@ class OmniGen2TransformerBlock(nn.Module): return hidden_states -class OmniGen2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): +class OmniGen2Transformer2DModel( + ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, OstrisModelMixin +): + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["noise_refiner", "ref_image_refiner", "context_refiner", "layers"] + """ OmniGen2 Transformer 2D Model. diff --git a/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py b/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py index b6f2137..ffea929 100644 --- a/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py +++ b/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py @@ -126,9 +126,7 @@ class PRXPixelT2IModel(BaseModel): self.print_and_status_update("Loading transformer") # from_pretrained reads config.json (bottleneck_size, resolution_embeds, # in_channels=3, ...) and the safetensors in one shot. - transformer = PRXTransformer2DModel.from_pretrained( - model_path, subfolder="transformer", torch_dtype=dtype - ) + transformer = PRXTransformer2DModel.load_model(model_path, dtype=dtype) transformer.to(dtype=dtype) flush() diff --git a/extensions_built_in/diffusion_models/prx_pixel_t2i/src/transformer_prx.py b/extensions_built_in/diffusion_models/prx_pixel_t2i/src/transformer_prx.py index e83d8a1..1dd6dee 100644 --- a/extensions_built_in/diffusion_models/prx_pixel_t2i/src/transformer_prx.py +++ b/extensions_built_in/diffusion_models/prx_pixel_t2i/src/transformer_prx.py @@ -40,6 +40,8 @@ from diffusers.models.attention_dispatch import dispatch_attention_fn from diffusers.models.embeddings import get_timestep_embedding from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.modeling_utils import ModelMixin + +from toolkit.models.v2._mixin import OstrisModelMixin from diffusers.models.normalization import RMSNorm @@ -666,7 +668,7 @@ def seq2img(seq: torch.Tensor, patch_size: int, shape: torch.Tensor) -> torch.Te return seq -class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): +class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin, OstrisModelMixin): r""" Transformer-based 2D model for text to image generation. @@ -702,6 +704,12 @@ class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): `PRXResolutionEmbedder`. Used by the PRX-7B variant. """ + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["blocks"] + config_name = "config.json" _supports_gradient_checkpointing = True diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py index 5199321..54116a5 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py @@ -7,28 +7,25 @@ from toolkit import train_tools from toolkit.config_modules import GenerateImageConfig, ModelConfig from PIL import Image from toolkit.models.base_model import BaseModel +from toolkit.models.v2.vae.qwen_image import QwenImageVAE, QwenImageVAEHolderMixin from toolkit.basic import flush from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.accelerator import get_accelerator, unwrap_model -from optimum.quanto import freeze, QTensor -from toolkit.util.quantize import quantize, get_qtype, quantize_model +from toolkit.util.quantize import quantize_model import torch.nn.functional as F from toolkit.memory_management import MemoryManager from safetensors.torch import load_file from diffusers import ( QwenImagePipeline, - QwenImageTransformer2DModel, AutoencoderKLQwenImage, ) -from transformers import ( - Qwen2_5_VLForConditionalGeneration, - Qwen2Tokenizer, - Qwen2VLProcessor, -) +from transformers import Qwen2VLProcessor +from toolkit.models.v2.diffusion_models.qwen_image import QwenImageTransformer2DModel +from toolkit.models.v2.text_encoders.qwen25_vl import Qwen25VLTextEncoder from tqdm import tqdm from toolkit.util.qwen_vae_gradient_checkpointing import patch_qwen_vae_gradient_checkpointing @@ -55,7 +52,7 @@ scheduler_config = { } -class QwenImageModel(BaseModel): +class QwenImageModel(QwenImageVAEHolderMixin, BaseModel): arch = "qwen_image" _qwen_image_keep_visual = False _qwen_pipeline = QwenImagePipeline @@ -97,31 +94,14 @@ class QwenImageModel(BaseModel): self.print_and_status_update("Loading transformer") - if model_path.endswith(".safetensors"): - # load the safetensors file - transformer = QwenImageTransformer2DModel.from_single_file( - model_path, - config="Qwen/Qwen-Image", - subfolder="transformer", - torch_dtype=model_dtype, - ) - transformer.to(model_dtype) + if not model_path.endswith(".safetensors") and 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 - 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 = QwenImageTransformer2DModel.from_pretrained( - transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype - ) + transformer = QwenImageTransformer2DModel.load_model(model_path, dtype=model_dtype) if self.model_config.quantize: self.print_and_status_update("Quantizing Transformer") @@ -145,41 +125,18 @@ class QwenImageModel(BaseModel): flush() self.print_and_status_update("Text Encoder") - tokenizer = Qwen2Tokenizer.from_pretrained( - base_model_path, subfolder="tokenizer", torch_dtype=dtype - ) - text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( - base_model_path, subfolder="text_encoder", torch_dtype=dtype - ) + tokenizer = Qwen25VLTextEncoder.load_tokenizer(base_model_path, use_fast=False) + text_encoder = Qwen25VLTextEncoder.load_model(base_model_path, dtype=dtype) # remove the visual model as it is not needed for image generation self.processor = None if not self._qwen_image_keep_visual: - text_encoder.model.visual = None + text_encoder.drop_vision_tower() - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - - text_encoder.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Text Encoder") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() + self.prepare_text_encoder(text_encoder, dtype=dtype) self.print_and_status_update("Loading VAE") - vae = AutoencoderKLQwenImage.from_pretrained( - base_model_path, subfolder="vae", torch_dtype=dtype - ) + vae = QwenImageVAE.load_model(base_model_path, dtype=dtype) self.noise_scheduler = QwenImageModel.get_train_scheduler() @@ -417,69 +374,3 @@ class QwenImageModel(BaseModel): lora_keys_use_comfy_prefix = True - def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None): - if device is None: - device = self.vae_device_torch - if dtype is None: - dtype = self.vae_torch_dtype - - # Move to vae to device if on cpu - if self.vae.device == torch.device("cpu"): - self.vae.to(device) - self.vae.eval() - self.vae.requires_grad_(False) - # move to device and dtype - image_list = [image.to(device, dtype=dtype) for image in image_list] - images = torch.stack(image_list).to(device, dtype=dtype) - # it uses wan vae, so add dim for frame count - - images = images.unsqueeze(2) - latents = self.vae.encode(images).latent_dist.sample() - - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latents.device, latents.dtype) - ) - latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view( - 1, self.vae.config.z_dim, 1, 1, 1 - ).to(latents.device, latents.dtype) - - latents = (latents - latents_mean) * latents_std - latents = latents.to(device, dtype=dtype) - - latents = latents.squeeze(2) # remove the frame count dimension - - return latents - - def decode_latents(self, latents: torch.Tensor, device=None, dtype=None): - if device is None: - device = self.vae_device_torch - if dtype is None: - dtype = self.vae_torch_dtype - - if self.vae.device == torch.device("cpu"): - self.vae.to(device) - - latents = latents.to(device, dtype=dtype) - - # add frame count dim for wan vae - latents = latents.unsqueeze(2) - - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latents.device, latents.dtype) - ) - latents_std = ( - torch.tensor(self.vae.config.latents_std) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latents.device, latents.dtype) - ) - latents = latents * latents_std + latents_mean - - images = self.vae.decode(latents).sample - - images = images.squeeze(2) # remove the frame count dimension - - return images.to(device, dtype=dtype) diff --git a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py index 37bcb7a..3a26897 100644 --- a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py +++ b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py @@ -17,7 +17,7 @@ from toolkit.samplers.custom_flowmatch_sampler import ( ) from toolkit.util.quantize import quantize_model from .wan22_pipeline import Wan22Pipeline -from diffusers import WanTransformer3DModel +from toolkit.models.v2.diffusion_models.wan import WanTransformer3DModel from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO from torchvision.transforms import functional as TF @@ -287,11 +287,9 @@ class Wan2214bModel(Wan21): self.print_and_status_update("Loading transformer 1") dtype = self.torch_dtype - transformer_1 = WanTransformer3DModel.from_pretrained( - transformer_path_1, - subfolder=subfolder_1, - torch_dtype=dtype, - ).to(dtype=dtype) + transformer_1 = WanTransformer3DModel.load_model( + transformer_path_1, dtype=dtype, subfolder=subfolder_1 + ) flush() @@ -317,11 +315,9 @@ class Wan2214bModel(Wan21): self.print_and_status_update("Loading transformer 2") dtype = self.torch_dtype - transformer_2 = WanTransformer3DModel.from_pretrained( - transformer_path_2, - subfolder=subfolder_2, - torch_dtype=dtype, - ).to(dtype=dtype) + transformer_2 = WanTransformer3DModel.load_model( + transformer_path_2, dtype=dtype, subfolder=subfolder_2 + ) flush() 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 324b76a..67ffd90 100644 --- a/extensions_built_in/diffusion_models/z_image/z_image.py +++ b/extensions_built_in/diffusion_models/z_image/z_image.py @@ -13,19 +13,16 @@ from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.accelerator import unwrap_model -from optimum.quanto import freeze from toolkit.util.quantize import ( - quantize, - get_qtype, quantize_model, dequantize_if_quantized, ) from toolkit.memory_management import MemoryManager from toolkit.metadata import get_meta_for_safetensors +from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder +from toolkit.models.v2.vae.autoencoder_kl import KLVAE from safetensors.torch import load_file, save_file -from transformers import AutoTokenizer, Qwen3ForCausalLM -from diffusers import AutoencoderKL try: from diffusers import ZImagePipeline @@ -255,36 +252,12 @@ class ZImageModel(BaseModel): flush() self.print_and_status_update("Text Encoder") - tokenizer = AutoTokenizer.from_pretrained( - base_model_path, subfolder="tokenizer", torch_dtype=dtype - ) - text_encoder = Qwen3ForCausalLM.from_pretrained( - base_model_path, subfolder="text_encoder", torch_dtype=dtype - ) - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - - text_encoder.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Text Encoder") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() + tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path) + text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype) + self.prepare_text_encoder(text_encoder, dtype=dtype) self.print_and_status_update("Loading VAE") - vae = AutoencoderKL.from_pretrained( - base_model_path, subfolder="vae", torch_dtype=dtype - ) + vae = KLVAE.load_model(base_model_path, dtype=dtype) self.noise_scheduler = ZImageModel.get_train_scheduler() diff --git a/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py b/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py index 34fd183..46ca0d0 100644 --- a/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py +++ b/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py @@ -9,11 +9,10 @@ import torch.nn.functional as F import yaml from toolkit.basic import flush from toolkit.accelerator import unwrap_model -from optimum.quanto import freeze -from toolkit.util.quantize import quantize, get_qtype, quantize_model +from toolkit.util.quantize import quantize_model from toolkit.memory_management import MemoryManager -from transformers import AutoTokenizer, Qwen3ForCausalLM +from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder from toolkit.models.FakeVAE import FakeVAE from toolkit.paths import MODELS_PATH from safetensors.torch import load_file, save_file @@ -468,31 +467,9 @@ class ZImageL2PModel(ZImageModel): flush() self.print_and_status_update("Text Encoder") - tokenizer = AutoTokenizer.from_pretrained( - base_model_path, subfolder="tokenizer", torch_dtype=dtype - ) - text_encoder = Qwen3ForCausalLM.from_pretrained( - base_model_path, subfolder="text_encoder", torch_dtype=dtype - ) - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - - text_encoder.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Text Encoder") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() + tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path) + text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype) + self.prepare_text_encoder(text_encoder, dtype=dtype) self.print_and_status_update("Loading VAE") # vae = AutoencoderKL.from_pretrained( diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py index a4f3c0c..b4a3bb3 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py @@ -13,14 +13,13 @@ from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.accelerator import unwrap_model -from optimum.quanto import freeze -from toolkit.util.quantize import quantize, get_qtype, quantize_model +from toolkit.util.quantize import quantize_model from toolkit.memory_management import MemoryManager from safetensors.torch import load_file from optimum.quanto import QTensor from toolkit.metadata import get_meta_for_safetensors from safetensors.torch import load_file, save_file -from transformers import AutoTokenizer, Qwen3ForCausalLM +from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder from diffusers import AutoencoderKL from toolkit.models.FakeVAE import FakeVAE from .zeta_chroma_transformer import ZImageDCT, ZImageDCTParams, vae_flatten, vae_unflatten, prepare_latent_image_ids, make_text_position_ids @@ -142,31 +141,9 @@ class ZetaChromaModel(BaseModel): flush() self.print_and_status_update("Text Encoder") - tokenizer = AutoTokenizer.from_pretrained( - base_model_path, subfolder="tokenizer", torch_dtype=dtype - ) - text_encoder = Qwen3ForCausalLM.from_pretrained( - base_model_path, subfolder="text_encoder", torch_dtype=dtype - ) - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - - text_encoder.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Text Encoder") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() + tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path) + text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype) + self.prepare_text_encoder(text_encoder, dtype=dtype) self.print_and_status_update("Loading VAE") vae = FakeVAE(scaling_factor=1.0) diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index ba4a1d7..8f06dfb 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -1627,7 +1627,39 @@ class BaseModel: encoder.to(*args, **kwargs) else: self.text_encoder.to(*args, **kwargs) - + + def prepare_text_encoder(self, text_encoder, dtype=None): + """Standard post-load text-encoder policy: layer offloading, device + placement, then quantize_te. Skips quantization when the checkpoint + loaded pre-quantized.""" + from optimum.quanto import freeze + from toolkit.memory_management import MemoryManager + from toolkit.util.quantize import get_qtype, quantize + + dtype = dtype if dtype is not None else self.torch_dtype + if ( + self.model_config.layer_offloading + and self.model_config.layer_offloading_text_encoder_percent > 0 + ): + MemoryManager.attach( + text_encoder, + self.device_torch, + offload_percent=self.model_config.layer_offloading_text_encoder_percent, + ) + + text_encoder.to(self.device_torch, dtype=dtype) + flush() + + if self.model_config.quantize_te and not getattr( + text_encoder, "aitk_is_quantized", False + ): + self.print_and_status_update("Quantizing Text Encoder") + quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) + freeze(text_encoder) + flush() + return text_encoder + + def convert_lora_weights_before_save(self, state_dict): # can be overridden in child classes to convert weights before saving if self.lora_keys_use_comfy_prefix: diff --git a/toolkit/models/loaders/umt5.py b/toolkit/models/loaders/umt5.py index f53db3c..3c7fd30 100644 --- a/toolkit/models/loaders/umt5.py +++ b/toolkit/models/loaders/umt5.py @@ -1,46 +1,29 @@ from typing import List -import torch -from transformers import T5Tokenizer, UMT5EncoderModel -class PatchedT5Tokenizer(T5Tokenizer): - def __init__( - self, - vocab: str | list[tuple[str, float]] | None = None, - eos_token="", - unk_token="", - pad_token="", - _spm_precompiled_charsmap=None, - extra_ids=100, - additional_special_tokens=None, - **kwargs, - ): - super().__init__( - vocab=vocab, - eos_token=eos_token, - unk_token=unk_token, - pad_token=pad_token, - _spm_precompiled_charsmap=None, # this is passing a empty byte string for some reason now. - extra_ids=extra_ids, - additional_special_tokens=additional_special_tokens, - **kwargs, - ) +import torch + +from toolkit.models.v2.text_encoders.umt5 import ( + PatchedT5Tokenizer, + UMT5TextEncoder, +) + def get_umt5_encoder( model_path: str, tokenizer_subfolder: str = None, encoder_subfolder: str = None, torch_dtype: str = torch.bfloat16, - comfy_files: List[str] = [ - "text_encoders/umt5_xxl_fp16.safetensors", - "text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors", - ], -) -> UMT5EncoderModel: - """ - Load the UMT5 encoder model from the specified path. - """ - tokenizer = PatchedT5Tokenizer.from_pretrained(model_path, subfolder=tokenizer_subfolder) + # reserved for the comfy-weights flip (Phase 2); accepted for + # signature compatibility, not consulted yet + comfy_files: List[str] = None, +): + """Load the UMT5 tokenizer + encoder. Thin compatibility wrapper around + toolkit/models/v2/text_encoders/umt5.py.""" + tokenizer = UMT5TextEncoder.load_tokenizer( + model_path, subfolder=tokenizer_subfolder or "" + ) print(f"Using {model_path} for UMT5 encoder.") - text_encoder = UMT5EncoderModel.from_pretrained( - model_path, subfolder=encoder_subfolder, torch_dtype=torch_dtype + text_encoder = UMT5TextEncoder.load_model( + model_path, dtype=torch_dtype, subfolder=encoder_subfolder or "" ) return tokenizer, text_encoder diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index 796723f..87e0c9b 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -166,21 +166,88 @@ Every arch's components become v2 classes; if `name_or_path` is diffusers, it st loads via diffusers. Nothing about sources or outputs changes yet. Suggested order (worst duplication first), each including its loading test (see Testing): -- [ ] Shared TEs: `text_encoders/qwen3_vl.py`, `text_encoders/qwen3.py`, - `text_encoders/t5.py`, `text_encoders/clip.py` -- [ ] Shared VAEs: `vae/flux_kl.py` (kills flux2 + ideogram4 copies and the 4 - scattered `AutoencoderKL` loads), `vae/qwen_image.py` (mean/std handling - built in) -- [ ] z_image / z_image_l2p (already half-migrated) -- [ ] qwen_image family (qwen_image, qwen_image_edit, qwen_image_edit_plus) -- [ ] nucleus_image, krea2, ideogram4, mageflow -- [ ] chroma, chroma_radiance, zeta_chroma -- [ ] flux2, flux_kontext -- [ ] minimax_h3 (+ ref2va), ltx2 family -- [ ] wan21 / wan22 family -- [ ] hidream family, omnigen2 -- [ ] anima, boogu_image, ernie_image, f_light, prx_pixel_t2i -- [ ] audio_models (ace_step) +- [x] `text_encoders/qwen3.py` — Qwen3TextEncoder + `OstrisTransformersMixin` + backend + `BaseModel.prepare_text_encoder` policy helper; the 3 verbatim + TE stanzas (z_image, z_image_l2p, zeta_chroma) replaced. Verified with + real Z-Image weights (load + encode on GPU). +- [x] `text_encoders/qwen3_vl.py` — Qwen3VLTextEncoder with + `drop_vision_tower` / `patch_vision_patch_embed`; the 4 identical + `patch_qwen_vl_patch_embed` copies (krea2, mageflow, boogu_image, + Qwen3VLCaptioner) consolidated; TE loads migrated in krea2, mageflow, + nucleus_image. Still on their own paths: ideogram4 (loads via AutoModel), + minimax_h3 (custom truncated/prequantized comfy load — port later), + qwen_image (Qwen2.5-VL, needs its own class) +- [x] `text_encoders/t5.py`, `text_encoders/clip.py` — T5TextEncoder, + CLIPTextEncoder, CLIPTextEncoderWithProjection; migrated chroma ×2, + flux_kontext, f_light (T5 stanzas → `prepare_text_encoder`, fixing their + `qtype` → `qtype_te` bug) and hidream (CLIP ×2 + T5 with subfolder + overrides; slow-tokenizer classes preserved via `use_fast=False`) +- [x] `vae/qwen_image.py` — QwenImageVAE + QwenImageVAEHolderMixin (frame-dim + + latents mean/std handling built in, tiling opt-in via + `vae_decode_tiled_on_low_vram`); the triplicated encode/decode deleted + from qwen_image, nucleus_image, krea2 and all three VAE loads routed + through the v2 loader +- [x] `vae/autoencoder_kl.py` — KLVAE (diffusers AutoencoderKL through the + universal loader); migrated the scattered loads in chroma, flux_kontext, + f_light, hidream, z_image +- [x] `vae/flux2_kl.py` — the BFL-style Flux2 KL autoencoder unified from the + flux2 + ideogram4 copies (both files deleted; flux2's + encode/decode/small-decoder superset + ideogram4's diffusers key + converter). Verified bit-identical to both originals (weights, encode/ + decode outputs, converter mapping) and round-tripped real ae.safetensors + weights on GPU. Packing/normalization stays per-model — flux2 packs + `(c pi pj)` with BatchNorm running stats, ideogram4 packs `(ph pw c)` + with its latent_norm tables; the conventions are incompatible. +- [x] z_image — transformer, TE (qwen3), and VAE (KLVAE) all on v2 modules. + z_image_l2p still has its local progressive-transformer subclass + (rebasing it onto the v2 class deferred; its TE is migrated) +- [x] qwen_image family — `v2/diffusion_models/qwen_image.py` (single-file + loads stay on diffusers' from_single_file until the comfy flip) + + `v2/text_encoders/qwen25_vl.py` (slow tokenizer preserved); edit + variants inherit +- [x] nucleus_image — `v2/diffusion_models/nucleus_image.py`, TE stanza + collapsed to prepare_text_encoder +- [ ] krea2, ideogram4, mageflow — TE/VAE migrated; their custom local DiT + classes still to be rebased onto the mixin +- [x] chroma, chroma_radiance — both vendored Chroma classes now carry + `OstrisModelMixin` with the block-count sniff moved into a new + `aitk_config_from_state_dict` hook (mixin now supports checkpoint-derived + configs + `load_from_state_dict` for non-safetensors sources, used by + radiance's .pth path). zeta_chroma transformer left as-is: its config + depends on holder state (patch_size), not the checkpoint +- [x] flux_kontext — `v2/diffusion_models/flux.py` (FluxTransformer2DModel); + whole model now loads through v2 (transformer, T5, CLIP, KLVAE) +- [ ] flux2 — TE/VAE partially migrated (flux2_kl); custom DiT still local. + krea2/mageflow/ideogram4/zeta_chroma DiTs stay model-specific: their + configs come from model_kwargs / holder state, so the mixin adds nothing + until the comfy-weights flip (Phase 2) +- [ ] minimax_h3 (+ ref2va), ltx2 family — already on the shared resolver + + comfy_quant_import; the full mixin port waits for Phase 2, when the + mixin's single-file precision policy (stored-precision loading, fp32-key + protection) is settled to match their deliberate behavior +- [x] wan21 / wan22 family — `v2/diffusion_models/wan.py` + (WanTransformer3DModel, both wan22 dual loads included) + + `v2/text_encoders/umt5.py` (UMT5TextEncoder + PatchedT5Tokenizer; + `loaders/umt5.py` is now a thin compat shim, `comfy_files` still + reserved for Phase 2 — no local comfy umt5 file to verify the key + conversion against). wan21's TE `qtype` → `qtype_te` bug fixed via + prepare_text_encoder +- [x] hidream family — vendored transformer carries the mixin; + `v2/diffusion_models/hidream.py` wraps the diffusers class for + hidream_e1; both load via the switchable `hidream_transformer_class` + through `load_model` +- [x] omnigen2 — vendored transformer carries the mixin, load migrated +- [x] boogu_image, ernie_image, prx_pixel_t2i — their vendored diffusers-style + DiT classes now carry OstrisModelMixin (subfolder + block names on the + class) and the holders load via `load_model` +- [x] f_light — DiT class carries the mixin (`aitk_subfolder="dit_model"`), + load migrated +- [ ] anima — loads through diffusers modular pipelines (AnimaModularPipeline); + not a mixin fit, revisit at Phase 2 +- [ ] flux2 DiT — holder-config params classes (Flux2/Klein variants), defer + like krea2/mageflow +- [ ] ace_step — one bundled safetensors holds model+TE+VAE+tokenizer via its + own load_models; decomposing into v2 components is its own task - [ ] Per-model fixes folded in as each migrates: `qtype_te` bug, dequant-on-save (`dequantize_if_quantized` everywhere), raw-`quantize()` → `quantize_model()` diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index 4e82640..cacc49f 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -141,9 +141,13 @@ class OstrisModelMixin: 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) @@ -213,11 +217,34 @@ class OstrisModelMixin: config_path: Optional[str] = None, subfolder: Optional[str] = None, ): - config = cls._load_single_file_config(config_path, subfolder) - 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: @@ -290,23 +317,51 @@ class OstrisModelMixin: # ------------------------------------------------------------------ @classmethod - def load_tokenizer(cls, **kwargs): + 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 - if cls.aitk_tokenizer_repo is None: + 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( - cls.aitk_tokenizer_repo, subfolder=cls.aitk_tokenizer_subfolder, **kwargs + source, subfolder=subfolder or "", **kwargs ) @classmethod - def load_processor(cls, **kwargs): + def load_processor( + cls, + name_or_path: Optional[str] = None, + subfolder: Optional[str] = None, + **kwargs, + ): from transformers import AutoProcessor - if cls.aitk_processor_repo is None: + 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( - cls.aitk_processor_repo, subfolder=cls.aitk_processor_subfolder, **kwargs + source, subfolder=subfolder or "", **kwargs ) # ------------------------------------------------------------------ @@ -403,3 +458,26 @@ class OstrisModelMixin: 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) diff --git a/toolkit/models/v2/diffusion_models/flux.py b/toolkit/models/v2/diffusion_models/flux.py new file mode 100644 index 0000000..f2fb094 --- /dev/null +++ b/toolkit/models/v2/diffusion_models/flux.py @@ -0,0 +1,14 @@ +from diffusers import FluxTransformer2DModel as DiffusersFluxTransformer2DModel + +from .._mixin import OstrisModelMixin + + +class FluxTransformer2DModel(DiffusersFluxTransformer2DModel, OstrisModelMixin): + """Flux1-family DiT (flux, flux_kontext, chroma-adjacent finetunes in + diffusers layout).""" + + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["transformer_blocks", "single_transformer_blocks"] diff --git a/toolkit/models/v2/diffusion_models/hidream.py b/toolkit/models/v2/diffusion_models/hidream.py new file mode 100644 index 0000000..6ec8f3e --- /dev/null +++ b/toolkit/models/v2/diffusion_models/hidream.py @@ -0,0 +1,18 @@ +from diffusers.models import ( + HiDreamImageTransformer2DModel as DiffusersHiDreamImageTransformer2DModel, +) + +from .._mixin import OstrisModelMixin + + +class HiDreamImageTransformer2DModel( + DiffusersHiDreamImageTransformer2DModel, OstrisModelMixin +): + """The diffusers HiDream DiT (hidream_e1; the base hidream arch uses the + vendored copy in the hidream extension).""" + + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["double_stream_blocks", "single_stream_blocks"] diff --git a/toolkit/models/v2/diffusion_models/nucleus_image.py b/toolkit/models/v2/diffusion_models/nucleus_image.py new file mode 100644 index 0000000..62092ce --- /dev/null +++ b/toolkit/models/v2/diffusion_models/nucleus_image.py @@ -0,0 +1,15 @@ +from diffusers import ( + NucleusMoEImageTransformer2DModel as DiffusersNucleusMoEImageTransformer2DModel, +) + +from .._mixin import OstrisModelMixin + + +class NucleusMoEImageTransformer2DModel( + DiffusersNucleusMoEImageTransformer2DModel, OstrisModelMixin +): + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["transformer_blocks"] diff --git a/toolkit/models/v2/diffusion_models/qwen_image.py b/toolkit/models/v2/diffusion_models/qwen_image.py new file mode 100644 index 0000000..6d5b7d7 --- /dev/null +++ b/toolkit/models/v2/diffusion_models/qwen_image.py @@ -0,0 +1,29 @@ +from diffusers import ( + QwenImageTransformer2DModel as DiffusersQwenImageTransformer2DModel, +) + +from .._mixin import OstrisModelMixin + + +class QwenImageTransformer2DModel( + DiffusersQwenImageTransformer2DModel, OstrisModelMixin +): + aitk_subfolder = "transformer" + aitk_config_repo = "Qwen/Qwen-Image" + + @classmethod + def get_transformer_block_names(cls): + return ["transformer_blocks"] + + @classmethod + def _load_single_file(cls, file_path, dtype, config_path=None, subfolder=None): + # single-file checkpoints in the wild carry diffusers or original key + # layouts; diffusers' single-file machinery owns that conversion + model = cls.from_single_file( + file_path, + config=config_path if config_path is not None else cls.aitk_config_repo, + subfolder="transformer", + torch_dtype=dtype, + ) + model.to(dtype) + return model diff --git a/toolkit/models/v2/diffusion_models/wan.py b/toolkit/models/v2/diffusion_models/wan.py new file mode 100644 index 0000000..2ba9e88 --- /dev/null +++ b/toolkit/models/v2/diffusion_models/wan.py @@ -0,0 +1,13 @@ +from diffusers import WanTransformer3DModel as DiffusersWanTransformer3DModel + +from .._mixin import OstrisModelMixin + + +class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin): + """Wan 2.1/2.2 video DiT (wan22 loads two of these into its dual wrapper).""" + + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["blocks"] diff --git a/toolkit/models/v2/text_encoders/clip.py b/toolkit/models/v2/text_encoders/clip.py new file mode 100644 index 0000000..6aac3bb --- /dev/null +++ b/toolkit/models/v2/text_encoders/clip.py @@ -0,0 +1,27 @@ +from transformers import CLIPTextModel, CLIPTextModelWithProjection + +from .._mixin import OstrisTransformersMixin + + +class CLIPTextEncoder(CLIPTextModel, OstrisTransformersMixin): + """CLIP-L text encoder (flux-style checkpoints: text_encoder/ + + tokenizer/).""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + return ["text_model.encoder.layers"] + + +class CLIPTextEncoderWithProjection(CLIPTextModelWithProjection, OstrisTransformersMixin): + """CLIP text encoder with the projection head (SDXL / SD3 / HiDream style + checkpoints; the second encoder lives at text_encoder_2/ + tokenizer_2/).""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + return ["text_model.encoder.layers"] diff --git a/toolkit/models/v2/text_encoders/qwen25_vl.py b/toolkit/models/v2/text_encoders/qwen25_vl.py new file mode 100644 index 0000000..ba10fe3 --- /dev/null +++ b/toolkit/models/v2/text_encoders/qwen25_vl.py @@ -0,0 +1,22 @@ +from transformers import Qwen2_5_VLForConditionalGeneration + +from .._mixin import OstrisTransformersMixin + + +class Qwen25VLTextEncoder(Qwen2_5_VLForConditionalGeneration, OstrisTransformersMixin): + """Qwen2.5-VL conditioning stack (qwen_image family). Loads from a + checkpoint's text_encoder/ subfolder; the edit variants keep the vision + tower, plain t2i drops it.""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + return ["model.language_model.layers"] + + def drop_vision_tower(self): + """Text-only conditioning: the vision tower is dead weight.""" + if getattr(self.model, "visual", None) is not None: + self.model.visual = None + return self diff --git a/toolkit/models/v2/text_encoders/qwen3.py b/toolkit/models/v2/text_encoders/qwen3.py new file mode 100644 index 0000000..328712b --- /dev/null +++ b/toolkit/models/v2/text_encoders/qwen3.py @@ -0,0 +1,17 @@ +from transformers import Qwen3ForCausalLM + +from .._mixin import OstrisTransformersMixin + + +class Qwen3TextEncoder(Qwen3ForCausalLM, OstrisTransformersMixin): + """Qwen3 causal-LM text encoder (Z-Image family, Zeta-Chroma, ...). Loads + from a checkpoint's text_encoder/ subfolder, a hub repo, or a single + .safetensors file; the tokenizer rides in the checkpoint's tokenizer/ + subfolder.""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + return ["model.layers"] diff --git a/toolkit/models/v2/text_encoders/qwen3_vl.py b/toolkit/models/v2/text_encoders/qwen3_vl.py new file mode 100644 index 0000000..dd42e22 --- /dev/null +++ b/toolkit/models/v2/text_encoders/qwen3_vl.py @@ -0,0 +1,53 @@ +import torch +import torch.nn.functional as F +from transformers import Qwen3VLForConditionalGeneration + +from .._mixin import OstrisTransformersMixin + + +def patch_qwen_vl_patch_embed(model) -> int: + """Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain + linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and + falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear + (a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves. + Returns the number of patch_embed modules patched.""" + patched = 0 + for module in model.modules(): + proj = getattr(module, "proj", None) + if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple( + proj.stride + ): + + def fast_forward(hidden_states, _proj=proj): + w = _proj.weight.reshape(_proj.weight.shape[0], -1) + x = hidden_states.view(-1, w.shape[1]).to(w.dtype) + return F.linear(x, w, _proj.bias) + + module.forward = fast_forward + patched += 1 + return patched + + +class Qwen3VLTextEncoder(Qwen3VLForConditionalGeneration, OstrisTransformersMixin): + """Qwen3-VL conditioning stack (krea2, mageflow, nucleus_image, ideogram4, + minimax_h3, ...). Loads from a checkpoint's text_encoder/ subfolder or a + raw Qwen repo (pass subfolder=\"\" for the latter).""" + + aitk_subfolder = "text_encoder" + aitk_processor_subfolder = "processor" + + @classmethod + def get_transformer_block_names(cls): + return ["model.language_model.layers"] + + def drop_vision_tower(self): + """Text-only conditioning: the vision tower is dead weight — drop it to + free VRAM and skip loading its (bf16-slow) Conv3d patch_embed.""" + if getattr(self.model, "visual", None) is not None: + self.model.visual = None + return self + + def patch_vision_patch_embed(self) -> int: + """Keep the vision tower (reference images ride into the embeddings) + but swap its Conv3d patch_embed for an equivalent GEMM.""" + return patch_qwen_vl_patch_embed(self) diff --git a/toolkit/models/v2/text_encoders/t5.py b/toolkit/models/v2/text_encoders/t5.py new file mode 100644 index 0000000..1a5953e --- /dev/null +++ b/toolkit/models/v2/text_encoders/t5.py @@ -0,0 +1,16 @@ +from transformers import T5EncoderModel + +from .._mixin import OstrisTransformersMixin + + +class T5TextEncoder(T5EncoderModel, OstrisTransformersMixin): + """T5-XXL text encoder. Defaults to the flux-style checkpoint layout + (text_encoder_2/ + tokenizer_2/); pass subfolder overrides for checkpoints + that keep it at text_encoder/ + tokenizer/ (e.g. f-lite).""" + + aitk_subfolder = "text_encoder_2" + aitk_tokenizer_subfolder = "tokenizer_2" + + @classmethod + def get_transformer_block_names(cls): + return ["encoder.block"] diff --git a/toolkit/models/v2/text_encoders/umt5.py b/toolkit/models/v2/text_encoders/umt5.py new file mode 100644 index 0000000..d0b2e99 --- /dev/null +++ b/toolkit/models/v2/text_encoders/umt5.py @@ -0,0 +1,49 @@ +import torch +from transformers import T5Tokenizer, UMT5EncoderModel + +from .._mixin import OstrisTransformersMixin + + +class PatchedT5Tokenizer(T5Tokenizer): + def __init__( + self, + vocab=None, + eos_token="", + unk_token="", + pad_token="", + _spm_precompiled_charsmap=None, + extra_ids=100, + additional_special_tokens=None, + **kwargs, + ): + super().__init__( + vocab=vocab, + eos_token=eos_token, + unk_token=unk_token, + pad_token=pad_token, + _spm_precompiled_charsmap=None, # this is passing a empty byte string for some reason now. + extra_ids=extra_ids, + additional_special_tokens=additional_special_tokens, + **kwargs, + ) + + +class UMT5TextEncoder(UMT5EncoderModel, OstrisTransformersMixin): + """UMT5-XXL text encoder (wan family).""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + return ["encoder.block"] + + @classmethod + def load_tokenizer(cls, name_or_path=None, subfolder=None, **kwargs): + # T5's tokenizer needs the _spm_precompiled_charsmap patch + source = name_or_path if name_or_path is not None else cls.aitk_tokenizer_repo + if subfolder is None: + subfolder = cls.aitk_tokenizer_subfolder + return PatchedT5Tokenizer.from_pretrained( + source, subfolder=subfolder or "", **kwargs + ) diff --git a/toolkit/models/v2/vae/autoencoder_kl.py b/toolkit/models/v2/vae/autoencoder_kl.py new file mode 100644 index 0000000..0eb58fc --- /dev/null +++ b/toolkit/models/v2/vae/autoencoder_kl.py @@ -0,0 +1,10 @@ +from diffusers import AutoencoderKL + +from .._mixin import OstrisModelMixin + + +class KLVAE(AutoencoderKL, OstrisModelMixin): + """The diffusers AutoencoderKL (SD/SDXL/Flux1/Z-Image image VAEs), loaded + from a checkpoint's vae/ subfolder through the universal loader.""" + + aitk_subfolder = "vae" diff --git a/extensions_built_in/diffusion_models/ideogram4/src/vae.py b/toolkit/models/v2/vae/flux2_kl.py similarity index 83% rename from extensions_built_in/diffusion_models/ideogram4/src/vae.py rename to toolkit/models/v2/vae/flux2_kl.py index d03a2e2..e7f6b00 100644 --- a/extensions_built_in/diffusion_models/ideogram4/src/vae.py +++ b/toolkit/models/v2/vae/flux2_kl.py @@ -1,4 +1,13 @@ -"""Flux2 KL autoencoder.""" +"""The BFL-style Flux2 KL autoencoder (32ch latents, 2x2 pixel-shuffle +packing to 128ch), shared by the flux2 family and ideogram4. + +flux2 loads the raw BFL ae.safetensors layout and uses encode/decode (with the +BatchNorm running-stats latent normalization); ideogram4 loads diffusers-format +checkpoints via convert_diffusers_state_dict and drives encoder/decoder +directly with its own patchify + latent-norm tables. The two models pack the +128 latent channels in different orders — the packing/normalization stays +per-model. +""" from __future__ import annotations @@ -22,6 +31,17 @@ class AutoEncoderParams: num_res_blocks: int = 2 z_channels: int = 32 +@dataclass +class AutoEncoderSmallDecoderParams: + resolution: int = 256 + in_channels: int = 3 + ch: int = 128 + ch_encoder: int = 96 + out_ch: int = 3 + ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4]) + num_res_blocks: int = 2 + z_channels: int = 32 + def swish(x: Tensor) -> Tensor: return x * torch.sigmoid(x) @@ -194,30 +214,30 @@ class Encoder(nn.Module): for i_level in range(self.num_resolutions): for i_block in range(self.num_res_blocks): if torch.is_grad_enabled() and self.gradient_checkpointing: - h = ckpt.checkpoint(self.down[i_level].block[i_block], hs[-1]) # type: ignore[index, operator] - if len(self.down[i_level].attn) > 0: # type: ignore[arg-type] - h = ckpt.checkpoint(self.down[i_level].attn[i_block], h) # type: ignore[index, operator] + h = ckpt.checkpoint(self.down[i_level].block[i_block], hs[-1]) + if len(self.down[i_level].attn) > 0: + h = ckpt.checkpoint(self.down[i_level].attn[i_block], h) else: - h = self.down[i_level].block[i_block](hs[-1]) # type: ignore[index, operator] - if len(self.down[i_level].attn) > 0: # type: ignore[arg-type] - h = self.down[i_level].attn[i_block](h) # type: ignore[index, operator] + h = self.down[i_level].block[i_block](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) hs.append(h) if i_level != self.num_resolutions - 1: if torch.is_grad_enabled() and self.gradient_checkpointing: - hs.append(ckpt.checkpoint(self.down[i_level].downsample, hs[-1])) # type: ignore[operator] + hs.append(ckpt.checkpoint(self.down[i_level].downsample, hs[-1])) else: - hs.append(self.down[i_level].downsample(hs[-1])) # type: ignore[operator] + hs.append(self.down[i_level].downsample(hs[-1])) # middle h = hs[-1] if torch.is_grad_enabled() and self.gradient_checkpointing: - h = ckpt.checkpoint(self.mid.block_1, h) # type: ignore[operator] - h = ckpt.checkpoint(self.mid.attn_1, h) # type: ignore[operator] - h = ckpt.checkpoint(self.mid.block_2, h) # type: ignore[operator] + h = ckpt.checkpoint(self.mid.block_1, h) + h = ckpt.checkpoint(self.mid.attn_1, h) + h = ckpt.checkpoint(self.mid.block_2, h) else: - h = self.mid.block_1(h) # type: ignore[operator] - h = self.mid.attn_1(h) # type: ignore[operator] - h = self.mid.block_2(h) # type: ignore[operator] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) # end h = self.norm_out(h) h = swish(h) @@ -300,13 +320,13 @@ class Decoder(nn.Module): # middle if torch.is_grad_enabled() and self.gradient_checkpointing: - h = ckpt.checkpoint(self.mid.block_1, h) # type: ignore[operator] - h = ckpt.checkpoint(self.mid.attn_1, h) # type: ignore[operator] - h = ckpt.checkpoint(self.mid.block_2, h) # type: ignore[operator] + h = ckpt.checkpoint(self.mid.block_1, h) + h = ckpt.checkpoint(self.mid.attn_1, h) + h = ckpt.checkpoint(self.mid.block_2, h) else: - h = self.mid.block_1(h) # type: ignore[operator] - h = self.mid.attn_1(h) # type: ignore[operator] - h = self.mid.block_2(h) # type: ignore[operator] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) # cast to proper dtype h = h.to(upscale_dtype) @@ -314,18 +334,18 @@ class Decoder(nn.Module): for i_level in reversed(range(self.num_resolutions)): for i_block in range(self.num_res_blocks + 1): if torch.is_grad_enabled() and self.gradient_checkpointing: - h = ckpt.checkpoint(self.up[i_level].block[i_block], h) # type: ignore[index, operator] - if len(self.up[i_level].attn) > 0: # type: ignore[arg-type] - h = ckpt.checkpoint(self.up[i_level].attn[i_block], h) # type: ignore[index, operator] + h = ckpt.checkpoint(self.up[i_level].block[i_block], h) + if len(self.up[i_level].attn) > 0: + h = ckpt.checkpoint(self.up[i_level].attn[i_block], h) else: - h = self.up[i_level].block[i_block](h) # type: ignore[index, operator] - if len(self.up[i_level].attn) > 0: # type: ignore[arg-type] - h = self.up[i_level].attn[i_block](h) # type: ignore[index, operator] + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) if i_level != 0: if torch.is_grad_enabled() and self.gradient_checkpointing: - h = ckpt.checkpoint(self.up[i_level].upsample, h) # type: ignore[operator] + h = ckpt.checkpoint(self.up[i_level].upsample, h) else: - h = self.up[i_level].upsample(h) # type: ignore[operator] + h = self.up[i_level].upsample(h) # end h = self.norm_out(h) @@ -346,10 +366,13 @@ class AutoEncoder(nn.Module): num_res_blocks=params.num_res_blocks, z_channels=params.z_channels, ) + decoder_ch = params.ch + if hasattr(params, "ch_encoder"): + decoder_ch = params.ch_encoder self.decoder = Decoder( resolution=params.resolution, in_channels=params.in_channels, - ch=params.ch, + ch=decoder_ch, out_ch=params.out_ch, ch_mult=params.ch_mult, num_res_blocks=params.num_res_blocks, @@ -369,7 +392,7 @@ class AutoEncoder(nn.Module): self._gradient_checkpointing = False @property - def gradient_checkpointing(self) -> bool: + def gradient_checkpointing(self): return self._gradient_checkpointing @gradient_checkpointing.setter @@ -378,18 +401,52 @@ class AutoEncoder(nn.Module): self.encoder.gradient_checkpointing = value self.decoder.gradient_checkpointing = value + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + def enable_gradient_checkpointing(self): self.gradient_checkpointing = True self.encoder.enable_gradient_checkpointing() self.decoder.enable_gradient_checkpointing() - @property - def device(self) -> torch.device: - return next(self.parameters()).device + def normalize(self, z): + self.bn.eval() + return self.bn(z) - @property - def dtype(self) -> torch.dtype: - return next(self.parameters()).dtype + def inv_normalize(self, z): + self.bn.eval() + s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.bn_eps) + m = self.bn.running_mean.view(1, -1, 1, 1) + return z * s + m + + def encode(self, x: Tensor) -> Tensor: + moments = self.encoder(x) + mean = torch.chunk(moments, 2, dim=1)[0] + + z = rearrange( + mean, + "... c (i pi) (j pj) -> ... (c pi pj) i j", + pi=self.ps[0], + pj=self.ps[1], + ) + z = self.normalize(z) + return z + + def decode(self, z: Tensor) -> Tensor: + z = self.inv_normalize(z) + z = rearrange( + z, + "... (c pi pj) i j -> ... c (i pi) (j pj)", + pi=self.ps[0], + pj=self.ps[1], + ) + dec = self.decoder(z) + return dec _NUM_RESOLUTIONS = 4 diff --git a/toolkit/models/v2/vae/qwen_image.py b/toolkit/models/v2/vae/qwen_image.py new file mode 100644 index 0000000..af6bc59 --- /dev/null +++ b/toolkit/models/v2/vae/qwen_image.py @@ -0,0 +1,86 @@ +import torch +from diffusers import AutoencoderKLQwenImage + +from .._mixin import OstrisModelMixin + + +class QwenImageVAE(AutoencoderKLQwenImage, OstrisModelMixin): + """The Qwen-Image (wan-style video) VAE, shared by the qwen_image family, + nucleus_image and krea2.""" + + aitk_subfolder = "vae" + + +class QwenImageVAEHolderMixin: + """BaseModel-side encode_images/decode_latents for models whose self.vae + is the Qwen-Image VAE: it is a video VAE, so images ride in a single-frame + dim and latents are normalized with the config's latents_mean/std.""" + + # tile the decode when low_vram (decode only; encode stays untiled) + vae_decode_tiled_on_low_vram = False + + def encode_images(self, image_list, device=None, dtype=None): + if device is None: + device = self.vae_device_torch + if dtype is None: + dtype = self.vae_torch_dtype + + if self.vae.device == torch.device("cpu"): + self.vae.to(device) + self.vae.eval() + self.vae.requires_grad_(False) + + image_list = [image.to(device, dtype=dtype) for image in image_list] + images = torch.stack(image_list).to(device, dtype=dtype) + images = images.unsqueeze(2) # add the frame dim + latents = self.vae.encode(images).latent_dist.sample() + + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view( + 1, self.vae.config.z_dim, 1, 1, 1 + ).to(latents.device, latents.dtype) + + latents = (latents - latents_mean) * latents_std + latents = latents.squeeze(2) # drop the frame dim + return latents.to(device, dtype=dtype) + + def decode_latents(self, latents: torch.Tensor, device=None, dtype=None): + if device is None: + device = self.vae_device_torch + if dtype is None: + dtype = self.vae_torch_dtype + + if self.vae.device == torch.device("cpu"): + self.vae.to(device) + + latents = latents.to(device, dtype=dtype) + latents = latents.unsqueeze(2) # add the frame dim + + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents = latents * latents_std + latents_mean + + # full-resolution decode spikes VRAM; models opt in to tiling it + tiled = self.vae_decode_tiled_on_low_vram and self.model_config.low_vram + if tiled: + self.vae.enable_tiling() + try: + images = self.vae.decode(latents).sample + finally: + if tiled: + self.vae.disable_tiling() + + images = images.squeeze(2) # drop the frame dim + return images.to(device, dtype=dtype) diff --git a/toolkit/models/wan21/wan21.py b/toolkit/models/wan21/wan21.py index 3f7f953..1fec78d 100644 --- a/toolkit/models/wan21/wan21.py +++ b/toolkit/models/wan21/wan21.py @@ -10,7 +10,8 @@ from toolkit.memory_management.manager import MemoryManager from toolkit.models.base_model import BaseModel from toolkit.prompt_utils import PromptEmbeds from transformers import AutoTokenizer, UMT5EncoderModel -from diffusers import WanPipeline, WanTransformer3DModel, AutoencoderKL +from diffusers import WanPipeline, AutoencoderKL +from toolkit.models.v2.diffusion_models.wan import WanTransformer3DModel from .autoencoder_kl_wan import AutoencoderKLWan import os import sys @@ -29,8 +30,6 @@ import os import copy from toolkit.config_modules import ModelConfig, GenerateImageConfig, ModelArch import torch -from optimum.quanto import freeze, qfloat8, QTensor, qint4 -from toolkit.util.quantize import quantize, get_qtype from diffusers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler from typing import TYPE_CHECKING, List from toolkit.accelerator import unwrap_model @@ -44,7 +43,7 @@ from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback from typing import Any, Callable, Dict, List, Optional, Union from toolkit.models.wan21.wan_lora_convert import convert_to_diffusers, convert_to_original from toolkit.util.quantize import quantize_model -from toolkit.models.loaders.umt5 import get_umt5_encoder +from toolkit.models.v2.text_encoders.umt5 import UMT5TextEncoder # for generation only? scheduler_configUniPC = { @@ -344,11 +343,9 @@ class Wan21(BaseModel): def load_wan_transformer(self, transformer_path, subfolder=None): self.print_and_status_update("Loading transformer") dtype = self.torch_dtype - transformer = WanTransformer3DModel.from_pretrained( - transformer_path, - subfolder=subfolder, - torch_dtype=dtype, - ).to(dtype=dtype) + transformer = WanTransformer3DModel.load_model( + transformer_path, dtype=dtype, subfolder=subfolder + ) if self.model_config.split_model_over_gpus: raise ValueError( @@ -418,29 +415,9 @@ class Wan21(BaseModel): self.print_and_status_update("Loading UMT5EncoderModel") - tokenizer, text_encoder = get_umt5_encoder( - model_path=te_path, - tokenizer_subfolder="tokenizer", - encoder_subfolder="text_encoder", - torch_dtype=dtype, - comfy_files=self._comfy_te_file - ) - - text_encoder.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing UMT5EncoderModel") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype)) - freeze(text_encoder) - flush() - - if self.model_config.layer_offloading and self.model_config.layer_offloading_text_encoder_percent > 0: - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent - ) + tokenizer = UMT5TextEncoder.load_tokenizer(te_path) + text_encoder = UMT5TextEncoder.load_model(te_path, dtype=dtype) + self.prepare_text_encoder(text_encoder, dtype=dtype) if self.model_config.low_vram: print("Moving transformer back to GPU") From 9ed2e0b8e7ff018affb8783753ebee7e64b62314 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 27 Aug 2026 14:31:53 -0600 Subject: [PATCH 03/14] Test model loading --- .gitignore | 3 +- .../diffusion_models/chroma/chroma_model.py | 10 +- .../chroma/chroma_radiance_model.py | 10 +- testing/test_model_loading.py | 286 ++++++++++++++++++ toolkit/models/v2/PLANNING.md | 33 +- 5 files changed, 326 insertions(+), 16 deletions(-) create mode 100644 testing/test_model_loading.py diff --git a/.gitignore b/.gitignore index cfc9185..d54f259 100644 --- a/.gitignore +++ b/.gitignore @@ -191,4 +191,5 @@ aitk_db.db-shm /data .claude original_repo -.next \ No newline at end of file +.next +testing/.model_test_outputs \ No newline at end of file diff --git a/extensions_built_in/diffusion_models/chroma/chroma_model.py b/extensions_built_in/diffusion_models/chroma/chroma_model.py index 0cd0166..d9d2025 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_model.py @@ -50,10 +50,12 @@ class FakeConfig: self.patch_size = 1 class FakeCLIP(torch.nn.Module): - def __init__(self): + def __init__(self, device='cuda'): super().__init__() self.dtype = torch.bfloat16 - self.device = 'cuda' + # the pipeline derives its execution device from this attribute; + # nn.Module.to() does not update it + self.device = device self.text_model = None self.tokenizer = None self.model_max_length = 77 @@ -180,8 +182,8 @@ class ChromaModel(BaseModel): self.prepare_text_encoder(text_encoder_2, dtype=dtype) # self.print_and_status_update("Loading CLIP") - text_encoder = FakeCLIP() - tokenizer = FakeCLIP() + text_encoder = FakeCLIP(device=self.device_torch) + tokenizer = FakeCLIP(device=self.device_torch) text_encoder.to(self.device_torch, dtype=dtype) self.noise_scheduler = ChromaModel.get_train_scheduler() diff --git a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py index 0a6133b..d95dddc 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py @@ -50,10 +50,12 @@ class FakeConfig: self.patch_size = 1 class FakeCLIP(torch.nn.Module): - def __init__(self): + def __init__(self, device='cuda'): super().__init__() self.dtype = torch.bfloat16 - self.device = 'cuda' + # the pipeline derives its execution device from this attribute; + # nn.Module.to() does not update it + self.device = device self.text_model = None self.tokenizer = None self.model_max_length = 77 @@ -182,8 +184,8 @@ class ChromaRadianceModel(BaseModel): self.prepare_text_encoder(text_encoder_2, dtype=dtype) # self.print_and_status_update("Loading CLIP") - text_encoder = FakeCLIP() - tokenizer = FakeCLIP() + text_encoder = FakeCLIP(device=self.device_torch) + tokenizer = FakeCLIP(device=self.device_torch) text_encoder.to(self.device_torch, dtype=dtype) self.noise_scheduler = ChromaRadianceModel.get_train_scheduler() diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py new file mode 100644 index 0000000..9ad0882 --- /dev/null +++ b/testing/test_model_loading.py @@ -0,0 +1,286 @@ +"""Per-arch model loading + inference smoke test (toolkit/models/v2/PLANNING.md). + +Loads one registered arch through its normal loading path (the same +get_model_class -> ModelClass(...).load_model() flow training uses), runs one +small sample generation, and asserts an output file was produced. + +Usage: + python testing/test_model_loading.py --arch zimage # one arch, in-process + python testing/test_model_loading.py --all # every registered arch, + # one subprocess each (full + # unload between archs) + --allow-download permit hub downloads (default: HF_HUB_OFFLINE=1, so archs + whose weights are not local/cached report SKIP) + --list list registered archs + --device cuda:0 + +Add a new model type by adding an entry to MODEL_TESTS. +""" + +import argparse +import glob +import json +import os +import subprocess +import sys + +TOOLKIT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, TOOLKIT_ROOT) + +from dotenv import load_dotenv + +# repo .env carries HF_TOKEN / HF_HOME / MODELS_PATH etc., same as run.py +load_dotenv(os.path.join(TOOLKIT_ROOT, ".env")) + +OUTPUT_ROOT = os.path.join(TOOLKIT_ROOT, "testing", ".model_test_outputs") + +# arch -> {"model": ModelConfig kwargs, "sample": GenerateImageConfig kwargs} +# Keep samples tiny: this asserts the load/encode/denoise/decode/save path +# works, not quality. +IMG = {"width": 512, "height": 512, "num_inference_steps": 8, "seed": 42} +VID = {"width": 256, "height": 256, "num_inference_steps": 6, "seed": 42, "num_frames": 9} + +MODEL_TESTS = { + "zimage": { + "model": {"name_or_path": "Tongyi-MAI/Z-Image-Turbo"}, + "sample": {**IMG, "guidance_scale": 1.0}, + }, + "qwen_image": { + # 20B: quantize to fit a 32GB card + "model": {"name_or_path": "Qwen/Qwen-Image", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 4.0}, + }, + "krea2": { + "model": {"name_or_path": "krea/Krea-2-Turbo", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "guidance_scale": 1.0}, + }, + "boogu_image": { + # native ~1024; 512/low-step/high-CFG degenerates to a black frame + "model": {"name_or_path": "Boogu/Boogu-Image-0.1-Base", "quantize": True, "quantize_te": True}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + }, + "ernie_image": { + "model": {"name_or_path": "baidu/ERNIE-Image", "quantize": True, "quantize_te": True}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + }, + "mageflow": { + "model": {"name_or_path": "microsoft/Mage-Flow-Base", "quantize": True, "quantize_te": True}, + "sample": IMG, + }, + "ideogram4": { + "model": {"name_or_path": "ideogram-ai/ideogram-4-fp8", "quantize": True, "quantize_te": True}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + }, + "hidream_o1": { + "model": {"name_or_path": "HiDream-ai/HiDream-O1-Image", "quantize": True, "quantize_te": True}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 28, "guidance_scale": 5.0, "seed": 42}, + }, + "anima": { + "model": {"name_or_path": "circlestone-labs/Anima-Base-v1.0-Diffusers"}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.5, "seed": 42}, + }, + "wan21": { + "model": {"name_or_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"}, + "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 5.0, "seed": 42, "num_frames": 17}, + }, + "wan22_5b": { + "model": {"name_or_path": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "quantize": True, "quantize_te": True}, + "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 5.0, "seed": 42, "num_frames": 17}, + }, + "ltx2.3": { + # even quantized, the 22B stack does not fit a 32GB card — needs the big GPU + "model": {"name_or_path": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", "quantize": True, "quantize_te": True}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 3.0, "seed": 42, "num_frames": 25}, + }, + # single-file / comfy-layout archs: weights resolve under MODELS_PATH (or + # download there with --allow-download) + "chroma": { + "model": {"name_or_path": "lodestones/Chroma1-HD", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 26, "guidance_scale": 4.0}, + }, + "flux_kontext": { + "model": {"name_or_path": "black-forest-labs/FLUX.1-Kontext-dev", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 2.5}, + "needs_control_image": True, + }, + "flux2_klein_4b": { + "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-4B", "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0}, + }, +} + +SKIP_MARKERS = ( + "couldn't connect", + "offline mode", + "hf_hub_offline", + "cannot find the requested files", + "not found in cache", + "localentrynotfounderror", + "does not appear to have a file named", + "404 client error", + "entrynotfounderror", + "gatedrepoerror", + "cannot access gated repo", + "repositorynotfounderror", +) + + +def classify_error(err: BaseException) -> str: + text = f"{type(err).__name__}: {err}".lower() + if any(m in text for m in SKIP_MARKERS): + return "SKIP" + if isinstance(err, FileNotFoundError): + return "SKIP" + return "FAIL" + + +def run_one(arch: str, device: str, allow_download: bool) -> dict: + entry = MODEL_TESTS[arch] + out_dir = os.path.join(OUTPUT_ROOT, arch.replace("/", "_").replace(":", "_")) + os.makedirs(out_dir, exist_ok=True) + for old in glob.glob(os.path.join(out_dir, "*")): + os.remove(old) + + from toolkit.config_modules import GenerateImageConfig, ModelConfig + from toolkit.util.get_model import get_model_class + + model_config = ModelConfig(arch=arch, dtype="bf16", **entry["model"]) + ModelClass = get_model_class(model_config) + # get_model_class silently falls back to the legacy SD class on an + # unknown arch; that is never what a registered test wants + if getattr(ModelClass, "arch", None) not in (arch, model_config.arch): + raise ValueError( + f"arch {arch!r} resolved to {ModelClass.__name__} " + f"(arch={getattr(ModelClass, 'arch', None)!r}) — registry mismatch" + ) + + sampler = None + if hasattr(ModelClass, "get_train_scheduler"): + sampler = ModelClass.get_train_scheduler() + + sd = ModelClass( + device=device, + model_config=model_config, + dtype="bf16", + noise_scheduler=sampler, + ) + sd.load_model() + + sample_kwargs = dict(entry["sample"]) + if entry.get("needs_control_image"): + # edit/kontext models require a control image; a flat gray input is fine + from PIL import Image + + ctrl_path = os.path.join(out_dir, ".ctrl.png") + Image.new( + "RGB", (sample_kwargs["width"], sample_kwargs["height"]), (128, 128, 128) + ).save(ctrl_path) + sample_kwargs["ctrl_img"] = ctrl_path + gen = GenerateImageConfig( + prompt="a photo of a cat sitting on a wooden table", + output_folder=out_dir, + # the GenerateImageConfig default for output_ext is the Literal type + # alias itself; real callers always pass one + output_ext="png", + **sample_kwargs, + ) + sd.generate_images([gen]) + + produced = [ + p + for p in glob.glob(os.path.join(out_dir, "*")) + if os.path.isfile(p) and os.path.getsize(p) > 1024 and not p.endswith(".txt") + ] + if not produced: + raise RuntimeError(f"no output file produced in {out_dir}") + return {"arch": arch, "status": "PASS", "outputs": produced} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--arch", type=str, default=None) + parser.add_argument("--all", action="store_true") + parser.add_argument("--list", action="store_true") + parser.add_argument("--device", type=str, default="cuda:0") + parser.add_argument("--allow-download", action="store_true") + parser.add_argument("--json-result", type=str, default=None) + args = parser.parse_args() + + if args.list: + for arch in MODEL_TESTS: + print(arch) + return + + os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") + if not args.allow_download: + os.environ.setdefault("HF_HUB_OFFLINE", "1") + + if args.arch is not None: + if args.arch not in MODEL_TESTS: + raise SystemExit( + f"arch {args.arch!r} is not registered; --list shows options" + ) + try: + result = run_one(args.arch, args.device, args.allow_download) + except BaseException as err: + status = classify_error(err) + result = {"arch": args.arch, "status": status, "error": f"{type(err).__name__}: {err}"} + if status == "FAIL": + import traceback + + traceback.print_exc() + if args.json_result: + with open(args.json_result, "w") as f: + json.dump(result, f) + print(f"[{result['status']}] {args.arch}" + (f" — {result.get('error', '')}" if result["status"] != "PASS" else "")) + if result["status"] == "FAIL": + sys.exit(1) + return + + if not args.all: + parser.print_help() + return + + # --all: one subprocess per arch so every model fully unloads (clean CUDA + # teardown) before the next loads + results = [] + for arch in MODEL_TESTS: + print(f"\n===== {arch} =====") + result_path = os.path.join(OUTPUT_ROOT, f".{arch.replace('/', '_')}.result.json") + cmd = [ + sys.executable, + os.path.abspath(__file__), + "--arch", + arch, + "--device", + args.device, + "--json-result", + result_path, + ] + if args.allow_download: + cmd.append("--allow-download") + proc = subprocess.run(cmd, cwd=TOOLKIT_ROOT) + if os.path.exists(result_path): + with open(result_path) as f: + results.append(json.load(f)) + os.remove(result_path) + else: + results.append( + {"arch": arch, "status": "FAIL", "error": f"subprocess died (exit {proc.returncode})"} + ) + + print("\n===== summary =====") + counts = {"PASS": 0, "FAIL": 0, "SKIP": 0} + for r in results: + counts[r["status"]] = counts.get(r["status"], 0) + 1 + line = f"[{r['status']}] {r['arch']}" + if r["status"] != "PASS": + line += f" — {r.get('error', '')[:160]}" + print(line) + print(f"\n{counts['PASS']} passed, {counts['FAIL']} failed, {counts['SKIP']} skipped") + if counts["FAIL"]: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index 87e0c9b..b5d7fd1 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -267,14 +267,33 @@ loads via diffusers. Nothing about sources or outputs changes yet. Suggested ord ## Testing -- [ ] `testing/` (or `tests/`) harness: for each migrated arch, load the model via - its v2 modules and run one small inference pass (single low-step sample; video - models at minimum frame count). One arch at a time, full unload between archs. -- [ ] Weights resolved through the normal resolver against `MODELS_PATH` - (GPU + local-weights test, not CI-portable at first; skip archs whose weights - are absent rather than failing). +- [x] `testing/test_model_loading.py`: per-arch load + one small sample through + the normal training-style flow (get_model_class → load_model → + generate_images). `--arch X` runs one in-process; `--all` runs every + registered arch in its own subprocess (full unload between archs). + 15 archs registered so far — add each model type as it migrates. +- [x] Missing weights skip rather than fail: default is HF_HUB_OFFLINE=1 and + hub/file errors classify as SKIP; `--allow-download` opts into fetching. + (GPU + local-weights test, not CI-portable.) +- [x] Full sweep run 2026-08-27: 14/15 PASS (zimage, qwen_image, krea2, + boogu_image, ernie_image, ideogram4, hidream_o1, anima, wan21, wan22_5b, + chroma, flux_kontext, flux2_klein_4b, ltx2.3 — the quantized 22B ltx + stack doesn't fit 32GB, needs the 96GB card). mageflow blocked + upstream: microsoft/Mage-Flow-Base 404s on the hub (cached locally, so + it runs offline — recheck whether the repo moved/went private). +- [x] Registry carries realistic per-arch sample settings (native res, steps, + CFG) so sweep outputs are visually verifiable, not just "a file + exists". Verified: all 14 produce proper generations. Findings from + the quality pass: boogu emits a black frame below native res at + low-step/high-CFG (settings regime, present pre-restructure, not a + migration bug); chroma's FakeCLIP hardcoded device 'cuda' broke any + non-cuda:0 run (pre-existing, fixed — FakeCLIP now takes the real + device); ideogram4's fp8 release renders its own "blocked by safety + filter" card for a plain cat prompt (model behavior, not a bug — + investigate its trigger). - [ ] Round-trip test per model: load → save comfy format → reload from the save → - outputs match (bf16) / load cleanly (quantized saves). + outputs match (bf16) / load cleanly (quantized saves). Lands with the + Phase 2 comfy save path. - [ ] Each newly migrated model adds its test in the same PR as its migration. ## TODO / look at later From 9113420b615807ec873c900a308cb211988f2507 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 27 Aug 2026 15:37:13 -0600 Subject: [PATCH 04/14] Phase 2 --- .../diffusion_models/qwen_image/qwen_image.py | 29 ++- .../diffusion_models/z_image/z_image.py | 45 ++--- toolkit/models/v2/PLANNING.md | 94 ++++++++- toolkit/models/v2/_mixin.py | 123 ++++++++++-- .../models/v2/diffusion_models/qwen_image.py | 25 ++- toolkit/models/v2/diffusion_models/wan.py | 73 +++++++ toolkit/models/v2/diffusion_models/z_image.py | 54 +++++- toolkit/models/v2/resolver.py | 72 +++++++ toolkit/util/comfy_quant_export.py | 104 ++++++++++ toolkit/util/comfy_quant_import.py | 181 +++++++++++++++++- toolkit/util/float8_quant.py | 53 +++++ toolkit/util/nvfp4_quant.py | 19 ++ toolkit/util/ostris_quant.py | 3 + 13 files changed, 799 insertions(+), 76 deletions(-) create mode 100644 toolkit/util/comfy_quant_export.py create mode 100644 toolkit/util/float8_quant.py diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py index 54116a5..f1ad25e 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py @@ -17,6 +17,7 @@ from toolkit.accelerator import get_accelerator, unwrap_model from toolkit.util.quantize import quantize_model import torch.nn.functional as F from toolkit.memory_management import MemoryManager +from toolkit.metadata import get_meta_for_safetensors from safetensors.torch import load_file from diffusers import ( @@ -101,9 +102,17 @@ class QwenImageModel(QwenImageVAEHolderMixin, BaseModel): if os.path.exists(te_folder_path): base_model_path = model_path - transformer = QwenImageTransformer2DModel.load_model(model_path, dtype=model_dtype) + transformer = QwenImageTransformer2DModel.load_model( + model_path, + dtype=model_dtype, + use_comfy_weights=self.model_config.model_kwargs.get( + "use_comfy_weights", True + ), + ) - if self.model_config.quantize: + if self.model_config.quantize and not getattr( + transformer, "aitk_is_quantized", False + ): self.print_and_status_update("Quantizing Transformer") quantize_model(self, transformer) flush() @@ -350,17 +359,17 @@ class QwenImageModel(QwenImageVAEHolderMixin, BaseModel): return False def save_model(self, output_path, meta, save_dtype): - # only save the unet + # comfy-format single-file save (diffusers keys ARE the comfy layout + # for qwen image); prequantized layers keep their quantized storage transformer: QwenImageTransformer2DModel = unwrap_model(self.model) - transformer.save_pretrained( - save_directory=os.path.join(output_path, "transformer"), - safe_serialization=True, + if not output_path.endswith(".safetensors"): + output_path += ".safetensors" + transformer.save_model( + output_path, + dtype=save_dtype, + metadata=get_meta_for_safetensors(meta, name=self.arch), ) - meta_path = os.path.join(output_path, "aitk_meta.yaml") - with open(meta_path, "w") as f: - yaml.dump(meta, f) - def get_loss_target(self, *args, **kwargs): noise = kwargs.get("noise") batch = kwargs.get("batch") 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 67ffd90..bc2a30e 100644 --- a/extensions_built_in/diffusion_models/z_image/z_image.py +++ b/extensions_built_in/diffusion_models/z_image/z_image.py @@ -3,7 +3,6 @@ from typing import List, Optional import huggingface_hub import torch -import yaml from toolkit.config_modules import GenerateImageConfig, ModelConfig, NetworkConfig from toolkit.lora_special import LoRASpecialNetwork from toolkit.models.base_model import BaseModel @@ -13,15 +12,12 @@ from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.accelerator import unwrap_model -from toolkit.util.quantize import ( - quantize_model, - dequantize_if_quantized, -) +from toolkit.util.quantize import quantize_model from toolkit.memory_management import MemoryManager from toolkit.metadata import get_meta_for_safetensors from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder from toolkit.models.v2.vae.autoencoder_kl import KLVAE -from safetensors.torch import load_file, save_file +from safetensors.torch import load_file try: @@ -200,6 +196,9 @@ class ZImageModel(BaseModel): qtype=qtype, quantize_device=self.device_torch, config_path=base_model_path if self.is_single_file else None, + use_comfy_weights=self.model_config.model_kwargs.get( + "use_comfy_weights", True + ), ) flush() @@ -400,31 +399,17 @@ class ZImageModel(BaseModel): return ZImageTransformer2DModel.get_quantization_exclude_modules() def save_model(self, output_path, meta, save_dtype): + # comfy-format single-file save (the standard save format regardless + # of how the model was loaded); prequantized layers keep their + # quantized storage transformer: ZImageTransformer2DModel = unwrap_model(self.model) - if self.is_single_file: - # loaded from a single-file checkpoint, save back in that format - sd = transformer.state_dict() - save_dict = {} - for key, value in sd.items(): - # dequantize any quantized (e.g. torchao) weights so we save plain tensors - save_dict[key] = ( - dequantize_if_quantized(value).clone().to("cpu", dtype=save_dtype) - ) - save_dict = transformer.convert_state_dict_on_save(save_dict) - - if not output_path.endswith(".safetensors"): - output_path += ".safetensors" - meta = get_meta_for_safetensors(meta, name=self.arch) - save_file(save_dict, output_path, metadata=meta) - else: - transformer.save_pretrained( - save_directory=os.path.join(output_path, "transformer"), - safe_serialization=True, - ) - - meta_path = os.path.join(output_path, "aitk_meta.yaml") - with open(meta_path, "w") as f: - yaml.dump(meta, f) + if not output_path.endswith(".safetensors"): + output_path += ".safetensors" + transformer.save_model( + output_path, + dtype=save_dtype, + metadata=get_meta_for_safetensors(meta, name=self.arch), + ) def get_loss_target(self, *args, **kwargs): noise = kwargs.get("noise") diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index b5d7fd1..0048e19 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -252,10 +252,89 @@ loads via diffusers. Nothing about sources or outputs changes yet. Suggested ord (`dequantize_if_quantized` everywhere), raw-`quantize()` → `quantize_model()` ### Phase 2 — comfy weights become the preferred source -- [ ] Wire `comfy_weight_names` per model; standard-repo `name_or_path` + existing - comfy weights → load comfy -- [ ] Comfy-format save (bf16 + convrot8/nvfp4 quantized) as the default - full-weight save + +Decisions: +- Comfy weights come from the Comfy-Org hub repos (per-model repos, comfy + layout nested under `split_files/` — stripped when placing files into + MODELS_PATH). Repos ship several precision variants of each component. +- **Selection preference: convrot8 > float8 mixed > float8 > bf16 > fp16** + (`resolver.comfy_precision_rank`; nvfp4/unmarked rank last and are only + used when explicitly listed). **Local-first**: the best-ranked LOCAL + candidate wins; only when no candidate is local is the best-ranked one + downloaded. +- Per-model candidate lists (`aitk_comfy_weight_names`) hold only the + variants the class can actually digest. Constraint discovered: comfy + convrot/quantized files carry markers on the ORIGINAL module layout (e.g. + z_image's fused attention.qkv) — diffusers-layout classes with split + modules can't attach them until they grow fused-layout support; until then + those models list bf16/fp8 variants only. (Vendored comfy-layout classes — + minimax/ltx pattern — take convrot directly.) +- The standard repo still supplies the config; local dirs and unregistered + repos load as-is; `model_kwargs.use_comfy_weights: false` opts out. + +- [x] Mechanism: `resolver.comfy_precision_rank` / `comfy_local_rel` / + `resolve_comfy_candidates` + `OstrisModelMixin.resolve_comfy_weights`, + integrated into `load_model` (comfy file preferred for registered + standard repos, downloaded into the shared comfy layout) +- [x] First wired model: z_image (Comfy-Org/z_image_turbo, bf16 candidate). + Verified end-to-end against the real shared ComfyUI folder + (MODELS_PATH=/mnt/Models/comfy_models): standard-repo name_or_path + resolved to the locally-present comfy file and produced the identical + generation to the diffusers-shards load +- [x] qwen_image wired: its comfy files use the diffusers key layout + directly — `fp8mixed` (float8-mixed, rank 1) attaches its 839 + `float8_e4m3fn` markers straight onto the class; candidates fp8mixed → + fp8_e4m3fn (raw cast) → bf16. Verified with real weights: loaded the + shared folder's local fp8 file (local-first, no download) and generated + correctly. Holder skips re-quantization for prequantized checkpoints. +- [x] New `float8_e4m3fn` Ostris backend (toolkit/util/float8_quant.py): + ComfyUI's fp8 + per-tensor-scale storage with dequantized matmul, in + get_ostris_quantizer + comfy import/export. Round-trip verified. +- [x] wan family wired: comfy wan files (original key layout) convert via + diffusers' own `convert_wan_transformer_to_diffusers` (rename-only, so + quantized weight/scale keys ride along with their modules). Candidate + keys support `(repo, subfolder)` tuples for wan2.2 A14B's dual DiTs + (transformer = high noise, transformer_2 = low noise) and per-entry + comfy-repo overrides ({"repo": ..., "files": [...]}) since wan2.1 and + 2.2 files live in different Comfy-Org repos. Wired: 2.2 TI2V-5B, + T2V/I2V-A14B (fp8_scaled), 2.1 T2V 1.3B/14B, I2V 480P/720P. Verified + with real weights: wan21 1.3B (downloaded comfy bf16) and wan22 5B + (local comfy fp16) both load through the converter and generate video. + Fix along the way: the mixin's meta build now uses accelerate + init_empty_weights (params meta, buffers real) so init-computed + non-persistent buffers like wan's rope tables materialize. +- [x] Legacy ComfyUI scaled-fp8 support (`scaled_fp8` marker + per-layer + fp8 weight / scalar scale_weight, e.g. every wan *_fp8_scaled file): + imports onto the float8 backend; scale_input (activation quant) is + dropped, matmuls run dequantized. +- [ ] Wire remaining archs' candidate lists (chroma/others as their key + conversions are verified per file); wan comfy-format save needs the + inverse key mapping (defer with the other save flips) +- [x] Fused-layout quantized attach for diffusers-split classes: + `split_fused_quantized_keys` / `fuse_split_quantized_keys` + (comfy_quant_import) do exact out-dim row surgery on quantized comfy + entries for all three formats (int8 rows+scales slice; fp8 scalar and + nvfp4 per-tensor scales shared; nvfp4 block scales + unswizzle→split→reswizzle). z_image's load/save converters use them, so + its convrot8 candidate is live and top-ranked. Unit-verified exact both + directions. +- [x] Comfy-format save: `save_model` auto-keeps quantized storage + (comfy_quant markers) for convrot8 / nvfp4 / convrotcomfyw4a4 layers via + `toolkit/util/comfy_quant_export.py` (inverse of comfy_quant_import; + nvfp4 nibbles re-swapped + scales re-swizzled to the cuBLAS tile + layout), plain layers save at bf16; partially-exportable models fall + back to dequantized. Round-trip verified: save → mixin reload → outputs + match for convrot8, nvfp4, and plain layers. +- [x] Save unification started: z_image and qwen_image holders now save + comfy-format single files via the mixin regardless of how they loaded + (z_image's dual-style branch deleted). Real round trip verified: the + published z_image int8_convrot file loads (270 quantized linears, + split-attach), resaves to the IDENTICAL 857-key comfy layout with + bit-exact fused qkv weights/scales/markers, and the reload's quantized + forward is bit-identical — toolkit saves are byte-compatible with + ComfyUI. +- [ ] Flip the remaining per-arch `save_model` overrides as each arch's + save-side key conversion is in place - [ ] Publish/verify comfy repacks per model as they flip ### Phase 3 — live server @@ -291,9 +370,10 @@ loads via diffusers. Nothing about sources or outputs changes yet. Suggested ord device); ideogram4's fp8 release renders its own "blocked by safety filter" card for a plain cat prompt (model behavior, not a bug — investigate its trigger). -- [ ] Round-trip test per model: load → save comfy format → reload from the save → - outputs match (bf16) / load cleanly (quantized saves). Lands with the - Phase 2 comfy save path. +- [x] Round-trip verified for the first comfy-save arch: z_image convrot + load → comfy save → identical key set + bit-exact quantized entries vs + the published file → reload → bit-identical quantized forward. Extend + per arch as saves flip. - [ ] Each newly migrated model adds its test in the same PR as its migration. ## TODO / look at later diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index cacc49f..33b9fa9 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -51,12 +51,15 @@ class OstrisModelMixin: # .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 + # ---- 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 - # (ComfyUI folder layout: diffusion_models/, text_encoders/, vae/, ...) - aitk_comfy_weight_names: Dict[str, str] = {} + # 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 @@ -110,7 +113,12 @@ class OstrisModelMixin: @classmethod def aitk_from_config(cls, config): - with torch.device("meta"): + # 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): return cls.from_config(config) # ------------------------------------------------------------------ @@ -128,6 +136,7 @@ class OstrisModelMixin: exclude_quant_modules: Optional[List[str]] = None, config_path: Optional[str] = None, subfolder: Optional[str] = None, + use_comfy_weights: bool = True, **kwargs, ): """Load a model universally from a given name or path. @@ -136,6 +145,13 @@ class OstrisModelMixin: 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. @@ -149,6 +165,18 @@ class OstrisModelMixin: 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) + 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( @@ -241,7 +269,9 @@ class OstrisModelMixin: (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) + has_quant_markers = "scaled_fp8" in state_dict or 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) @@ -299,17 +329,42 @@ class OstrisModelMixin: # ------------------------------------------------------------------ @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 + 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, + ) -> 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). - rel_path = cls.aitk_comfy_weight_names.get(name_or_path) - if rel_path is None: + 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_file( - rel_path, repo_id=cls.aitk_comfy_repo, local_only=True + return resolve_comfy_candidates( + candidates, + repo_id=repo_id, + hf_token=hf_token, + status_fn=status_fn, + local_only=local_only, ) # ------------------------------------------------------------------ @@ -374,21 +429,47 @@ class OstrisModelMixin: 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 weights are - dequantized to full precision; dtype, when given, casts the floating - point tensors. (Quantized-storage saves — comfy_quant markers — come - with Phase 2.)""" + 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) @@ -479,5 +560,7 @@ class OstrisTransformersMixin(OstrisModelMixin): @classmethod def aitk_from_config(cls, config): - with torch.device("meta"): + from accelerate import init_empty_weights + + with init_empty_weights(include_buffers=False): return cls(config) diff --git a/toolkit/models/v2/diffusion_models/qwen_image.py b/toolkit/models/v2/diffusion_models/qwen_image.py index 6d5b7d7..3c4e689 100644 --- a/toolkit/models/v2/diffusion_models/qwen_image.py +++ b/toolkit/models/v2/diffusion_models/qwen_image.py @@ -11,13 +11,36 @@ class QwenImageTransformer2DModel( aitk_subfolder = "transformer" aitk_config_repo = "Qwen/Qwen-Image" + aitk_comfy_repo = "Comfy-Org/Qwen-Image_ComfyUI" + # the comfy files use the diffusers key layout directly (no conversion); + # fp8mixed carries float8_e4m3fn comfy_quant markers that attach straight + # onto this class's modules + aitk_comfy_weight_names = { + "Qwen/Qwen-Image": [ + "split_files/diffusion_models/qwen_image_fp8mixed.safetensors", + # raw fp8 cast (no markers, diffusers keys) — loads via from_single_file + "split_files/diffusion_models/qwen_image_fp8_e4m3fn.safetensors", + "split_files/diffusion_models/qwen_image_bf16.safetensors", + ], + } + @classmethod def get_transformer_block_names(cls): return ["transformer_blocks"] @classmethod def _load_single_file(cls, file_path, dtype, config_path=None, subfolder=None): - # single-file checkpoints in the wild carry diffusers or original key + from safetensors import safe_open + + with safe_open(file_path, framework="pt") as f: + has_markers = any(k.endswith(".comfy_quant") for k in f.keys()) + if has_markers: + # comfy prequantized checkpoint (diffusers key layout): the mixin + # path attaches the quantized layers + return super()._load_single_file( + file_path, dtype, config_path=config_path, subfolder=subfolder + ) + # other single-file checkpoints carry diffusers or original key # layouts; diffusers' single-file machinery owns that conversion model = cls.from_single_file( file_path, diff --git a/toolkit/models/v2/diffusion_models/wan.py b/toolkit/models/v2/diffusion_models/wan.py index 2ba9e88..98ddd1c 100644 --- a/toolkit/models/v2/diffusion_models/wan.py +++ b/toolkit/models/v2/diffusion_models/wan.py @@ -8,6 +8,79 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin): aitk_subfolder = "transformer" + aitk_comfy_repo = "Comfy-Org/Wan_2.2_ComfyUI_Repackaged" + # comfy wan files use the original key layout; convert_state_dict_on_load + # renames them (pure substring renames, so the legacy scaled_fp8 weight / + # scale keys ride along with their modules). wan2.2 A14B checkpoints hold + # two DiTs, keyed by (repo, subfolder): transformer = high noise, + # transformer_2 = low noise. wan2.1 entries override the comfy repo. + aitk_comfy_weight_names = { + "Wan-AI/Wan2.2-TI2V-5B-Diffusers": [ + "split_files/diffusion_models/wan2.2_ti2v_5B_fp16.safetensors", + ], + ("Wan-AI/Wan2.2-T2V-A14B-Diffusers", "transformer"): [ + "split_files/diffusion_models/wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors", + ], + ("Wan-AI/Wan2.2-T2V-A14B-Diffusers", "transformer_2"): [ + "split_files/diffusion_models/wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors", + ], + ("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "transformer"): [ + "split_files/diffusion_models/wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors", + ], + ("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "transformer_2"): [ + "split_files/diffusion_models/wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors", + ], + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": { + "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged", + "files": [ + "split_files/diffusion_models/wan2.1_t2v_1.3B_bf16.safetensors", + "split_files/diffusion_models/wan2.1_t2v_1.3B_fp16.safetensors", + ], + }, + "Wan-AI/Wan2.1-T2V-14B-Diffusers": { + "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged", + "files": [ + "split_files/diffusion_models/wan2.1_t2v_14B_fp8_scaled.safetensors", + "split_files/diffusion_models/wan2.1_t2v_14B_bf16.safetensors", + "split_files/diffusion_models/wan2.1_t2v_14B_fp16.safetensors", + ], + }, + "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": { + "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged", + "files": [ + "split_files/diffusion_models/wan2.1_i2v_480p_14B_fp8_scaled.safetensors", + "split_files/diffusion_models/wan2.1_i2v_480p_14B_bf16.safetensors", + "split_files/diffusion_models/wan2.1_i2v_480p_14B_fp16.safetensors", + ], + }, + "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": { + "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged", + "files": [ + "split_files/diffusion_models/wan2.1_i2v_720p_14B_fp8_scaled.safetensors", + "split_files/diffusion_models/wan2.1_i2v_720p_14B_bf16.safetensors", + "split_files/diffusion_models/wan2.1_i2v_720p_14B_fp16.safetensors", + ], + }, + } + @classmethod def get_transformer_block_names(cls): return ["blocks"] + + @classmethod + def convert_state_dict_on_load(cls, state_dict): + # original/comfy wan keys -> diffusers layout via diffusers' own + # converter (rename-only for the base/i2v/vace variants) + is_original = any( + ".self_attn." in k + or ".cross_attn." in k + or k.startswith("model.diffusion_model.") + for k in state_dict + ) + if not is_original: + return state_dict + from diffusers.loaders.single_file_utils import ( + convert_wan_transformer_to_diffusers, + ) + + return convert_wan_transformer_to_diffusers(dict(state_dict)) diff --git a/toolkit/models/v2/diffusion_models/z_image.py b/toolkit/models/v2/diffusion_models/z_image.py index 5ee32c4..794c3ec 100644 --- a/toolkit/models/v2/diffusion_models/z_image.py +++ b/toolkit/models/v2/diffusion_models/z_image.py @@ -11,6 +11,17 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix # repo to pull the config from when loading a single-file checkpoint aitk_config_repo = "Tongyi-MAI/Z-Image-Turbo" + aitk_comfy_repo = "Comfy-Org/z_image_turbo" + # the int8_convrot file marks the FUSED attention.qkv modules; the load + # converter splits those entries exactly into to_q/to_k/to_v (row-sliced + # qdata + scales), so it attaches onto this split layout directly + aitk_comfy_weight_names = { + "Tongyi-MAI/Z-Image-Turbo": [ + "split_files/diffusion_models/z_image_turbo_int8_convrot.safetensors", + "split_files/diffusion_models/z_image_turbo_bf16.safetensors", + ], + } + @classmethod def get_transformer_block_names(cls): return ["layers"] @@ -36,6 +47,22 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix @classmethod def convert_state_dict_on_load(cls, state_dict): """Convert a single-file Z-Image checkpoint to diffusers transformer keys.""" + from toolkit.util.comfy_quant_import import split_fused_quantized_keys + + state_dict = dict(state_dict) + # quantized fused qkv entries (comfy convrot/fp8/nvfp4 files) split + # exactly into the three projections: row-consecutive weights/scales + for marker_key in [ + k for k in list(state_dict) if k.endswith(".attention.qkv.comfy_quant") + ]: + prefix = marker_key[: -len(".comfy_quant")] + base = prefix[: -len(".qkv")] + split_fused_quantized_keys( + state_dict, + prefix, + [f"{base}.to_q", f"{base}.to_k", f"{base}.to_v"], + ) + new_sd = {} for key, value in state_dict.items(): k = key @@ -47,9 +74,10 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix 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") + # module-prefix renames so weight/bias/scale/marker keys all map + k = k.replace(".attention.out.", ".attention.to_out.0.") + k = k.replace(".attention.q_norm.", ".attention.norm_q.") + k = k.replace(".attention.k_norm.", ".attention.norm_k.") if k.startswith("x_embedder."): k = "all_x_embedder.2-1." + k[len("x_embedder.") :] elif k.startswith("final_layer."): @@ -60,6 +88,20 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix @classmethod def convert_state_dict_on_save(cls, state_dict): """Convert a diffusers transformer state dict back to the single-file layout.""" + from toolkit.util.comfy_quant_import import fuse_split_quantized_keys + + state_dict = dict(state_dict) + # quantized split projections fuse back into the single-file qkv entry + for marker_key in [ + k for k in list(state_dict) if k.endswith(".attention.to_q.comfy_quant") + ]: + base = marker_key[: -len(".to_q.comfy_quant")] + fuse_split_quantized_keys( + state_dict, + [f"{base}.to_q", f"{base}.to_k", f"{base}.to_v"], + f"{base}.qkv", + ) + new_sd = {} qkv_cache = {} for key, value in state_dict.items(): @@ -90,9 +132,9 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix 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") + k = k.replace(".attention.to_out.0.", ".attention.out.") + k = k.replace(".attention.norm_q.", ".attention.q_norm.") + k = k.replace(".attention.norm_k.", ".attention.k_norm.") 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."): diff --git a/toolkit/models/v2/resolver.py b/toolkit/models/v2/resolver.py index fda8ff0..adef8de 100644 --- a/toolkit/models/v2/resolver.py +++ b/toolkit/models/v2/resolver.py @@ -16,6 +16,78 @@ from typing import Callable, Iterable, Optional from toolkit.paths import MODELS_PATH +def comfy_precision_rank(filename: str) -> int: + """Load-preference rank for a comfy weight filename: + convrot8 (0) > float8 mixed (1) > float8 (2) > bf16 (3) > fp16 (4) > + anything else, e.g. nvfp4 or unmarked (5).""" + name = os.path.basename(filename).lower() + if "convrot" in name: + return 0 + is_fp8 = "fp8" in name or "float8" in name or "e4m3" in name + if is_fp8 and "mixed" in name: + return 1 + if is_fp8: + return 2 + if "bf16" in name: + return 3 + if "fp16" in name: + return 4 + return 5 + + +def comfy_local_rel(repo_rel: str) -> str: + """Repo file path -> ComfyUI models-folder path. Comfy-Org repos nest the + comfy layout under a packaging prefix (split_files/, non_official/) that is + not part of the shared models folder layout.""" + for prefix in ("split_files/", "non_official/"): + if repo_rel.startswith(prefix): + return repo_rel[len(prefix):] + return repo_rel + + +def resolve_comfy_candidates( + candidates: Iterable[str], + repo_id: str, + hf_token: Optional[str] = None, + status_fn: Optional[Callable[[str], None]] = None, + local_only: bool = False, +) -> Optional[str]: + """Pick the best comfy weight file among precision variants of one + component (repo-relative paths, ranked by comfy_precision_rank then list + order). The best-ranked LOCAL candidate wins; only when no candidate is + local is the best-ranked one downloaded to its comfy-layout location + under MODELS_PATH.""" + candidates = list(candidates) + ordered = sorted( + candidates, key=lambda c: (comfy_precision_rank(c), candidates.index(c)) + ) + for repo_rel in ordered: + found = resolve_comfy_file( + comfy_local_rel(repo_rel), repo_id, local_only=True + ) + if found is not None: + return found + if local_only: + return None + + import huggingface_hub + + best = ordered[0] + local_rel = comfy_local_rel(best) + if status_fn is not None: + status_fn(f"Downloading {best} from {repo_id} into {MODELS_PATH}") + path = huggingface_hub.hf_hub_download( + repo_id=repo_id, filename=best, token=hf_token, local_dir=MODELS_PATH + ) + target = os.path.join(MODELS_PATH, local_rel) + if os.path.abspath(path) != os.path.abspath(target): + # move out of the packaging prefix into the shared comfy layout + os.makedirs(os.path.dirname(target), exist_ok=True) + os.replace(path, target) + return target + return path + + def find_file_recursive(root_dir: str, filename: str) -> Optional[str]: """First (breadth-stable, sorted) match of ``filename`` anywhere under ``root_dir``.""" diff --git a/toolkit/util/comfy_quant_export.py b/toolkit/util/comfy_quant_export.py new file mode 100644 index 0000000..657c2c6 --- /dev/null +++ b/toolkit/util/comfy_quant_export.py @@ -0,0 +1,104 @@ +"""Export toolkit-quantized modules into ComfyUI ``comfy_quant`` checkpoints — +the inverse of toolkit/util/comfy_quant_import.py. + +Every quantized OstrisLinear whose backend has a comfy storage format emits +its quantized tensors plus the ``.comfy_quant`` uint8 JSON marker: + + - convrot8 (int8_tensorwise + convrot): weight int8 [out, in], fp32 + weight_scale [out, 1] (comfy_kitchen's per-channel convention) + - nvfp4: high-nibble-first packed fp4 pairs, e4m3 block scales re-swizzled + to the cuBLAS 128x4 tile layout, fp32 weight_scale_2 per-tensor scale and + optional AWQ pre_quant_scale + - float8_e4m3fn: fp8_e4m3 weight + one fp32 per-tensor weight_scale + - convrotcomfyw4a4: via convrot_quant.export_comfy_convrot_w4a4 + +Biases are NOT emitted here — they are ordinary parameters and flow through +the regular state_dict path. +""" + +import json +from typing import Dict, List, Tuple + +import torch + +from toolkit.util.ostris_quant import OstrisLinear + + +def comfy_quant_marker(conf: dict) -> torch.Tensor: + return torch.tensor(list(json.dumps(conf).encode("utf-8")), dtype=torch.uint8) + + +@torch.no_grad() +def export_comfy_quantized_layers( + root: torch.nn.Module, +) -> Tuple[Dict[str, torch.Tensor], List[str], List[str]]: + """Comfy-format state-dict entries for every quantized OstrisLinear in + ``root``. Returns ``(entries, exported_names, unexportable_names)`` — + entries are keyed by module path in root's layout (run + convert_state_dict_on_save afterwards for the checkpoint layout); + unexportable_names lists quantized modules whose backend has no comfy + storage format (the caller decides whether to dequantize instead).""" + from toolkit.util.convrot_quant import ( + ConvRotComfyW4A4Quantizer, + export_comfy_convrot_w4a4, + ) + from toolkit.util.nvfp4_quant import swap_nvfp4_nibbles, swizzle_nvfp4_scales + + entries: Dict[str, torch.Tensor] = {} + exported: List[str] = [] + unexportable: List[str] = [] + + for name, module in root.named_modules(): + if not isinstance(module, OstrisLinear): + continue + + if hasattr(module, "cr8_qdata"): + rot = int(getattr(module, "cr8_rot_size", 1) or 1) + conf = {"format": "int8_tensorwise"} + if rot > 1: + conf.update({"convrot": True, "convrot_groupsize": rot}) + entries[f"{name}.weight"] = module.cr8_qdata.detach().cpu().contiguous() + entries[f"{name}.weight_scale"] = ( + module.cr8_scales.view(torch.float32) + .detach() + .cpu() + .reshape(module.out_features, 1) + .contiguous() + ) + entries[f"{name}.comfy_quant"] = comfy_quant_marker(conf) + elif hasattr(module, "nv4_qdata"): + entries[f"{name}.weight"] = swap_nvfp4_nibbles( + module.nv4_qdata.detach().cpu() + ) + scales = module.nv4_scales.view(torch.float8_e4m3fn).detach().cpu() + entries[f"{name}.weight_scale"] = swizzle_nvfp4_scales( + scales.reshape(module.out_features, module.in_features // 16) + ).view(torch.float8_e4m3fn) + entries[f"{name}.weight_scale_2"] = ( + module.nv4_pts.view(torch.float32).detach().cpu().reshape(()) + ) + if hasattr(module, "nv4_pre_scale"): + entries[f"{name}.pre_quant_scale"] = ( + module.nv4_pre_scale.view(torch.float32).detach().cpu().contiguous() + ) + entries[f"{name}.comfy_quant"] = comfy_quant_marker({"format": "nvfp4"}) + elif hasattr(module, "f8_qdata"): + entries[f"{name}.weight"] = module.f8_qdata.detach().cpu().contiguous() + entries[f"{name}.weight_scale"] = ( + module.f8_scale.view(torch.float32).detach().cpu().reshape(()) + ) + entries[f"{name}.comfy_quant"] = comfy_quant_marker( + {"format": "float8_e4m3fn", "full_precision_matrix_mult": True} + ) + elif isinstance(module.ostris_quantizer, ConvRotComfyW4A4Quantizer): + layer_entries = export_comfy_convrot_w4a4(module, f"{name}.") + layer_entries.pop(f"{name}.bias", None) + entries.update( + {k: v.detach().cpu() if torch.is_tensor(v) else v for k, v in layer_entries.items()} + ) + else: + unexportable.append(name) + continue + exported.append(name) + + return entries, exported, unexportable diff --git a/toolkit/util/comfy_quant_import.py b/toolkit/util/comfy_quant_import.py index 833d6b2..6cd16b6 100644 --- a/toolkit/util/comfy_quant_import.py +++ b/toolkit/util/comfy_quant_import.py @@ -71,6 +71,135 @@ class Int8Embedding(torch.nn.Module): return out.to(input_ids.device).reshape(*input_ids.shape, self.embedding_dim) +@torch.no_grad() +def split_fused_quantized_keys( + state_dict: Dict[str, torch.Tensor], + prefix: str, + dst_prefixes, +) -> Dict[str, torch.Tensor]: + """Split one fused quantized comfy entry (``.weight`` / + ``.weight_scale`` / ``.comfy_quant`` / ...) into equal row ranges under + ``dst_prefixes`` (out-dim concat order). Exact for every supported format: + int8 rows and their per-row scales slice; fp8's per-tensor scale and + nvfp4's weight_scale_2 / pre_quant_scale are shared by every split; nvfp4 + block scales are unswizzled, row-split, and re-swizzled. Mutates and + returns state_dict. Used by classes whose module layout splits a fused + checkpoint projection (e.g. qkv -> to_q/to_k/to_v).""" + from toolkit.util.nvfp4_quant import swizzle_nvfp4_scales + + marker = state_dict.pop(f"{prefix}.comfy_quant") + conf = parse_comfy_quant_blob(marker) + fmt = conf.get("format") + + weight = state_dict.pop(f"{prefix}.weight") + scale = state_dict.pop(f"{prefix}.weight_scale", None) + pts = state_dict.pop(f"{prefix}.weight_scale_2", None) + pre = state_dict.pop(f"{prefix}.pre_quant_scale", None) + bias = state_dict.pop(f"{prefix}.bias", None) + state_dict.pop(f"{prefix}.input_scale", None) + + n = len(dst_prefixes) + if weight.shape[0] % n != 0: + raise ValueError( + f"{prefix}: fused out dim {weight.shape[0]} does not split into {n}" + ) + rows = weight.shape[0] // n + + scale_parts = None + if scale is not None: + if fmt == "nvfp4": + in_features = weight.shape[1] * 2 # packed fp4 pairs + full = unswizzle_nvfp4_scales( + scale.view(torch.float8_e4m3fn), weight.shape[0], in_features // 16 + ) + scale_parts = [ + swizzle_nvfp4_scales(p).view(torch.float8_e4m3fn) + for p in full.split(rows, dim=0) + ] + elif scale.ndim == 0 or scale.numel() == 1: + scale_parts = [scale.clone() for _ in range(n)] + else: + scale_parts = list(scale.reshape(weight.shape[0], -1).split(rows, dim=0)) + + for i, dst in enumerate(dst_prefixes): + state_dict[f"{dst}.comfy_quant"] = marker.clone() + state_dict[f"{dst}.weight"] = weight[i * rows : (i + 1) * rows].contiguous() + if scale_parts is not None: + state_dict[f"{dst}.weight_scale"] = scale_parts[i].contiguous() + if pts is not None: + state_dict[f"{dst}.weight_scale_2"] = pts.clone() + if pre is not None: + state_dict[f"{dst}.pre_quant_scale"] = pre.clone() + if bias is not None: + state_dict[f"{dst}.bias"] = bias[i * rows : (i + 1) * rows].contiguous() + return state_dict + + +@torch.no_grad() +def fuse_split_quantized_keys( + state_dict: Dict[str, torch.Tensor], + src_prefixes, + prefix: str, +) -> Dict[str, torch.Tensor]: + """Inverse of split_fused_quantized_keys: concatenate N split quantized + comfy entries back into one fused entry (out-dim concat in src order). + All parts must share the same format config; fp8 parts must share the same + per-tensor scale (true for entries produced by the splitter). Mutates and + returns state_dict.""" + from toolkit.util.nvfp4_quant import swizzle_nvfp4_scales + + markers = [state_dict.pop(f"{p}.comfy_quant") for p in src_prefixes] + confs = [parse_comfy_quant_blob(m) for m in markers] + if any(c != confs[0] for c in confs[1:]): + raise ValueError(f"{prefix}: split parts carry different quant configs") + fmt = confs[0].get("format") + + weights = [state_dict.pop(f"{p}.weight") for p in src_prefixes] + scales = [state_dict.pop(f"{p}.weight_scale", None) for p in src_prefixes] + ptss = [state_dict.pop(f"{p}.weight_scale_2", None) for p in src_prefixes] + pres = [state_dict.pop(f"{p}.pre_quant_scale", None) for p in src_prefixes] + biases = [state_dict.pop(f"{p}.bias", None) for p in src_prefixes] + + state_dict[f"{prefix}.comfy_quant"] = markers[0] + weight = torch.cat(weights, dim=0).contiguous() + state_dict[f"{prefix}.weight"] = weight + if scales[0] is not None: + if fmt == "nvfp4": + in_features = weight.shape[1] * 2 + rows = [w.shape[0] for w in weights] + full = torch.cat( + [ + unswizzle_nvfp4_scales( + s.view(torch.float8_e4m3fn), r, in_features // 16 + ) + for s, r in zip(scales, rows) + ], + dim=0, + ) + state_dict[f"{prefix}.weight_scale"] = swizzle_nvfp4_scales(full).view( + torch.float8_e4m3fn + ) + elif scales[0].ndim == 0 or scales[0].numel() == 1: + if any( + not torch.equal(s.reshape(-1), scales[0].reshape(-1)) for s in scales[1:] + ): + raise ValueError( + f"{prefix}: per-tensor scales differ across split parts" + ) + state_dict[f"{prefix}.weight_scale"] = scales[0] + else: + state_dict[f"{prefix}.weight_scale"] = torch.cat( + [s.reshape(w.shape[0], -1) for s, w in zip(scales, weights)], dim=0 + ).contiguous() + if ptss[0] is not None: + state_dict[f"{prefix}.weight_scale_2"] = ptss[0] + if pres[0] is not None: + state_dict[f"{prefix}.pre_quant_scale"] = pres[0] + if biases[0] is not None: + state_dict[f"{prefix}.bias"] = torch.cat(biases, dim=0).contiguous() + return state_dict + + def _to_ostris(module: torch.nn.Linear, quantizer, orig_dtype: torch.dtype) -> OstrisLinear: if "weight" in module._parameters: del module._parameters["weight"] @@ -128,7 +257,19 @@ def import_comfy_quantized_layers( "expected nn.Linear or nn.Embedding" ) - if fmt == "int8_tensorwise": + if fmt == "float8_e4m3fn": + # fp8_e4m3 weight + fp32 per-tensor scale, dequantized matmul + from toolkit.util.float8_quant import Float8Quantizer + + quantizer = get_ostris_quantizer("float8_e4m3fn") + Float8Quantizer.attach_( + module, + weight.view(torch.float8_e4m3fn) + if weight.dtype != torch.float8_e4m3fn + else weight, + weight_scale, + ) + elif fmt == "int8_tensorwise": rot = int(conf.get("convrot_groupsize", 256)) if conf.get("convrot") else 1 quantizer = get_ostris_quantizer("convrot8") module.register_buffer("cr8_qdata", weight.contiguous(), persistent=False) @@ -158,7 +299,7 @@ def import_comfy_quantized_layers( else: raise ValueError( f"Unsupported comfy quant format {fmt!r} on {prefix} " - "(supported: int8_tensorwise, nvfp4)" + "(supported: int8_tensorwise, nvfp4, float8_e4m3fn)" ) # drop unused calibration extras if present @@ -174,4 +315,40 @@ def import_comfy_quantized_layers( ) converted += 1 + # legacy ComfyUI scaled-fp8 checkpoints (e.g. the wan *_fp8_scaled files): + # a top-level ``scaled_fp8`` marker tensor plus per-layer fp8 ``weight`` + # and scalar fp32 ``scale_weight`` — the float8 backend's exact storage. + # ``scale_input`` (activation quant) is dropped; matmuls run dequantized. + if "scaled_fp8" in state_dict: + from toolkit.util.float8_quant import Float8Quantizer + + state_dict.pop("scaled_fp8") + for scale_key in [k for k in state_dict if k.endswith(".scale_weight")]: + prefix = scale_key[: -len(".scale_weight")] + module_path = key_map(prefix) if key_map is not None else prefix + module = root.get_submodule(module_path) + if not isinstance(module, torch.nn.Linear): + raise ValueError( + f"scaled_fp8 entry {prefix} points at {type(module).__name__}, " + "expected nn.Linear" + ) + weight = state_dict.pop(f"{prefix}.weight") + scale = state_dict.pop(scale_key) + state_dict.pop(f"{prefix}.scale_input", None) + quantizer = get_ostris_quantizer("float8_e4m3fn") + Float8Quantizer.attach_( + module, + weight + if weight.dtype == torch.float8_e4m3fn + else weight.view(torch.float8_e4m3fn), + scale, + ) + _to_ostris(module, quantizer, orig_dtype) + bias = state_dict.pop(f"{prefix}.bias", None) + if bias is not None and module.bias is not None: + module._parameters["bias"] = torch.nn.Parameter( + bias.detach().clone(), requires_grad=False + ) + converted += 1 + return state_dict, converted diff --git a/toolkit/util/float8_quant.py b/toolkit/util/float8_quant.py new file mode 100644 index 0000000..0ac2ff4 --- /dev/null +++ b/toolkit/util/float8_quant.py @@ -0,0 +1,53 @@ +"""ComfyUI-style float8 weight storage as an Ostris backend. + +Matches the comfy_quant ``{"format": "float8_e4m3fn", +"full_precision_matrix_mult": true}`` layout: the weight stored as +torch.float8_e4m3fn plus one fp32 per-tensor scale, matmuls running on the +dequantized weight (W8A16 numerics). Used both to import comfy fp8/fp8-mixed +checkpoints and to quantize/export in that format. +""" + +from typing import Optional + +import torch + +from toolkit.util.ostris_quant import OstrisLinear, OstrisQuantizer + +FLOAT8_QTYPES = ["float8_e4m3fn"] + +F8_MAX = torch.finfo(torch.float8_e4m3fn).max + + +class Float8Quantizer(OstrisQuantizer): + """fp8_e4m3 weight + fp32 per-tensor scale, dequantized matmul.""" + + def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None: + scale = (weight_fp32.abs().max() / F8_MAX).clamp(min=1e-12) + q = (weight_fp32 / scale).clamp(-F8_MAX, F8_MAX).to(torch.float8_e4m3fn) + self.attach_(module, q, scale) + + @staticmethod + def attach_( + module: torch.nn.Module, + qweight: torch.Tensor, # float8_e4m3fn (out, in) + scale: torch.Tensor, # fp32 scalar + ) -> None: + """Register the quantized representation on the module. Used both by + quantize_ and by importers of pre-quantized checkpoints.""" + module.register_buffer("f8_qdata", qweight.contiguous(), persistent=False) + module.register_buffer( + "f8_scale", + scale.detach().float().clone().reshape(1).view(torch.uint8), + persistent=False, + ) + + def dequantize(self, module: "OstrisLinear") -> torch.Tensor: + scale = module.f8_scale.view(torch.float32)[0] + return module.f8_qdata.to(torch.float32) * scale + + @torch.no_grad() + def requantize_(self, module: "OstrisLinear", fp_weight: torch.Tensor) -> None: + w = fp_weight.to(torch.float32) + scale = (w.abs().max() / F8_MAX).clamp(min=1e-12) + module.f8_qdata.copy_((w / scale).clamp(-F8_MAX, F8_MAX).to(torch.float8_e4m3fn)) + module.f8_scale.copy_(scale.reshape(1).view(torch.uint8)) diff --git a/toolkit/util/nvfp4_quant.py b/toolkit/util/nvfp4_quant.py index 72d7155..ef33147 100644 --- a/toolkit/util/nvfp4_quant.py +++ b/toolkit/util/nvfp4_quant.py @@ -59,6 +59,25 @@ def swap_nvfp4_nibbles(packed: torch.Tensor) -> torch.Tensor: return ((packed << 4) | (packed >> 4)).contiguous() +def swizzle_nvfp4_scales(scales: torch.Tensor) -> torch.Tensor: + """Inverse of unswizzle_nvfp4_scales: row-major (rows, cols) block scales + into the cuBLAS 128x4-tile layout ComfyUI checkpoints store (comfy_kitchen's + ``to_blocked``). Pads to tile boundaries when needed.""" + rows, cols = scales.shape + n_row_blocks = (rows + 127) // 128 + n_col_blocks = (cols + 3) // 4 + padded_rows = n_row_blocks * 128 + padded_cols = n_col_blocks * 4 + if (padded_rows, padded_cols) != (rows, cols): + padded = scales.new_zeros(padded_rows, padded_cols) + padded[:rows, :cols] = scales + scales = padded + x = scales.reshape(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3) + x = x.reshape(n_row_blocks, n_col_blocks, 4, 32, 4) + x = x.transpose(2, 3).reshape(-1, 32, 16) + return x.reshape(-1).contiguous() + + class Nvfp4Quantizer(OstrisQuantizer): """Block-16 nvfp4 weights, full-precision activations. One instance is shareable across modules.""" diff --git a/toolkit/util/ostris_quant.py b/toolkit/util/ostris_quant.py index 6384674..e5bd9c6 100644 --- a/toolkit/util/ostris_quant.py +++ b/toolkit/util/ostris_quant.py @@ -186,6 +186,7 @@ def get_ostris_quantizer(qtype: str) -> Optional[OstrisQuantizer]: """Resolve a qtype string to a quantizer backend instance, or None if the qtype does not belong to a custom backend. Add new backends here.""" from toolkit.util.convrot_quant import CONVROT_QTYPES, get_convrot_quantizer + from toolkit.util.float8_quant import FLOAT8_QTYPES, Float8Quantizer from toolkit.util.nvfp4_quant import NVFP4_QTYPES, Nvfp4Quantizer from toolkit.util.orbit_quant import ORBIT_QTYPES, OrbitQuantizer from toolkit.util.orbit_vq_quant import ORBIT_VQ_QTYPES, OrbitVQQuantizer @@ -200,6 +201,8 @@ def get_ostris_quantizer(qtype: str) -> Optional[OstrisQuantizer]: quantizer = get_convrot_quantizer(qtype) elif qtype in NVFP4_QTYPES: quantizer = Nvfp4Quantizer() + elif qtype in FLOAT8_QTYPES: + quantizer = Float8Quantizer() elif qtype in UINTX_QTYPES: quantizer = UIntXQuantizer(UINTX_QTYPES[qtype]) if quantizer is not None: From 45886f01b29e4b47f0166b50fb566315a1097691 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 27 Aug 2026 16:16:35 -0600 Subject: [PATCH 05/14] Phase 2 --- toolkit/models/v2/PLANNING.md | 8 ++++++-- toolkit/models/v2/text_encoders/umt5.py | 25 +++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index 0048e19..7b17bef 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -382,8 +382,12 @@ Decisions: raw `quantize()` call sites and silently dropped by `quantize_model()`; the ARA path inside `quantize_model` hardcodes `uint8`. Decide the unified behavior when consolidating. -- [ ] `toolkit/models/loaders/umt5.py` dead `comfy_files` param (wan21 comfy-TE - no-op) — fix when wan migrates. +- [x] wan comfy-TE resolved for real: UMT5TextEncoder carries comfy + candidates (fp8_e4m3fn_scaled via the legacy importer, fp16), files + already in transformers key layout (spiece blob dropped, tied + embed_tokens materialized). Verified: wan21 samples with the local + comfy fp8 TE. The loaders/umt5.py `comfy_files` param stays as a + no-op shim for old callers. - [ ] Registry hardening: error (don't fall back to SD1) on unknown arch; lazy per-arch imports; single source of truth shared with the UI's `options.tsx` model list. diff --git a/toolkit/models/v2/text_encoders/umt5.py b/toolkit/models/v2/text_encoders/umt5.py index d0b2e99..8bd02e0 100644 --- a/toolkit/models/v2/text_encoders/umt5.py +++ b/toolkit/models/v2/text_encoders/umt5.py @@ -33,11 +33,36 @@ class UMT5TextEncoder(UMT5EncoderModel, OstrisTransformersMixin): aitk_subfolder = "text_encoder" aitk_tokenizer_subfolder = "tokenizer" + aitk_config_repo = "ai-toolkit/umt5_xxl_encoder" + + aitk_comfy_repo = "Comfy-Org/Wan_2.1_ComfyUI_repackaged" + # comfy umt5 files already use the transformers key layout; the + # fp8_e4m3fn_scaled variant is the legacy scaled-fp8 format (handled by + # the importer's float8 backend) + aitk_comfy_weight_names = { + "ai-toolkit/umt5_xxl_encoder": [ + "split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors", + "split_files/text_encoders/umt5_xxl_fp16.safetensors", + ], + } @classmethod def get_transformer_block_names(cls): return ["encoder.block"] + @classmethod + def convert_state_dict_on_load(cls, state_dict): + # drop the embedded sentencepiece blob and materialize the tied + # embed_tokens reference + state_dict = dict(state_dict) + state_dict.pop("spiece_model", None) + if ( + "encoder.embed_tokens.weight" not in state_dict + and "shared.weight" in state_dict + ): + state_dict["encoder.embed_tokens.weight"] = state_dict["shared.weight"] + return state_dict + @classmethod def load_tokenizer(cls, name_or_path=None, subfolder=None, **kwargs): # T5's tokenizer needs the _spm_precompiled_charsmap patch From 520d96aac3274fd68c28fe052988638ce749a8c0 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 27 Aug 2026 18:18:31 -0600 Subject: [PATCH 06/14] Phase 2 --- .../diffusion_models/anima/anima.py | 2 +- .../diffusion_models/chroma/chroma_model.py | 36 ++++----- .../chroma/chroma_radiance_model.py | 63 +++++---------- .../diffusion_models/f_light/f_light.py | 17 ++-- .../diffusion_models/flux2/flux2_model.py | 14 +--- .../diffusion_models/flux2/src/model.py | 4 +- .../flux_kontext/flux_kontext.py | 17 ++-- .../diffusion_models/ideogram4/ideogram4.py | 7 +- .../ideogram4/src/transformer.py | 4 +- .../diffusion_models/krea2/krea2.py | 12 +-- .../diffusion_models/krea2/src/mmdit.py | 4 +- .../diffusion_models/ltx2/ltx2.py | 30 +------ .../diffusion_models/mageflow/mageflow.py | 13 +-- .../mageflow/src/transformer.py | 4 +- .../diffusion_models/minimax_h3/minimax_h3.py | 42 +--------- .../minimax_h3/src/transformer.py | 28 ++++++- .../diffusion_models/omnigen2/__init__.py | 8 +- .../diffusion_models/wan22/wan22_14b_model.py | 21 ++--- .../zeta_chroma/zeta_chroma_model.py | 11 +-- .../zeta_chroma/zeta_chroma_transformer.py | 4 +- toolkit/models/v2/PLANNING.md | 80 ++++++++++++++----- toolkit/models/v2/_mixin.py | 41 ++++++++-- .../models/v2/diffusion_models/qwen_image.py | 10 ++- toolkit/models/v2/diffusion_models/wan.py | 51 ++++++++++++ toolkit/models/wan21/wan21.py | 18 ++--- toolkit/util/get_model.py | 28 ++++++- toolkit/util/quantize.py | 17 +++- 27 files changed, 333 insertions(+), 253 deletions(-) diff --git a/extensions_built_in/diffusion_models/anima/anima.py b/extensions_built_in/diffusion_models/anima/anima.py index 9d3772b..b22089a 100644 --- a/extensions_built_in/diffusion_models/anima/anima.py +++ b/extensions_built_in/diffusion_models/anima/anima.py @@ -271,7 +271,7 @@ class AnimaModel(BaseModel): flush() self.print_and_status_update("Quantizing Text Conditioner") - quantize(text_conditioner, weights=get_qtype(self.model_config.qtype)) + quantize(text_conditioner, weights=get_qtype(self.model_config.qtype_te)) freeze(text_conditioner) flush() diff --git a/extensions_built_in/diffusion_models/chroma/chroma_model.py b/extensions_built_in/diffusion_models/chroma/chroma_model.py index d9d2025..6a3aa9a 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_model.py @@ -11,10 +11,9 @@ from toolkit.basic import flush # from toolkit.pixel_shuffle_encoder import AutoencoderPixelMixer from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler -from toolkit.dequantize import patch_dequantization_on_save from toolkit.accelerator import unwrap_model -from optimum.quanto import freeze, QTensor -from toolkit.util.quantize import quantize, get_qtype +from optimum.quanto import QTensor +from toolkit.util.quantize import quantize, quantize_model from .pipeline import ChromaPipeline, prepare_latent_image_ids from einops import rearrange, repeat import random @@ -67,6 +66,9 @@ class FakeCLIP(torch.nn.Module): class ChromaModel(BaseModel): arch = "chroma" + def get_transformer_block_names(self): + return ["double_blocks", "single_blocks"] + def __init__( self, device, @@ -163,13 +165,10 @@ class ChromaModel(BaseModel): transformer.config.num_single_layers = transformer.params.depth_single_blocks if self.model_config.quantize: - # patch the state dict method - patch_dequantization_on_save(transformer) - quantization_type = get_qtype(self.model_config.qtype) + # block-streaming quantize (handles dequant-on-save patching, + # excludes, ARA, and quantize_kwargs) self.print_and_status_update("Quantizing transformer") - quantize(transformer, weights=quantization_type, - **self.model_config.quantize_kwargs) - freeze(transformer) + quantize_model(self, transformer) transformer.to(self.device_torch) else: transformer.to(self.device_torch, dtype=dtype) @@ -389,19 +388,16 @@ class ChromaModel(BaseModel): return self.text_encoder[1].encoder.block[0].layer[0].SelfAttention.q.weight.requires_grad def save_model(self, output_path, meta, save_dtype): + # comfy-format single-file save via the mixin (chroma's class keys ARE + # the original layout); handles torchao/Ostris dequant, not just quanto if not output_path.endswith(".safetensors"): - output_path = output_path + ".safetensors" - # only save the unet + output_path = output_path + ".safetensors" transformer: Chroma = unwrap_model(self.model) - state_dict = transformer.state_dict() - save_dict = {} - for k, v in state_dict.items(): - if isinstance(v, QTensor): - v = v.dequantize() - save_dict[k] = v.clone().to('cpu', dtype=save_dtype) - - meta = get_meta_for_safetensors(meta, name='chroma') - save_file(save_dict, output_path, metadata=meta) + transformer.save_model( + output_path, + dtype=save_dtype, + metadata=get_meta_for_safetensors(meta, name="chroma"), + ) def get_loss_target(self, *args, **kwargs): noise = kwargs.get('noise') diff --git a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py index d95dddc..2fcc390 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py @@ -10,10 +10,9 @@ from toolkit.basic import flush # from toolkit.pixel_shuffle_encoder import AutoencoderPixelMixer from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler -from toolkit.dequantize import patch_dequantization_on_save from toolkit.accelerator import unwrap_model -from optimum.quanto import freeze, QTensor -from toolkit.util.quantize import quantize, get_qtype +from optimum.quanto import QTensor +from toolkit.util.quantize import quantize, quantize_model from .pipeline import ChromaPipeline, prepare_latent_image_ids from einops import rearrange, repeat import random @@ -37,36 +36,16 @@ scheduler_config = { "use_dynamic_shifting": True } -class FakeConfig: - # for diffusers compatability - def __init__(self): - self.attention_head_dim = 128 - self.guidance_embeds = True - self.in_channels = 64 - self.joint_attention_dim = 4096 - self.num_attention_heads = 24 - self.num_layers = 19 - self.num_single_layers = 38 - self.patch_size = 1 - -class FakeCLIP(torch.nn.Module): - def __init__(self, device='cuda'): - super().__init__() - self.dtype = torch.bfloat16 - # the pipeline derives its execution device from this attribute; - # nn.Module.to() does not update it - self.device = device - self.text_model = None - self.tokenizer = None - self.model_max_length = 77 - - def forward(self, *args, **kwargs): - return torch.zeros(1, 1, 1).to(self.device) +# shared with the base chroma model (identical stubs) +from .chroma_model import FakeCLIP, FakeConfig class ChromaRadianceModel(BaseModel): arch = "chroma_radiance" + def get_transformer_block_names(self): + return ["double_blocks", "single_blocks"] + def __init__( self, device, @@ -165,13 +144,10 @@ class ChromaRadianceModel(BaseModel): transformer.config.num_single_layers = transformer.params.depth_single_blocks if self.model_config.quantize: - # patch the state dict method - patch_dequantization_on_save(transformer) - quantization_type = get_qtype(self.model_config.qtype) + # block-streaming quantize (handles dequant-on-save patching, + # excludes, ARA, and quantize_kwargs) self.print_and_status_update("Quantizing transformer") - quantize(transformer, weights=quantization_type, - **self.model_config.quantize_kwargs) - freeze(transformer) + quantize_model(self, transformer) transformer.to(self.device_torch) else: transformer.to(self.device_torch, dtype=dtype) @@ -373,19 +349,16 @@ class ChromaRadianceModel(BaseModel): return False def save_model(self, output_path, meta, save_dtype): + # comfy-format single-file save via the mixin (chroma's class keys ARE + # the original layout); handles torchao/Ostris dequant, not just quanto if not output_path.endswith(".safetensors"): - output_path = output_path + ".safetensors" - # only save the unet + output_path = output_path + ".safetensors" transformer: Chroma = unwrap_model(self.model) - state_dict = transformer.state_dict() - save_dict = {} - for k, v in state_dict.items(): - if isinstance(v, QTensor): - v = v.dequantize() - save_dict[k] = v.clone().to('cpu', dtype=save_dtype) - - meta = get_meta_for_safetensors(meta, name='chroma') - save_file(save_dict, output_path, metadata=meta) + transformer.save_model( + output_path, + dtype=save_dtype, + metadata=get_meta_for_safetensors(meta, name="chroma"), + ) def get_loss_target(self, *args, **kwargs): noise = kwargs.get('noise') diff --git a/extensions_built_in/diffusion_models/f_light/f_light.py b/extensions_built_in/diffusion_models/f_light/f_light.py index 2f3795c..018b97b 100644 --- a/extensions_built_in/diffusion_models/f_light/f_light.py +++ b/extensions_built_in/diffusion_models/f_light/f_light.py @@ -11,10 +11,9 @@ from toolkit.models.v2.vae.autoencoder_kl import KLVAE from toolkit.basic import flush from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler -from toolkit.dequantize import patch_dequantization_on_save from toolkit.accelerator import unwrap_model -from optimum.quanto import freeze, QTensor -from toolkit.util.quantize import quantize, get_qtype +from optimum.quanto import QTensor +from toolkit.util.quantize import quantize, quantize_model from .src import FLitePipeline, DiT if TYPE_CHECKING: @@ -34,6 +33,9 @@ scheduler_config = { class FLiteModel(BaseModel): arch = "f-lite" + def get_transformer_block_names(self): + return ["blocks"] + def __init__( self, device, @@ -79,13 +81,10 @@ class FLiteModel(BaseModel): transformer.to(self.quantize_device, dtype=dtype) if self.model_config.quantize: - # patch the state dict method - patch_dequantization_on_save(transformer) - quantization_type = get_qtype(self.model_config.qtype) + # block-streaming quantize (handles dequant-on-save patching, + # excludes, ARA, and quantize_kwargs) self.print_and_status_update("Quantizing transformer") - quantize(transformer, weights=quantization_type, - **self.model_config.quantize_kwargs) - freeze(transformer) + quantize_model(self, transformer) transformer.to(self.device_torch) else: transformer.to(self.device_torch, dtype=dtype) diff --git a/extensions_built_in/diffusion_models/flux2/flux2_model.py b/extensions_built_in/diffusion_models/flux2/flux2_model.py index 20aedf0..b61f119 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_model.py @@ -111,7 +111,7 @@ class Flux2Model(BaseModel): if self.model_config.quantize_te: self.print_and_status_update("Quantizing Mistral") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype)) + quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) freeze(text_encoder) flush() @@ -136,9 +136,6 @@ class Flux2Model(BaseModel): transformer_path = model_path self.print_and_status_update("Loading transformer") - with torch.device("meta"): - transformer = Flux2(self.get_flux2_params()) - # use local path if provided if os.path.exists(os.path.join(transformer_path, self.flux2_te_filename)): transformer_path = os.path.join(transformer_path, self.flux2_te_filename) @@ -152,12 +149,9 @@ class Flux2Model(BaseModel): ) transformer_state_dict = load_file(transformer_path, device="cpu") - - # cast to dtype - for key in transformer_state_dict: - transformer_state_dict[key] = transformer_state_dict[key].to(dtype) - - transformer.load_state_dict(transformer_state_dict, assign=True) + transformer = Flux2.load_from_state_dict( + transformer_state_dict, dtype, config=self.get_flux2_params() + ) if self.model_config.quantize: # patch the state dict method diff --git a/extensions_built_in/diffusion_models/flux2/src/model.py b/extensions_built_in/diffusion_models/flux2/src/model.py index cffae9f..5c04503 100644 --- a/extensions_built_in/diffusion_models/flux2/src/model.py +++ b/extensions_built_in/diffusion_models/flux2/src/model.py @@ -1,4 +1,6 @@ import torch + +from toolkit.models.v2._mixin import OstrisModelMixin from einops import rearrange from torch import Tensor, nn import torch.utils.checkpoint as ckpt @@ -54,7 +56,7 @@ class FakeConfig: self.patch_size = 1 -class Flux2(nn.Module): +class Flux2(nn.Module, OstrisModelMixin): def __init__(self, params: Flux2Params): super().__init__() self.config = FakeConfig() diff --git a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py index 927477f..535e9fb 100644 --- a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py +++ b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py @@ -17,11 +17,10 @@ from toolkit.basic import flush from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler from toolkit.models.flux import add_model_gpu_splitter_to_flux, bypass_flux_guidance, restore_flux_guidance -from toolkit.dequantize import patch_dequantization_on_save from toolkit.accelerator import get_accelerator, unwrap_model -from optimum.quanto import freeze, QTensor +from optimum.quanto import QTensor from toolkit.util.mask import generate_random_mask, random_dialate_mask -from toolkit.util.quantize import quantize, get_qtype +from toolkit.util.quantize import quantize, quantize_model from einops import rearrange, repeat import random import torch.nn.functional as F @@ -44,6 +43,9 @@ scheduler_config = { class FluxKontextModel(BaseModel): arch = "flux_kontext" + def get_transformer_block_names(self): + return ["transformer_blocks", "single_transformer_blocks"] + def __init__( self, device, @@ -95,13 +97,10 @@ class FluxKontextModel(BaseModel): transformer.to(self.quantize_device, dtype=dtype) if self.model_config.quantize: - # patch the state dict method - patch_dequantization_on_save(transformer) - quantization_type = get_qtype(self.model_config.qtype) + # block-streaming quantize (handles dequant-on-save patching, + # excludes, ARA, and quantize_kwargs) self.print_and_status_update("Quantizing transformer") - quantize(transformer, weights=quantization_type, - **self.model_config.quantize_kwargs) - freeze(transformer) + quantize_model(self, transformer) transformer.to(self.device_torch) else: transformer.to(self.device_torch, dtype=dtype) diff --git a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py index 2048bdd..d84c6ab 100644 --- a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py +++ b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py @@ -238,9 +238,6 @@ class Ideogram4Model(BaseModel): self.print_and_status_update("Loading transformer") transformer_config = Ideogram4Config() - with torch.device("meta"): - transformer = Ideogram4Transformer2DModel(transformer_config) - self.print_and_status_update(" - fetching transformer weights") state_dict = _load_component_state_dict( base, "transformer", "diffusion_pytorch_model" @@ -250,7 +247,9 @@ class Ideogram4Model(BaseModel): state_dict, dtype, self.device_torch, self.model_config.low_vram ) self.print_and_status_update(" - loading transformer state dict") - transformer.load_state_dict(state_dict, assign=True) + transformer = Ideogram4Transformer2DModel.load_from_state_dict( + state_dict, dtype, config=transformer_config + ) del state_dict flush() diff --git a/extensions_built_in/diffusion_models/ideogram4/src/transformer.py b/extensions_built_in/diffusion_models/ideogram4/src/transformer.py index 3173bda..4abc5c8 100644 --- a/extensions_built_in/diffusion_models/ideogram4/src/transformer.py +++ b/extensions_built_in/diffusion_models/ideogram4/src/transformer.py @@ -12,6 +12,8 @@ import math from dataclasses import dataclass import torch + +from toolkit.models.v2._mixin import OstrisModelMixin import torch.nn as nn import torch.nn.functional as F from torch.utils.checkpoint import checkpoint @@ -355,7 +357,7 @@ class Ideogram4FinalLayer(nn.Module): return self.linear(self.norm_final(x) * scale) -class Ideogram4Transformer2DModel(nn.Module): +class Ideogram4Transformer2DModel(nn.Module, OstrisModelMixin): """Ideogram 4 flow-matching transformer.""" def __init__(self, config: Ideogram4Config) -> None: diff --git a/extensions_built_in/diffusion_models/krea2/krea2.py b/extensions_built_in/diffusion_models/krea2/krea2.py index b0fa0a6..90462d9 100644 --- a/extensions_built_in/diffusion_models/krea2/krea2.py +++ b/extensions_built_in/diffusion_models/krea2/krea2.py @@ -217,21 +217,15 @@ class Krea2Model(QwenImageVAEHolderMixin, BaseModel): mmdit_kwargs.update(self.model_config.model_kwargs.get("mmdit_config", {})) config = SingleMMDiTConfig(**mmdit_kwargs) - # Build on meta, then materialize straight from the checkpoint. - with torch.device("meta"): - transformer = SingleStreamDiT(config) - self.print_and_status_update(" - fetching transformer weights") state_dict = _load_mmdit_state_dict( self.model_config.name_or_path, self.model_config.model_kwargs.get("checkpoint_filename", None), ) - state_dict = { - k: (v.to(dtype) if v.is_floating_point() else v) - for k, v in state_dict.items() - } self.print_and_status_update(" - loading transformer state dict") - transformer.load_state_dict(state_dict, strict=True, assign=True) + transformer = SingleStreamDiT.load_from_state_dict( + state_dict, dtype, config=config + ) del state_dict flush() return transformer diff --git a/extensions_built_in/diffusion_models/krea2/src/mmdit.py b/extensions_built_in/diffusion_models/krea2/src/mmdit.py index bb6619a..5e7cfcb 100644 --- a/extensions_built_in/diffusion_models/krea2/src/mmdit.py +++ b/extensions_built_in/diffusion_models/krea2/src/mmdit.py @@ -20,6 +20,8 @@ import math from dataclasses import dataclass import torch + +from toolkit.models.v2._mixin import OstrisModelMixin import torch.nn as nn import torch.nn.functional as F from einops import rearrange @@ -408,7 +410,7 @@ class SingleStreamBlock(nn.Module): return x -class SingleStreamDiT(nn.Module): +class SingleStreamDiT(nn.Module, OstrisModelMixin): def __init__(self, config: SingleMMDiTConfig): super().__init__() self.config = config diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index 2821e23..ddd6f3f 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -1279,7 +1279,7 @@ class LTX25Model(LTX2Model): from the state dict. Works unchanged for bf16 checkpoints, where no quant markers exist and everything strict-loads.""" from toolkit.util.comfy_quant_import import import_comfy_quantized_layers - from toolkit.util.ostris_quant import OstrisLinear + from toolkit.models.v2._mixin import OstrisModelMixin state_dict, num_quantized = import_comfy_quantized_layers( module, state_dict, orig_dtype=self.torch_dtype @@ -1288,32 +1288,8 @@ class LTX25Model(LTX2Model): self.print_and_status_update( f" - attached {num_quantized} pre-quantized ConvRot layers to {name}" ) - result = module.load_state_dict(state_dict, assign=True, strict=False) - # quantized linears hold their weight as backend buffers and had their - # bias assigned by the importer, so both report as "missing" here - quantized_param_keys = set() - for mod_name, m in module.named_modules(): - if isinstance(m, OstrisLinear): - quantized_param_keys.add(f"{mod_name}.weight") - if m.bias is not None: - quantized_param_keys.add(f"{mod_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"LTX-2.5 {name} load mismatch: missing {bad_missing[:8]}, " - f"unexpected {result.unexpected_keys[:8]}" - ) - # nothing may be left on the meta device (e.g. a bias the importer - # should have filled) - leftover_meta = [ - param_name - for param_name, p in module.named_parameters() - if p.is_meta - ] - if leftover_meta: - raise ValueError( - f"LTX-2.5 {name} load left meta parameters: {leftover_meta[:8]}" - ) + # whitelist for quantized weights + leftover-meta check + OstrisModelMixin._load_state_dict_with_quantized(module, state_dict) return num_quantized def _load_gemma4_text_encoder(self, te_path: str, te_state_dict: dict): diff --git a/extensions_built_in/diffusion_models/mageflow/mageflow.py b/extensions_built_in/diffusion_models/mageflow/mageflow.py index 49b3530..1f11192 100644 --- a/extensions_built_in/diffusion_models/mageflow/mageflow.py +++ b/extensions_built_in/diffusion_models/mageflow/mageflow.py @@ -178,22 +178,13 @@ class MageFlowModel(BaseModel): structure.update(self.model_config.model_kwargs.get("transformer_config", {})) params = MageFlowParams(**structure) - # Build on meta, then materialize straight from the checkpoint. - with torch.device("meta"): - transformer = MageFlow(params) - self.print_and_status_update(" - fetching transformer weights") state_dict = load_file( self._get_model_file("transformer/diffusion_pytorch_model.safetensors") ) - state_dict = { - k: (v.to(dtype) if v.is_floating_point() else v) - for k, v in state_dict.items() - } self.print_and_status_update(" - loading transformer state dict") - transformer.load_state_dict(state_dict, strict=True, assign=True) - # The RoPE tables are plain tensors (not buffers/params), so the meta - # init left them unmaterialized — rebuild them for real. + transformer = MageFlow.load_from_state_dict(state_dict, dtype, config=params) + # rebuild the plain-tensor RoPE tables at full precision transformer.reset_rope() del state_dict flush() diff --git a/extensions_built_in/diffusion_models/mageflow/src/transformer.py b/extensions_built_in/diffusion_models/mageflow/src/transformer.py index 6b0d247..7ce1a76 100644 --- a/extensions_built_in/diffusion_models/mageflow/src/transformer.py +++ b/extensions_built_in/diffusion_models/mageflow/src/transformer.py @@ -28,6 +28,8 @@ from dataclasses import dataclass from typing import Any import torch + +from toolkit.models.v2._mixin import OstrisModelMixin import torch.nn as nn from torch import Tensor from torch.utils.checkpoint import checkpoint @@ -681,7 +683,7 @@ class MageFlowParams: patch_size: int = 1 -class MageFlow(nn.Module): +class MageFlow(nn.Module, OstrisModelMixin): def __init__(self, params: MageFlowParams): super().__init__() self.params = params diff --git a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py index ab0d7fa..05a0d4a 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py @@ -373,46 +373,12 @@ class MinimaxH3Model(BaseModel): self.invert_assistant_lora = False def _load_transformer(self) -> MiniMaxH3Transformer: - dtype = self.torch_dtype dit_path = self._resolve_comfy_file(self._dit_component()) self.print_and_status_update(f"Loading transformer from {dit_path}") - state_dict = load_file(dit_path) - - params = MiniMaxH3TransformerParams() - table = state_dict.get("adaln_t_table", None) - if table is not None: - # pruned checkpoint: factored timestep table instead of the MLP - params.adaln_t_table_size = table.shape[0] - params.time_embed_dim = table.shape[1] - - with torch.device("meta"): - transformer = MiniMaxH3Transformer(params) - - # attach the pre-quantized (int8 ConvRot) linears onto the toolkit's - # quantization backends; the rest loads at its stored precision (the - # checkpoint's bf16/fp16/fp32 mix is deliberate) - state_dict, num_quantized = import_comfy_quantized_layers( - transformer, state_dict, orig_dtype=dtype - ) - if num_quantized: - self.print_and_status_update( - f" - attached {num_quantized} pre-quantized ConvRot layers" - ) - result = transformer.load_state_dict(state_dict, assign=True, strict=False) - quantized_weight_keys = { - f"{name}.weight" - for name, m in transformer.named_modules() - if isinstance(m, OstrisLinear) - } - bad_missing = [k for k in result.missing_keys if k not in quantized_weight_keys] - if bad_missing or result.unexpected_keys: - raise ValueError( - f"MiniMax-H3 transformer load mismatch: missing {bad_missing[:8]}, " - f"unexpected {result.unexpected_keys[:8]}" - ) - del state_dict - flush() - return transformer + # the mixin single-file path: config sniffed from the checkpoint + # (adaln_t_table), pre-quantized ConvRot linears attached, everything + # else at its stored precision (the bf16/fp16/fp32 mix is deliberate) + return MiniMaxH3Transformer.load_model(dit_path, dtype=self.torch_dtype) def _load_text_encoder(self): from accelerate import init_empty_weights diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py b/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py index f958268..809ff03 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py +++ b/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py @@ -30,6 +30,8 @@ import math from dataclasses import dataclass from typing import Optional, Tuple +from toolkit.models.v2._mixin import OstrisModelMixin + import torch import torch.nn.functional as F from torch import nn @@ -356,7 +358,31 @@ class MiniMaxH3FinalLayer(nn.Module): return self.video_out(h), self.audio_out(h) -class MiniMaxH3Transformer(nn.Module): +class MiniMaxH3Transformer(nn.Module, OstrisModelMixin): + # comfy checkpoints carry a deliberate bf16/fp16/fp32 mix + aitk_cast_on_load = False + + @classmethod + def aitk_config_from_state_dict(cls, state_dict): + params = MiniMaxH3TransformerParams() + table = state_dict.get("adaln_t_table", None) + if table is not None: + # pruned checkpoint: factored timestep table instead of the MLP + params.adaln_t_table_size = table.shape[0] + params.time_embed_dim = table.shape[1] + return params + + @classmethod + def aitk_from_config(cls, config): + from accelerate import init_empty_weights + + with init_empty_weights(include_buffers=False): + return cls(config) + + @classmethod + def get_transformer_block_names(cls): + return ["blocks"] + def __init__(self, params: Optional[MiniMaxH3TransformerParams] = None): super().__init__() if params is None: diff --git a/extensions_built_in/diffusion_models/omnigen2/__init__.py b/extensions_built_in/diffusion_models/omnigen2/__init__.py index 77741be..09b00f9 100644 --- a/extensions_built_in/diffusion_models/omnigen2/__init__.py +++ b/extensions_built_in/diffusion_models/omnigen2/__init__.py @@ -13,7 +13,7 @@ from toolkit.samplers.custom_flowmatch_sampler import ( ) from toolkit.accelerator import unwrap_model from optimum.quanto import freeze -from toolkit.util.quantize import quantize, get_qtype +from toolkit.util.quantize import quantize, get_qtype, quantize_model from .src.pipelines.omnigen2.pipeline_omnigen2 import OmniGen2Pipeline from .src.models.transformers import OmniGen2Transformer2DModel from .src.models.transformers.repo import OmniGen2RotaryPosEmbed @@ -105,9 +105,9 @@ class OmniGen2Model(BaseModel): if self.model_config.quantize: self.print_and_status_update("Quantizing transformer") - quantization_type = get_qtype(self.model_config.qtype) - quantize(transformer, weights=quantization_type) - freeze(transformer) + quantize_model(self, transformer) + if not self.low_vram: + transformer.to(self.device_torch) if self.low_vram: # unload it for now diff --git a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py index 3a26897..97ccd52 100644 --- a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py +++ b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py @@ -22,6 +22,7 @@ from toolkit.models.v2.diffusion_models.wan import WanTransformer3DModel from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO from torchvision.transforms import functional as TF +from toolkit.metadata import get_meta_for_safetensors from toolkit.models.wan21.wan21 import Wan21 from .wan22_5b_model import ( scheduler_config, @@ -431,20 +432,20 @@ class Wan2214bModel(Wan21): return False def save_model(self, output_path, meta, save_dtype): + # comfy-format single-file saves, one per DiT (comfy convention: + # separate high/low noise files) transformer_combo: DualWanTransformer3DModel = unwrap_model(self.model) - transformer_combo.transformer_1.save_pretrained( - save_directory=os.path.join(output_path, "transformer"), - safe_serialization=True, + base = output_path + if base.endswith(".safetensors"): + base = base[: -len(".safetensors")] + metadata = get_meta_for_safetensors(meta, name=self.arch) + transformer_combo.transformer_1.save_model( + f"{base}_high_noise.safetensors", dtype=save_dtype, metadata=metadata ) - transformer_combo.transformer_2.save_pretrained( - save_directory=os.path.join(output_path, "transformer_2"), - safe_serialization=True, + transformer_combo.transformer_2.save_model( + f"{base}_low_noise.safetensors", dtype=save_dtype, metadata=metadata ) - meta_path = os.path.join(output_path, "aitk_meta.yaml") - with open(meta_path, "w") as f: - yaml.dump(meta, f) - def save_lora( self, state_dict: Dict[str, torch.Tensor], diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py index b4a3bb3..d9af932 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py @@ -92,10 +92,6 @@ class ZetaChromaModel(BaseModel): transformer_state_dict = load_file(transformer_path, device="cpu") - # cast to dtype - for key in transformer_state_dict: - transformer_state_dict[key] = transformer_state_dict[key].to(dtype) - # Auto-detect use_x0 from checkpoint use_x0 = "__x0__" in transformer_state_dict @@ -107,10 +103,9 @@ class ZetaChromaModel(BaseModel): use_x0=use_x0, ) - with torch.device("meta"): - transformer = ZImageDCT(model_params) - - transformer.load_state_dict(transformer_state_dict, assign=True) + transformer = ZImageDCT.load_from_state_dict( + transformer_state_dict, dtype, config=model_params + ) del transformer_state_dict transformer.to(self.quantize_device, dtype=dtype) diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py index e2db51d..03995b1 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py @@ -6,6 +6,8 @@ from typing import List, Optional import math import torch + +from toolkit.models.v2._mixin import OstrisModelMixin import torch.nn as nn import torch.nn.functional as F from torch import Tensor @@ -449,7 +451,7 @@ class SimpleMLPAdaLN(nn.Module): return self.final_layer(x) -class ZImageDCT(nn.Module): +class ZImageDCT(nn.Module, OstrisModelMixin): def __init__(self, params: ZImageDCTParams): super().__init__() self.config = FakeConfig() diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index 7b17bef..74096aa 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -207,8 +207,11 @@ loads via diffusers. Nothing about sources or outputs changes yet. Suggested ord variants inherit - [x] nucleus_image — `v2/diffusion_models/nucleus_image.py`, TE stanza collapsed to prepare_text_encoder -- [ ] krea2, ideogram4, mageflow — TE/VAE migrated; their custom local DiT - classes still to be rebased onto the mixin +- [x] krea2, ideogram4, mageflow — DiTs on the mixin via the `config=` + passthrough (holder builds config from model_kwargs/holder state, mixin + does build/markers/whitelist/casting; plain nn.Module classes now work + with the default builder). krea2 + ideogram4 verified by harness; + mageflow untestable while its repo 404s - [x] chroma, chroma_radiance — both vendored Chroma classes now carry `OstrisModelMixin` with the block-count sniff moved into a new `aitk_config_from_state_dict` hook (mixin now supports checkpoint-derived @@ -217,14 +220,23 @@ loads via diffusers. Nothing about sources or outputs changes yet. Suggested ord depends on holder state (patch_size), not the checkpoint - [x] flux_kontext — `v2/diffusion_models/flux.py` (FluxTransformer2DModel); whole model now loads through v2 (transformer, T5, CLIP, KLVAE) -- [ ] flux2 — TE/VAE partially migrated (flux2_kl); custom DiT still local. - krea2/mageflow/ideogram4/zeta_chroma DiTs stay model-specific: their - configs come from model_kwargs / holder state, so the mixin adds nothing - until the comfy-weights flip (Phase 2) -- [ ] minimax_h3 (+ ref2va), ltx2 family — already on the shared resolver + - comfy_quant_import; the full mixin port waits for Phase 2, when the - mixin's single-file precision policy (stored-precision loading, fp32-key - protection) is settled to match their deliberate behavior +- [x] flux2 + zeta_chroma DiTs on the mixin via `config=` passthrough + (flux2_klein_4b verified by harness). Every arch's DiT now loads + through the mixin except: ltx2 family (one-file→two-modules split, + helper delegated), anima (diffusers modular pipeline), ace_step + (bundled single-file loader), z_image_l2p's local subclass, and the + grandfathered legacy stable_diffusion_model archs +- [x] minimax_h3 (+ ref2va) transformer ported to the mixin: config sniffed + from the checkpoint via aitk_config_from_state_dict (adaln_t_table), + marker attach + stored-precision load via the new + `aitk_cast_on_load = False` knob. Verified on the real pruned convrot + file: 200 ConvRot linears, pruned table detected, fp32/fp16/bf16 mix + preserved, no meta leftovers. Its TE stays custom (50-layer truncation + + key_map). ltx2.5's `_load_quantized_module` now delegates to the + mixin's whitelist/meta helper (~30 lines deleted); its full port is + blocked on the one-comfy-file → transformer+connectors split, which + doesn't fit the per-class single-file shape — revisit with the live + server's component model - [x] wan21 / wan22 family — `v2/diffusion_models/wan.py` (WanTransformer3DModel, both wan22 dual loads included) + `v2/text_encoders/umt5.py` (UMT5TextEncoder + PatchedT5Tokenizer; @@ -307,9 +319,16 @@ Decisions: fp8 weight / scalar scale_weight, e.g. every wan *_fp8_scaled file): imports onto the float8 backend; scale_input (activation quant) is dropped, matmuls run dequantized. +- [x] wan comfy-format saves: `convert_state_dict_on_save` inverts diffusers' + rename table (base/t2v/i2v; vace/animate excluded — their reverse + mappings collide). Round-trip verified on both real comfy files + (exact; the 2.1 file's legacy model.diffusion_model. prefix drops per + the modern convention) and with real weights (load → save → 825-key + original-layout file → reload bit-equal). wan21 + wan22_5b save one + comfy file; wan22_14b saves the comfy-standard _high_noise/_low_noise + pair instead of two diffusers folders. - [ ] Wire remaining archs' candidate lists (chroma/others as their key - conversions are verified per file); wan comfy-format save needs the - inverse key mapping (defer with the other save flips) + conversions are verified per file) - [x] Fused-layout quantized attach for diffusers-split classes: `split_fused_quantized_keys` / `fuse_split_quantized_keys` (comfy_quant_import) do exact out-dim row surgery on quantized comfy @@ -333,6 +352,14 @@ Decisions: bit-exact fused qkv weights/scales/markers, and the reload's quantized forward is bit-identical — toolkit saves are byte-compatible with ComfyUI. +- [x] chroma + chroma_radiance saves flipped to the mixin (their class keys + ARE the original layout) — also fixes their quanto-only dequant bug + (torchao/Ostris weights now dequantize on save). Tiny-model round trip + verified. Save flips so far: z_image, qwen_image, wan21, wan22_5b, + wan22_14b (dual files), chroma ×2. +- [ ] flux_kontext comfy wiring deferred: its Comfy-Org repo ships a single + legacy-fp8 file in fused BFL layout — needs the flux fused-split + conversion (split_fused_quantized_keys pattern + BFL↔diffusers maps) - [ ] Flip the remaining per-arch `save_model` overrides as each arch's save-side key conversion is in place - [ ] Publish/verify comfy repacks per model as they flip @@ -354,6 +381,13 @@ Decisions: - [x] Missing weights skip rather than fail: default is HF_HUB_OFFLINE=1 and hub/file errors classify as SKIP; `--allow-download` opts into fetching. (GPU + local-weights test, not CI-portable.) +- [x] Final certification sweep (2026-08-27, post-polish): 14/14 runnable + archs PASS — comfy-source loads (zimage convrot8, qwen fp8, wan ×2 + + fp8 umt5 TE), all ported holder-config DiTs, and the migrated + quantize_model paths (chroma, flux_kontext, f_light block-streamed) in + one run; mageflow remains the upstream 404 skip. One regression caught + and fixed: qwen's _load_single_file override needed the new config + kwarg. - [x] Full sweep run 2026-08-27: 14/15 PASS (zimage, qwen_image, krea2, boogu_image, ernie_image, ideogram4, hidream_o1, anima, wan21, wan22_5b, chroma, flux_kontext, flux2_klein_4b, ltx2.3 — the quantized 22B ltx @@ -378,20 +412,28 @@ Decisions: ## TODO / look at later -- [ ] Quantize-path consolidation quirks: `quantize_kwargs` is honored only by the - raw `quantize()` call sites and silently dropped by `quantize_model()`; the - ARA path inside `quantize_model` hardcodes `uint8`. Decide the unified - behavior when consolidating. +- [x] Quantize consolidation: `quantize_model` now honors `quantize_kwargs` + (blocks + extras) and tolerates missing block names; the chroma ×2, + flux_kontext, f_light, omnigen2 raw-quantize sites migrated onto it + (gaining block streaming, excludes, ARA, dequant-on-save patching) with + holder block names added. Remaining raw sites are legacy/extension + (flex2, cogview4, stable_diffusion_model). The ARA uint8 hardcode + stands — revisit if a non-uint8 ARA base is ever wanted. +- [x] Last known `qtype_te` bugs fixed (flux2's Mistral TE, anima's + text_conditioner) — 9/9 sites from the survey now correct outside the + grandfathered legacy monolith (cogview4/legacy SD remain as-is). - [x] wan comfy-TE resolved for real: UMT5TextEncoder carries comfy candidates (fp8_e4m3fn_scaled via the legacy importer, fp16), files already in transformers key layout (spiece blob dropped, tied embed_tokens materialized). Verified: wan21 samples with the local comfy fp8 TE. The loaders/umt5.py `comfy_files` param stays as a no-op shim for old callers. -- [ ] Registry hardening: error (don't fall back to SD1) on unknown arch; lazy +- [x] Registry hardening: unknown archs now raise with the known-arch list + (legacy monolith archs whitelisted via LEGACY_ARCHS). Still open: lazy per-arch imports; single source of truth shared with the UI's `options.tsx` model list. -- [ ] Fake/stub components: consolidate on `toolkit/models/FakeVAE.py` / - `toolkit/unloader.py`, delete local copies. +- [x] Stub dedup where identical: chroma_radiance imports FakeCLIP/FakeConfig + from chroma_model. The other Fake* copies (hidream_o1, flux2, zeta) + carry model-specific values — left in place. - [ ] Vendored upstream code (hidream/src, omnigen2/src, ltx2 converter's private comfy-quant parser): dedupe against toolkit utils where practical. diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index 33b9fa9..e4d58de 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -67,6 +67,12 @@ class OstrisModelMixin: 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 + # ---- state set by the loader / quantizer ---- aitk_is_quantized: bool = False aitk_qtype: Optional[str] = None @@ -119,7 +125,10 @@ class OstrisModelMixin: from accelerate import init_empty_weights with init_empty_weights(include_buffers=False): - return cls.from_config(config) + if hasattr(cls, "from_config"): + return cls.from_config(config) + # plain nn.Module classes take their params/config object directly + return cls(config) # ------------------------------------------------------------------ # loading @@ -135,6 +144,7 @@ class OstrisModelMixin: 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, @@ -180,7 +190,11 @@ class OstrisModelMixin: 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 + file_path, + dtype=dtype, + config_path=config_path, + config=config, + subfolder=subfolder, ) else: if os.path.isdir(name_or_path): @@ -243,11 +257,16 @@ class OstrisModelMixin: file_path: str, dtype: torch.dtype, config_path: Optional[str] = None, + config=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 + state_dict, + dtype, + config_path=config_path, + config=config, + subfolder=subfolder, ) @classmethod @@ -263,16 +282,20 @@ class OstrisModelMixin: state_dict: Dict[str, torch.Tensor], dtype: torch.dtype, config_path: Optional[str] = None, + config=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).""" + 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 ) - config = cls.aitk_config_from_state_dict(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) @@ -288,11 +311,15 @@ class OstrisModelMixin: ) cls._load_state_dict_with_quantized(model, state_dict) model.aitk_is_quantized = True - else: + elif cls.aitk_cast_on_load: for key, value in state_dict.items(): - state_dict[key] = value.to(dtype=dtype) + 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 diff --git a/toolkit/models/v2/diffusion_models/qwen_image.py b/toolkit/models/v2/diffusion_models/qwen_image.py index 3c4e689..56e3f3a 100644 --- a/toolkit/models/v2/diffusion_models/qwen_image.py +++ b/toolkit/models/v2/diffusion_models/qwen_image.py @@ -29,7 +29,9 @@ class QwenImageTransformer2DModel( return ["transformer_blocks"] @classmethod - def _load_single_file(cls, file_path, dtype, config_path=None, subfolder=None): + def _load_single_file( + cls, file_path, dtype, config_path=None, config=None, subfolder=None + ): from safetensors import safe_open with safe_open(file_path, framework="pt") as f: @@ -38,7 +40,11 @@ class QwenImageTransformer2DModel( # comfy prequantized checkpoint (diffusers key layout): the mixin # path attaches the quantized layers return super()._load_single_file( - file_path, dtype, config_path=config_path, subfolder=subfolder + file_path, + dtype, + config_path=config_path, + config=config, + subfolder=subfolder, ) # other single-file checkpoints carry diffusers or original key # layouts; diffusers' single-file machinery owns that conversion diff --git a/toolkit/models/v2/diffusion_models/wan.py b/toolkit/models/v2/diffusion_models/wan.py index 98ddd1c..dfa5958 100644 --- a/toolkit/models/v2/diffusion_models/wan.py +++ b/toolkit/models/v2/diffusion_models/wan.py @@ -84,3 +84,54 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin): ) return convert_wan_transformer_to_diffusers(dict(state_dict)) + + # inverse of diffusers' rename table, for the base/t2v/i2v variants (no + # vace/animate: their extra mappings collide in reverse). Order matters: + # longest/most-specific first, and the norm2/norm3 swap uses a placeholder. + _SAVE_RENAMES = [ + ("condition_embedder.time_embedder.linear_1", "time_embedding.0"), + ("condition_embedder.time_embedder.linear_2", "time_embedding.2"), + ("condition_embedder.text_embedder.linear_1", "text_embedding.0"), + ("condition_embedder.text_embedder.linear_2", "text_embedding.2"), + ("condition_embedder.time_proj", "time_projection.1"), + ("condition_embedder.image_embedder.norm1", "img_emb.proj.0"), + ("condition_embedder.image_embedder.ff.net.0.proj", "img_emb.proj.1"), + ("condition_embedder.image_embedder.ff.net.2", "img_emb.proj.3"), + ("condition_embedder.image_embedder.norm2", "img_emb.proj.4"), + ("ffn.net.0.proj", "ffn.0"), + ("ffn.net.2", "ffn.2"), + (".norm_added_k.", ".norm_k_img."), + (".add_k_proj.", ".k_img."), + (".add_v_proj.", ".v_img."), + (".to_out.0.", ".o."), + (".to_q.", ".q."), + (".to_k.", ".k."), + (".to_v.", ".v."), + ("attn2", "cross_attn"), + ("attn1", "self_attn"), + # norm2 <-> norm3 swap back + ("norm3", "norm__placeholder"), + ("norm2", "norm3"), + ("norm__placeholder", "norm2"), + ("proj_out", "head.head"), + ] + + @classmethod + def convert_state_dict_on_save(cls, state_dict): + """Diffusers layout back to the original/comfy wan key layout + (rename-only, so quantized weight/scale/marker keys ride along).""" + if any(".self_attn." in k or ".cross_attn." in k for k in state_dict): + return state_dict # already original layout + new_sd = {} + for key, value in state_dict.items(): + k = key + # scale_shift_table: blocks keep the name as `modulation`, the + # top-level one belongs to the output head + if k == "scale_shift_table": + k = "head.modulation" + elif k.endswith(".scale_shift_table"): + k = k[: -len("scale_shift_table")] + "modulation" + for src, dst in cls._SAVE_RENAMES: + k = k.replace(src, dst) + new_sd[k] = value + return new_sd diff --git a/toolkit/models/wan21/wan21.py b/toolkit/models/wan21/wan21.py index 1fec78d..5e07efc 100644 --- a/toolkit/models/wan21/wan21.py +++ b/toolkit/models/wan21/wan21.py @@ -44,6 +44,7 @@ from typing import Any, Callable, Dict, List, Optional, Union from toolkit.models.wan21.wan_lora_convert import convert_to_diffusers, convert_to_original from toolkit.util.quantize import quantize_model from toolkit.models.v2.text_encoders.umt5 import UMT5TextEncoder +from toolkit.metadata import get_meta_for_safetensors # for generation only? scheduler_configUniPC = { @@ -680,17 +681,16 @@ class Wan21(BaseModel): return False def save_model(self, output_path, meta, save_dtype): - # only save the unet - transformer: Wan21 = unwrap_model(self.model) - transformer.save_pretrained( - save_directory=os.path.join(output_path, 'transformer'), - safe_serialization=True, + # comfy-format single-file save (original wan key layout) + transformer = unwrap_model(self.model) + if not output_path.endswith(".safetensors"): + output_path += ".safetensors" + transformer.save_model( + output_path, + dtype=save_dtype, + metadata=get_meta_for_safetensors(meta, name=self.arch), ) - meta_path = os.path.join(output_path, 'aitk_meta.yaml') - with open(meta_path, 'w') as f: - yaml.dump(meta, f) - def get_loss_target(self, *args, **kwargs): noise = kwargs.get('noise') batch = kwargs.get('batch') diff --git a/toolkit/util/get_model.py b/toolkit/util/get_model.py index 545175d..3cd9e79 100644 --- a/toolkit/util/get_model.py +++ b/toolkit/util/get_model.py @@ -41,10 +41,34 @@ def get_all_models() -> List[BaseModel]: return all_model_classes +# archs the legacy StableDiffusion monolith still serves (see the arch +# normalization in toolkit/config_modules.py) +LEGACY_ARCHS = { + "sd1", + "sd2", + "sd3", + "sdxl", + "pixart", + "pixart_sigma", + "auraflow", + "flux", + "lumina2", + "vega", + "ssd", +} + + def get_model_class(config: ModelConfig): all_models = get_all_models() for ModelClass in all_models: if ModelClass.arch == config.arch: return ModelClass - # default to the legacy model - return StableDiffusion + if config.arch in LEGACY_ARCHS: + return StableDiffusion + # a typo'd or unregistered arch used to silently fall back to SD1; error + # instead (a broken extension import also lands here — its error was + # printed during get_all_models) + known = sorted({m.arch for m in all_models if m.arch} | LEGACY_ARCHS) + raise ValueError( + f"Unknown model arch {config.arch!r}. Known archs: {', '.join(known)}" + ) diff --git a/toolkit/util/quantize.py b/toolkit/util/quantize.py index b187616..db54f58 100644 --- a/toolkit/util/quantize.py +++ b/toolkit/util/quantize.py @@ -429,9 +429,10 @@ def quantize_model( # quantize model the original way without an accuracy recovery adapter # move and quantize only certain pieces at a time. quantization_type = get_qtype(base_model.model_config.qtype) + quantize_kwargs = base_model.model_config.quantize_kwargs or {} # all_blocks = list(model_to_quantize.transformer_blocks) all_blocks: List[torch.nn.Module] = [] - transformer_block_names = base_model.get_transformer_block_names() + transformer_block_names = base_model.get_transformer_block_names() or [] for name in transformer_block_names: # name may be a dotted path for models that nest their blocks # (e.g. hidream_o1's "model.language_model.layers"). @@ -456,7 +457,12 @@ def quantize_model( block.to(base_model.device_torch, dtype=base_model.torch_dtype, non_blocking=True) # exclude patterns with a leading wildcard (e.g. "*adaln_proj*") # also apply inside blocks, where names are block-relative - quantize(block, weights=quantization_type, exclude=exclude_modules) + quantize( + block, + weights=quantization_type, + exclude=exclude_modules, + **quantize_kwargs, + ) freeze(block) # NOT non_blocking: an async D2H allocates the cpu destination in pinned # memory, which the caching host allocator keeps forever (with power-of-2 @@ -472,5 +478,10 @@ def quantize_model( # device without having to move the transformer blocks to the device first base_model.print_and_status_update(" - quantizing extras") # model_to_quantize.to(base_model.device_torch, dtype=base_model.torch_dtype) - quantize(model_to_quantize, weights=quantization_type, exclude=exclude_modules) + quantize( + model_to_quantize, + weights=quantization_type, + exclude=exclude_modules, + **quantize_kwargs, + ) freeze(model_to_quantize) From c3bc8b0b4ea103908a4ef2a266f20d776861b21d Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 28 Aug 2026 06:56:05 -0600 Subject: [PATCH 07/14] Full test of ui arch compatability --- .../diffusion_models/ltx2/ltx2.py | 11 +-- testing/test_model_loading.py | 71 +++++++++++++++++++ toolkit/models/v2/PLANNING.md | 14 ++++ toolkit/models/v2/diffusion_models/wan.py | 13 ++++ 4 files changed, 105 insertions(+), 4 deletions(-) diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index ddd6f3f..94ca300 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -1374,8 +1374,12 @@ class LTX25Model(LTX2Model): transformer, transformer_sd, "transformer" ) del transformer_sd - if num_quantized_dit == 0: - transformer = transformer.to(dtype) + # cast to the compute dtype even when pre-quantized: the comfy file + # stores the scale_shift modulation tables in fp32, which promotes + # hidden states to fp32 and breaks the (bf16) diffusers attention + # linears. Quantized backends are immune (int8 qdata, uint8-viewed + # scales survive a dtype cast untouched). + transformer = transformer.to(dtype) flush() if self.model_config.quantize: @@ -1420,8 +1424,7 @@ class LTX25Model(LTX2Model): connectors, connectors_sd, "connectors" ) del connectors_sd, dit_sd - if num_quantized_connectors == 0: - connectors = connectors.to(dtype) + connectors = connectors.to(dtype) flush() # ---- text encoder (Gemma-4 12B, single comfy file) ---- diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index 9ad0882..77ae4ca 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -107,6 +107,77 @@ MODEL_TESTS = { "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-4B", "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0}, }, + # ---- coverage for every UI-default arch (big downloads on first run) ---- + "wan22_14b": { + "model": {"name_or_path": "ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True}, + "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 3.5, "seed": 42, "num_frames": 17}, + }, + "wan22_14b_i2v": { + "model": {"name_or_path": "ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True}, + "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 3.5, "seed": 42, "num_frames": 17}, + "needs_control_image": True, + }, + "wan21_i2v": { + "model": {"name_or_path": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", "quantize": True, "quantize_te": True}, + "sample": {"width": 480, "height": 480, "num_inference_steps": 20, "guidance_scale": 5.0, "seed": 42, "num_frames": 17}, + "needs_control_image": True, + }, + "hidream": { + "model": {"name_or_path": "HiDream-ai/HiDream-I1-Full", "quantize": True, "quantize_te": True}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 28, "guidance_scale": 5.0, "seed": 42}, + }, + "hidream_e1": { + # editing seq budget fits 768x768 (source+target concat) + "model": {"name_or_path": "HiDream-ai/HiDream-E1-1", "quantize": True, "quantize_te": True}, + "sample": {"width": 768, "height": 768, "num_inference_steps": 28, "guidance_scale": 5.0, "seed": 42}, + "needs_control_image": True, + }, + "nucleus_image": { + "model": {"name_or_path": "NucleusAI/Nucleus-Image", "quantize": True, "quantize_te": True}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + }, + "omnigen2": { + "model": {"name_or_path": "OmniGen2/OmniGen2", "quantize": True, "quantize_te": True}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + }, + "ltx2.5": { + "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 3.0, "seed": 42, "num_frames": 25}, + }, + "flux2": { + "model": {"name_or_path": "black-forest-labs/FLUX.2-dev", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0}, + }, + "flux2_klein_9b": { + "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-9B", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0}, + }, + "prx_pixel": { + "model": {"name_or_path": "Photoroom/prxpixel-t2i", "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0}, + }, + "zeta_chroma": { + "model": {"name_or_path": "lodestones/Zeta-Chroma/zeta-chroma-base-x0-pixel-dino-distance.safetensors", "extras_name_or_path": "Tongyi-MAI/Z-Image-Turbo", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 25, "guidance_scale": 4.0}, + }, + "zimage_l2p": { + "model": {"name_or_path": "zhen-nan/L2P/model-1k-merge.safetensors", "extras_name_or_path": "Tongyi-MAI/Z-Image-Turbo", "quantize_te": True}, + "sample": {**IMG, "guidance_scale": 1.0}, + }, + "qwen_image_edit": { + "model": {"name_or_path": "Qwen/Qwen-Image-Edit", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 4.0}, + "needs_control_image": True, + }, + "qwen_image_edit_plus": { + "model": {"name_or_path": "Qwen/Qwen-Image-Edit-2509", "quantize": True, "quantize_te": True}, + "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 4.0}, + "needs_control_image": True, + }, + "f-lite": { + "model": {"name_or_path": "Freepik/F-Lite", "quantize": True, "quantize_te": True}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + }, } SKIP_MARKERS = ( diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index 74096aa..3e943cc 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -408,6 +408,20 @@ Decisions: load → comfy save → identical key set + bit-exact quantized entries vs the published file → reload → bit-identical quantized forward. Extend per arch as saves flip. +- [x] Full UI-default coverage (2026-08-27): the registry now holds all 31 + UI-facing archs, and 30/31 load + generate through the new stack with + their real UI defaults (mageflow blocked upstream by its hub 404). + This round verified the previously-untested tail with real weights: + wan22_14b + i2v (comfy fp8 pairs via the legacy importer, candidate + keys added for the ai-toolkit bf16 default repos), wan21_i2v, hidream, + hidream_e1 (native 768² editing), nucleus, omnigen2, ltx2.5, flux2, + flux2_klein_9b, prx_pixel, zeta_chroma, zimage_l2p, both qwen edit + archs, f-lite. Fixes found by the run: ltx2.5's fp32 scale_shift + tables promoted hidden states into bf16 linears under the diffusers + class (ComfyUI casts per-op) — the DiT/connectors now cast to compute + dtype after quantized attach (ConvRot storage immune); harness configs + for e1 resolution and zeta/l2p extras_name_or_path corrected to match + the UI defaults. - [ ] Each newly migrated model adds its test in the same PR as its migration. ## TODO / look at later diff --git a/toolkit/models/v2/diffusion_models/wan.py b/toolkit/models/v2/diffusion_models/wan.py index dfa5958..67f62b6 100644 --- a/toolkit/models/v2/diffusion_models/wan.py +++ b/toolkit/models/v2/diffusion_models/wan.py @@ -30,6 +30,19 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin): ("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "transformer_2"): [ "split_files/diffusion_models/wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors", ], + # the UI defaults point at the ai-toolkit bf16 repacks; same comfy files + ("ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "transformer"): [ + "split_files/diffusion_models/wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors", + ], + ("ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "transformer_2"): [ + "split_files/diffusion_models/wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors", + ], + ("ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "transformer"): [ + "split_files/diffusion_models/wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors", + ], + ("ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "transformer_2"): [ + "split_files/diffusion_models/wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors", + ], "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": { "repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged", "files": [ From 702254688d585071a73959e0a892e1c7f9891dbd Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 28 Aug 2026 07:55:19 -0600 Subject: [PATCH 08/14] Added legacy paths --- .../audio_models/ace_step/src/model.py | 119 ++++++++++-------- .../diffusion_models/anima/anima.py | 38 ++++-- .../boogu_image/boogu_image.py | 10 +- .../ernie_image/ernie_image.py | 9 +- .../flux2/flux2_klein_model.py | 6 +- .../diffusion_models/flux2/flux2_model.py | 25 +--- .../diffusion_models/hidream/hidream_model.py | 6 +- .../hidream/hidream_o1_model.py | 5 +- .../diffusion_models/ideogram4/ideogram4.py | 8 +- .../ltx2/convert_ltx2_to_diffusers.py | 12 +- .../diffusion_models/ltx2/ltx2.py | 83 ++++++------ .../diffusion_models/mageflow/mageflow.py | 4 +- .../diffusion_models/mageflow/src/vae.py | 30 +++-- .../diffusion_models/minimax_h3/minimax_h3.py | 40 +----- .../minimax_h3/src/audio_vae.py | 23 +++- .../diffusion_models/minimax_h3/src/vae.py | 21 +++- .../diffusion_models/omnigen2/__init__.py | 12 +- .../prx_pixel_t2i/prx_pixel_t2i.py | 5 +- .../z_image/z_image_l2p_model.py | 2 +- testing/test_model_loading.py | 48 ++++++- toolkit/models/v2/PLANNING.md | 45 ++++++- toolkit/models/v2/_mixin.py | 22 ++++ .../models/v2/diffusion_models/auraflow.py | 13 ++ toolkit/models/v2/diffusion_models/cosmos.py | 15 +++ toolkit/models/v2/diffusion_models/ltx2.py | 34 +++++ toolkit/models/v2/diffusion_models/lumina2.py | 13 ++ toolkit/models/v2/diffusion_models/pixart.py | 24 ++++ toolkit/models/v2/diffusion_models/sd3.py | 11 ++ toolkit/models/v2/diffusion_models/unet.py | 13 ++ toolkit/models/v2/text_encoders/anima.py | 9 ++ toolkit/models/v2/text_encoders/gemma2.py | 14 +++ toolkit/models/v2/text_encoders/gemma3.py | 32 +++++ toolkit/models/v2/text_encoders/llama.py | 14 +++ toolkit/models/v2/text_encoders/mistral3.py | 25 ++++ toolkit/models/v2/text_encoders/qwen3.py | 13 +- toolkit/models/v2/text_encoders/qwen3_vl.py | 38 +++++- toolkit/models/v2/vae/autoencoder_kl_flux2.py | 10 ++ toolkit/models/v2/vae/flux2_kl.py | 11 +- toolkit/models/v2/vae/ltx2.py | 16 +++ toolkit/models/v2/vae/wan.py | 9 ++ .../models/v2/vision_encoders/clip_vision.py | 13 ++ toolkit/models/wan21/wan21.py | 7 +- toolkit/models/wan21/wan21_i2v.py | 13 +- toolkit/stable_diffusion_model.py | 50 ++++++++ toolkit/util/mixed_precision.py | 93 ++++++++++++++ 45 files changed, 840 insertions(+), 223 deletions(-) create mode 100644 toolkit/models/v2/diffusion_models/auraflow.py create mode 100644 toolkit/models/v2/diffusion_models/cosmos.py create mode 100644 toolkit/models/v2/diffusion_models/ltx2.py create mode 100644 toolkit/models/v2/diffusion_models/lumina2.py create mode 100644 toolkit/models/v2/diffusion_models/pixart.py create mode 100644 toolkit/models/v2/diffusion_models/sd3.py create mode 100644 toolkit/models/v2/diffusion_models/unet.py create mode 100644 toolkit/models/v2/text_encoders/anima.py create mode 100644 toolkit/models/v2/text_encoders/gemma2.py create mode 100644 toolkit/models/v2/text_encoders/gemma3.py create mode 100644 toolkit/models/v2/text_encoders/llama.py create mode 100644 toolkit/models/v2/text_encoders/mistral3.py create mode 100644 toolkit/models/v2/vae/autoencoder_kl_flux2.py create mode 100644 toolkit/models/v2/vae/ltx2.py create mode 100644 toolkit/models/v2/vae/wan.py create mode 100644 toolkit/models/v2/vision_encoders/clip_vision.py create mode 100644 toolkit/util/mixed_precision.py diff --git a/extensions_built_in/audio_models/ace_step/src/model.py b/extensions_built_in/audio_models/ace_step/src/model.py index 8bf9ffa..7199596 100644 --- a/extensions_built_in/audio_models/ace_step/src/model.py +++ b/extensions_built_in/audio_models/ace_step/src/model.py @@ -19,6 +19,8 @@ import os import time import torch + +from toolkit.models.v2._mixin import OstrisModelMixin import torch.nn.functional as F import torchaudio from safetensors.torch import load_file @@ -857,7 +859,35 @@ class DiTModel(nn.Module): # ── Top-level model ── -class AceStep15(nn.Module): +class AceStep15(nn.Module, OstrisModelMixin): + @classmethod + def load_from_state_dict(cls, state_dict, dtype=None, **kwargs): + cfg = infer_dit_config(state_dict) + model = cls( + hidden=cfg["hidden"], + inter=cfg["inter"], + heads=cfg["heads"], + kv=cfg["kv"], + head_dim=cfg["head_dim"], + n_dit=cfg["n_dit"], + n_lyric=cfg["n_lyric"], + n_timbre=cfg["n_timbre"], + enc_hidden=cfg["enc_hidden"], + enc_heads=cfg["enc_heads"], + enc_kv=cfg["enc_kv"], + enc_inter=cfg["enc_inter"], + ) + missing, unexpected = model.load_state_dict(state_dict, strict=False) + # tokenizer/detokenizer keys are expected to be unused (cover mode only) + unexpected = [ + k for k in unexpected if not k.startswith(("tokenizer.", "detokenizer.")) + ] + if missing: + print(f" DiT missing: {len(missing)} (first 3: {missing[:3]})") + if unexpected: + print(f" DiT unexpected: {len(unexpected)} (first 3: {unexpected[:3]})") + return model + def __init__( self, hidden=2048, @@ -1051,7 +1081,17 @@ class _SeqWrap(nn.Module): return self.layers(x) -class OobleckVAE(nn.Module): +class OobleckVAE(nn.Module, OstrisModelMixin): + @classmethod + def load_from_state_dict(cls, state_dict, dtype=None, **kwargs): + vae = cls() + m, u = vae.load_state_dict(state_dict, strict=False) + if m: + print(f" VAE missing: {len(m)} (first 3: {m[:3]})") + if u: + print(f" VAE unexpected: {len(u)}") + return vae + def __init__( self, in_ch=2, @@ -1152,7 +1192,28 @@ class OobleckVAE(nn.Module): # ═══════════════════════════════════════════════════════════════════════════════ -class TextEncoder(nn.Module): +class TextEncoder(nn.Module, OstrisModelMixin): + @classmethod + def load_from_state_dict(cls, state_dict, dtype=None, **kwargs): + from transformers import Qwen3Config, Qwen3Model + + qwen_cfg = Qwen3Config( + vocab_size=151669, + hidden_size=1024, + intermediate_size=3072, + num_hidden_layers=28, + num_attention_heads=16, + num_key_value_heads=8, + head_dim=128, + max_position_embeddings=32768, + rms_norm_eps=1e-6, + ) + qwen = Qwen3Model(qwen_cfg) + m2, u2 = qwen.load_state_dict(state_dict, strict=False) + if m2: + print(f" TE missing: {len(m2)} (first 3: {m2[:3]})") + return cls(qwen) + """Wraps Qwen3 weights loaded from AIO. Forward returns last_hidden_state.""" def __init__(self, qwen_model): @@ -1261,41 +1322,13 @@ def load_models(checkpoint_path, device="cuda", dtype=torch.bfloat16): for k, v in sd.items() if k.startswith("model.diffusion_model.") } - cfg = infer_dit_config(dit_sd) - model = AceStep15( - hidden=cfg["hidden"], - inter=cfg["inter"], - heads=cfg["heads"], - kv=cfg["kv"], - head_dim=cfg["head_dim"], - n_dit=cfg["n_dit"], - n_lyric=cfg["n_lyric"], - n_timbre=cfg["n_timbre"], - enc_hidden=cfg["enc_hidden"], - enc_heads=cfg["enc_heads"], - enc_kv=cfg["enc_kv"], - enc_inter=cfg["enc_inter"], - ) - missing, unexpected = model.load_state_dict(dit_sd, strict=False) - # tokenizer/detokenizer keys are expected to be unused (cover mode only) - unexpected = [ - k for k in unexpected if not k.startswith(("tokenizer.", "detokenizer.")) - ] - if missing: - print(f" DiT missing: {len(missing)} (first 3: {missing[:3]})") - if unexpected: - print(f" DiT unexpected: {len(unexpected)} (first 3: {unexpected[:3]})") + model = AceStep15.load_from_state_dict(dit_sd, dtype) model = model.to(device).to(dtype).eval() # --- VAE --- print(" Loading VAE...") vae_sd = {k.removeprefix("vae."): v for k, v in sd.items() if k.startswith("vae.")} - vae = OobleckVAE() - m, u = vae.load_state_dict(vae_sd, strict=False) - if m: - print(f" VAE missing: {len(m)} (first 3: {m[:3]})") - if u: - print(f" VAE unexpected: {len(u)}") + vae = OobleckVAE.load_from_state_dict(vae_sd, dtype) vae = vae.to(device).to(dtype).eval() # --- Text encoder (Qwen3-Embedding from AIO) --- @@ -1305,25 +1338,7 @@ def load_models(checkpoint_path, device="cuda", dtype=torch.bfloat16): for k, v in sd.items() if k.startswith("text_encoders.qwen3_06b.transformer.model.") } - # Load Qwen3 model structure from transformers, then override weights - from transformers import Qwen3Model, Qwen3Config - - qwen_cfg = Qwen3Config( - vocab_size=151669, - hidden_size=1024, - intermediate_size=3072, - num_hidden_layers=28, - num_attention_heads=16, - num_key_value_heads=8, - head_dim=128, - max_position_embeddings=32768, - rms_norm_eps=1e-6, - ) - qwen = Qwen3Model(qwen_cfg) - m2, u2 = qwen.load_state_dict(te_sd, strict=False) - if m2: - print(f" TE missing: {len(m2)} (first 3: {m2[:3]})") - te = TextEncoder(qwen).to(device).to(dtype).eval() + te = TextEncoder.load_from_state_dict(te_sd, dtype).to(device).to(dtype).eval() # Tokenizer — download from HF print(" Loading tokenizer...") diff --git a/extensions_built_in/diffusion_models/anima/anima.py b/extensions_built_in/diffusion_models/anima/anima.py index b22089a..0bd6c99 100644 --- a/extensions_built_in/diffusion_models/anima/anima.py +++ b/extensions_built_in/diffusion_models/anima/anima.py @@ -11,13 +11,16 @@ from toolkit.basic import flush from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.memory_management import MemoryManager from toolkit.models.base_model import BaseModel +from toolkit.models.v2.diffusion_models.cosmos import CosmosTransformer3DModel +from toolkit.models.v2.text_encoders.anima import AnimaTextConditioner +from toolkit.models.v2.text_encoders.qwen3 import Qwen3ModelEncoder +from toolkit.models.v2.vae.qwen_image import QwenImageVAE from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler from toolkit.util.quantize import get_qtype, quantize, quantize_model try: - from diffusers import AnimaAutoBlocks, AnimaModularPipeline, AnimaTextConditioner - from diffusers.models import CosmosTransformer3DModel + from diffusers import AnimaAutoBlocks, AnimaModularPipeline from diffusers.modular_pipelines import SequentialPipelineBlocks from diffusers.modular_pipelines.anima.modular_blocks_anima import AnimaCoreDenoiseStep, AnimaDecodeStep except ImportError as e: @@ -255,12 +258,31 @@ class AnimaModel(BaseModel): self.print_and_status_update("Loading Anima model") pipe: AnimaModularPipeline = AnimaAutoBlocks().init_pipeline(self.model_config.name_or_path) - load_kwargs = {"torch_dtype": dtype} - model_path = os.path.abspath(os.path.expanduser(str(self.model_config.name_or_path))) - if os.path.isdir(model_path): - load_kwargs["pretrained_model_name_or_path"] = model_path - pipe.load_components(**load_kwargs) - pipe.update_components(scheduler=self.get_train_scheduler()) + name = self.model_config.name_or_path + local_path = os.path.abspath(os.path.expanduser(str(name))) + if os.path.isdir(local_path): + name = local_path + + # components load individually through the v2 module classes and are + # handed to the modular pipeline + from transformers import AutoTokenizer + + self.print_and_status_update("Loading components") + transformer = CosmosTransformer3DModel.load_model(name, dtype=dtype) + vae = QwenImageVAE.load_model(name, dtype=dtype) + text_encoder = Qwen3ModelEncoder.load_model(name, dtype=dtype) + text_conditioner = AnimaTextConditioner.load_model(name, dtype=dtype) + tokenizer = AutoTokenizer.from_pretrained(name, subfolder="tokenizer") + t5_tokenizer = AutoTokenizer.from_pretrained(name, subfolder="t5_tokenizer") + pipe.update_components( + transformer=transformer, + vae=vae, + text_encoder=text_encoder, + text_conditioner=text_conditioner, + tokenizer=tokenizer, + t5_tokenizer=t5_tokenizer, + scheduler=self.get_train_scheduler(), + ) transformer = pipe.transformer text_conditioner = pipe.text_conditioner diff --git a/extensions_built_in/diffusion_models/boogu_image/boogu_image.py b/extensions_built_in/diffusion_models/boogu_image/boogu_image.py index 5f70298..2cb4fde 100644 --- a/extensions_built_in/diffusion_models/boogu_image/boogu_image.py +++ b/extensions_built_in/diffusion_models/boogu_image/boogu_image.py @@ -31,6 +31,8 @@ from toolkit.basic import flush from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.models.base_model import BaseModel from toolkit.models.v2.text_encoders.qwen3_vl import patch_qwen_vl_patch_embed +from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLModelEncoder +from toolkit.models.v2.vae.autoencoder_kl import KLVAE from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) @@ -186,8 +188,8 @@ class BooguImageModel(BaseModel): # AutoModel yields the inner Qwen3VLModel (the ``.model`` of the # *ForConditionalGeneration), whose last_hidden_state is exactly the # instruction feature the Boogu pipeline consumes. - text_encoder = AutoModel.from_pretrained( - te_path, subfolder=te_subfolder, torch_dtype=dtype, token=HF_TOKEN + text_encoder = Qwen3VLModelEncoder.load_model( + te_path, dtype=dtype, subfolder=te_subfolder, token=HF_TOKEN ) text_encoder.eval() text_encoder.requires_grad_(False) @@ -217,9 +219,7 @@ class BooguImageModel(BaseModel): # --- VAE (FLUX AutoencoderKL) --- self.print_and_status_update("Loading VAE") - vae = AutoencoderKL.from_pretrained( - base, subfolder="vae", torch_dtype=self.vae_torch_dtype, token=HF_TOKEN - ) + vae = KLVAE.load_model(base, dtype=self.vae_torch_dtype, token=HF_TOKEN) vae.to(self.vae_device_torch, dtype=self.vae_torch_dtype) vae.eval() vae.requires_grad_(False) diff --git a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py index c9d76a6..497207a 100644 --- a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py +++ b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py @@ -5,6 +5,8 @@ import torch import yaml from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.models.base_model import BaseModel +from toolkit.models.v2.text_encoders.mistral3 import Mistral3ModelEncoder +from toolkit.models.v2.vae.autoencoder_kl_flux2 import Flux2KLVAE from toolkit.basic import flush from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import ( @@ -116,8 +118,8 @@ class ErnieImageModel(BaseModel): tokenizer = AutoTokenizer.from_pretrained( base_model_path, subfolder="tokenizer", torch_dtype=dtype ) - text_encoder = AutoModel.from_pretrained( - base_model_path, subfolder="text_encoder", torch_dtype=dtype + text_encoder = Mistral3ModelEncoder.load_model( + base_model_path, subfolder="text_encoder", dtype=dtype ) if ( @@ -140,8 +142,7 @@ class ErnieImageModel(BaseModel): flush() self.print_and_status_update("Loading VAE") - vae = AutoencoderKLFlux2.from_pretrained( - base_model_path, subfolder="vae", torch_dtype=dtype + vae = Flux2KLVAE.load_model( base_model_path, dtype=dtype ).to(self.device_torch, dtype=dtype) self.noise_scheduler = ErnieImageModel.get_train_scheduler() diff --git a/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py b/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py index 864bb41..f565abe 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py @@ -1,5 +1,6 @@ from .flux2_model import Flux2Model from transformers import Qwen3ForCausalLM, Qwen2Tokenizer +from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder from optimum.quanto import freeze from toolkit.util.quantize import quantize, get_qtype from toolkit.config_modules import ModelConfig @@ -40,9 +41,8 @@ class Flux2KleinModel(Flux2Model): dtype = self.torch_dtype self.print_and_status_update("Loading Qwen3") - text_encoder: Qwen3ForCausalLM = Qwen3ForCausalLM.from_pretrained( - self.flux2_klein_te_path, - torch_dtype=dtype, + text_encoder = Qwen3TextEncoder.load_model( + self.flux2_klein_te_path, dtype=dtype, subfolder="" ) if self.model_config.quantize_te: self.print_and_status_update("Quantizing Qwen3") diff --git a/extensions_built_in/diffusion_models/flux2/flux2_model.py b/extensions_built_in/diffusion_models/flux2/flux2_model.py index b61f119..ec4c106 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_model.py @@ -19,6 +19,7 @@ from optimum.quanto import freeze, QTensor from toolkit.util.quantize import quantize, get_qtype, quantize_model from transformers import AutoProcessor, Mistral3ForConditionalGeneration +from toolkit.models.v2.text_encoders.mistral3 import Mistral3TextEncoder from .src.model import Flux2, Flux2Params from .src.pipeline import Flux2Pipeline from toolkit.models.v2.vae.flux2_kl import ( @@ -99,11 +100,8 @@ class Flux2Model(BaseModel): dtype = self.torch_dtype self.print_and_status_update("Loading Mistral") - text_encoder: Mistral3ForConditionalGeneration = ( - Mistral3ForConditionalGeneration.from_pretrained( - MISTRAL_PATH, - torch_dtype=dtype, - ) + text_encoder = Mistral3TextEncoder.load_model( + MISTRAL_PATH, dtype=dtype, subfolder="" ) text_encoder.to(self.device_torch, dtype=dtype) @@ -205,21 +203,8 @@ class Flux2Model(BaseModel): token=HF_TOKEN, ) - vae_state_dict = load_file(vae_path, device="cpu") - - autoencoder_params = AutoEncoderParams() - if vae_state_dict['decoder.up.0.block.0.conv1.bias'].shape[0] == 96: - # this is the small decoder version - autoencoder_params = AutoEncoderSmallDecoderParams() - - with torch.device("meta"): - vae = AutoEncoder(autoencoder_params) - - # cast to dtype - for key in vae_state_dict: - vae_state_dict[key] = vae_state_dict[key].to(dtype) - - vae.load_state_dict(vae_state_dict, assign=True) + # config sniffed from the checkpoint (small-decoder detection) + vae = AutoEncoder.load_model(vae_path, dtype=dtype) self.noise_scheduler = Flux2Model.get_train_scheduler() diff --git a/extensions_built_in/diffusion_models/hidream/hidream_model.py b/extensions_built_in/diffusion_models/hidream/hidream_model.py index 3d72f04..9a1925d 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_model.py @@ -10,6 +10,7 @@ from toolkit.config_modules import GenerateImageConfig, ModelConfig from PIL import Image from toolkit.models.base_model import BaseModel from toolkit.models.v2.text_encoders.t5 import T5TextEncoder +from toolkit.models.v2.text_encoders.llama import LlamaTextEncoder from toolkit.models.v2.text_encoders.clip import CLIPTextEncoderWithProjection from toolkit.models.v2.vae.autoencoder_kl import KLVAE from diffusers import TorchAoConfig @@ -96,11 +97,12 @@ class HidreamModel(BaseModel): use_fast=False ) - text_encoder_4 = LlamaForCausalLM.from_pretrained( + text_encoder_4 = LlamaTextEncoder.load_model( llama_model_path, + dtype=torch.bfloat16, + subfolder="", output_hidden_states=True, output_attentions=True, - torch_dtype=torch.bfloat16, ) text_encoder_4.to(self.device_torch, dtype=dtype) diff --git a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py index 4542bef..1b374ec 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py @@ -1,4 +1,5 @@ import os +from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder from typing import List, Optional import torch @@ -189,7 +190,7 @@ class HidreamO1Model(BaseModel): ) # transformer.load_state_dict(state_dict, assign=True) - transformer = Qwen3VLForConditionalGeneration.from_pretrained( + transformer = Qwen3VLTextEncoder.from_pretrained( None, config=Qwen3VLConfig(**model_config), state_dict=state_dict, @@ -197,7 +198,7 @@ class HidreamO1Model(BaseModel): ) del state_dict # free memory else: - transformer = Qwen3VLForConditionalGeneration.from_pretrained( + transformer = Qwen3VLTextEncoder.from_pretrained( model_path, torch_dtype=self.torch_dtype, ) diff --git a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py index d84c6ab..308892d 100644 --- a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py +++ b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py @@ -31,6 +31,7 @@ from toolkit.models.v2.vae.flux2_kl import ( AutoEncoderParams, convert_diffusers_state_dict, ) +from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLModelEncoder from .src.latent_norm import get_latent_norm from .src.pipeline import ( Ideogram4Pipeline, @@ -224,8 +225,8 @@ class Ideogram4Model(BaseModel): self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}") tokenizer = AutoTokenizer.from_pretrained(te_path, token=HF_TOKEN) - text_encoder = AutoModel.from_pretrained( - te_path, torch_dtype=dtype, token=HF_TOKEN + text_encoder = Qwen3VLModelEncoder.load_model( + te_path, dtype=dtype, subfolder="", token=HF_TOKEN ) flush() @@ -268,8 +269,7 @@ class Ideogram4Model(BaseModel): self.print_and_status_update("Loading VAE") vae_sd = _load_component_state_dict(base, "vae", "diffusion_pytorch_model") vae_sd = convert_diffusers_state_dict(vae_sd) - vae = AutoEncoder(AutoEncoderParams()) - vae.load_state_dict(vae_sd) + vae = AutoEncoder.load_from_state_dict(vae_sd, self.vae_torch_dtype) del vae_sd vae.to(self.vae_device_torch, dtype=dtype) vae.eval() diff --git a/extensions_built_in/diffusion_models/ltx2/convert_ltx2_to_diffusers.py b/extensions_built_in/diffusion_models/ltx2/convert_ltx2_to_diffusers.py index 091d258..0b52e27 100644 --- a/extensions_built_in/diffusion_models/ltx2/convert_ltx2_to_diffusers.py +++ b/extensions_built_in/diffusion_models/ltx2/convert_ltx2_to_diffusers.py @@ -7,12 +7,16 @@ from typing import Any, Dict, Tuple import torch from accelerate import init_empty_weights -from diffusers import ( - AutoencoderKLLTX2Audio, - AutoencoderKLLTX2Video, +from toolkit.models.v2.diffusion_models.ltx2 import ( + LTX2TextConnectors, LTX2VideoTransformer3DModel, + LTX2Vocoder, + LTX2VocoderWithBWE, +) +from toolkit.models.v2.vae.ltx2 import ( + LTX2AudioVAE as AutoencoderKLLTX2Audio, + LTX2VideoVAE as AutoencoderKLLTX2Video, ) -from diffusers.pipelines.ltx2 import LTX2TextConnectors, LTX2Vocoder, LTX2VocoderWithBWE from diffusers.utils.import_utils import is_accelerate_available diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index 94ca300..dc561a3 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -21,24 +21,29 @@ from optimum.quanto import freeze from toolkit.util.quantize import quantize, get_qtype, quantize_model from toolkit.memory_management import MemoryManager from toolkit.paths import MODELS_PATH +from toolkit.util.mixed_precision import attach_per_op_casting, pin_stored_fp32 from safetensors.torch import load_file from PIL import Image import huggingface_hub try: from diffusers import LTX2Pipeline, LTX2ImageToVideoPipeline - from diffusers.models.autoencoders import ( - AutoencoderKLLTX2Audio, - AutoencoderKLLTX2Video, - ) - from diffusers.models.transformers import LTX2VideoTransformer3DModel from diffusers.pipelines.ltx2.export_utils import encode_video from transformers import ( Gemma3ForConditionalGeneration, GemmaTokenizerFast, ) - from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder, LTX2VocoderWithBWE - from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors + from toolkit.models.v2.diffusion_models.ltx2 import ( + LTX2TextConnectors, + LTX2VideoTransformer3DModel, + LTX2Vocoder, + LTX2VocoderWithBWE, + ) + from toolkit.models.v2.vae.ltx2 import ( + LTX2AudioVAE as AutoencoderKLLTX2Audio, + LTX2VideoVAE as AutoencoderKLLTX2Video, + ) + from toolkit.models.v2.text_encoders.gemma3 import Gemma3TextEncoder from .convert_ltx2_to_diffusers import ( get_model_state_dict_from_combined_ckpt, convert_ltx2_transformer, @@ -293,20 +298,14 @@ class LTX2Model(BaseModel): ) transformer = transformer.to(dtype) 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") + if 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 - transformer = LTX2VideoTransformer3DModel.from_pretrained( - transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype - ) + transformer = LTX2VideoTransformer3DModel.load_model(model_path, dtype=dtype) if self.model_config.quantize: self.print_and_status_update("Quantizing Transformer") @@ -347,7 +346,7 @@ class LTX2Model(BaseModel): tokenizer = GemmaTokenizerFast.from_pretrained(base_te_path) with init_empty_weights(): - text_encoder = Gemma3ForConditionalGeneration( + text_encoder = Gemma3TextEncoder( Gemma3Config( **{ "boi_token_index": 255999, @@ -420,22 +419,22 @@ class LTX2Model(BaseModel): tokenizer = GemmaTokenizerFast.from_pretrained( self.model_config.te_name_or_path ) - text_encoder = Gemma3ForConditionalGeneration.from_pretrained( - self.model_config.te_name_or_path, dtype=dtype + text_encoder = Gemma3TextEncoder.load_model( + self.model_config.te_name_or_path, dtype=dtype, subfolder="" ) elif self.ltx_te_path is not None: # pull from model specific te tokenizer = GemmaTokenizerFast.from_pretrained(self.ltx_te_path) - text_encoder = Gemma3ForConditionalGeneration.from_pretrained( - self.ltx_te_path, dtype=dtype + text_encoder = Gemma3TextEncoder.load_model( + self.ltx_te_path, dtype=dtype, subfolder="" ) else: # using combo hf repo tokenizer = GemmaTokenizerFast.from_pretrained( self.model_config.name_or_path, subfolder="tokenizer" ) - text_encoder = Gemma3ForConditionalGeneration.from_pretrained( - self.model_config.name_or_path, subfolder="text_encoder", dtype=dtype + text_encoder = Gemma3TextEncoder.load_model( + self.model_config.name_or_path, dtype=dtype ) # remove the vision tower @@ -497,24 +496,16 @@ class LTX2Model(BaseModel): del combined_state_dict flush() else: - vae = AutoencoderKLLTX2Video.from_pretrained( - base_model_path, subfolder="vae", torch_dtype=dtype - ) - audio_vae = AutoencoderKLLTX2Audio.from_pretrained( - base_model_path, subfolder="audio_vae", torch_dtype=dtype - ) + vae = AutoencoderKLLTX2Video.load_model(base_model_path, dtype=dtype) + audio_vae = AutoencoderKLLTX2Audio.load_model(base_model_path, dtype=dtype) - connectors = LTX2TextConnectors.from_pretrained( - base_model_path, subfolder="connectors", torch_dtype=dtype - ) + connectors = LTX2TextConnectors.load_model(base_model_path, dtype=dtype) vocoder_cls = LTX2Vocoder if self.ltx_version in ("2.3", "2.5"): vocoder_cls = LTX2VocoderWithBWE - vocoder = vocoder_cls.from_pretrained( - base_model_path, subfolder="vocoder", torch_dtype=dtype - ) + vocoder = vocoder_cls.load_model(base_model_path, dtype=dtype) self.noise_scheduler = LTX2Model.get_train_scheduler() @@ -1299,7 +1290,7 @@ class LTX25Model(LTX2Model): dropped, matching how the Gemma-3 vision tower was discarded.""" from safetensors import safe_open from transformers import Gemma4TextConfig - from transformers.models.gemma4.modeling_gemma4 import Gemma4TextModel + from toolkit.models.v2.text_encoders.gemma3 import Gemma4TextEncoder with safe_open(te_path, framework="pt") as f: metadata = f.metadata() or {} @@ -1309,7 +1300,7 @@ class LTX25Model(LTX2Model): } with init_empty_weights(): - text_encoder = Gemma4TextModel(Gemma4TextConfig(**text_config)) + text_encoder = Gemma4TextEncoder(Gemma4TextConfig(**text_config)) def strip_model_prefix(key: str) -> str: return key[len("model.") :] if key.startswith("model.") else key @@ -1374,12 +1365,14 @@ class LTX25Model(LTX2Model): transformer, transformer_sd, "transformer" ) del transformer_sd - # cast to the compute dtype even when pre-quantized: the comfy file - # stores the scale_shift modulation tables in fp32, which promotes - # hidden states to fp32 and breaks the (bf16) diffusers attention - # linears. Quantized backends are immune (int8 qdata, uint8-viewed - # scales survive a dtype cast untouched). - transformer = transformer.to(dtype) + if num_quantized_dit == 0: + transformer = transformer.to(dtype) + else: + # mixed-precision comfy file (fp32 scale_shift tables next to bf16 + # weights): comfy-style per-op input casting, with the stored-fp32 + # pieces pinned so parent .to(dtype) casts cannot downcast them + attach_per_op_casting(transformer) + pin_stored_fp32(transformer) flush() if self.model_config.quantize: @@ -1424,7 +1417,11 @@ class LTX25Model(LTX2Model): connectors, connectors_sd, "connectors" ) del connectors_sd, dit_sd - connectors = connectors.to(dtype) + if num_quantized_connectors == 0: + connectors = connectors.to(dtype) + else: + attach_per_op_casting(connectors) + pin_stored_fp32(connectors) flush() # ---- text encoder (Gemma-4 12B, single comfy file) ---- diff --git a/extensions_built_in/diffusion_models/mageflow/mageflow.py b/extensions_built_in/diffusion_models/mageflow/mageflow.py index 1f11192..d54d628 100644 --- a/extensions_built_in/diffusion_models/mageflow/mageflow.py +++ b/extensions_built_in/diffusion_models/mageflow/mageflow.py @@ -230,8 +230,8 @@ class MageFlowModel(BaseModel): vae_path = self.model_config.model_kwargs.get("vae_path", None) if vae_path is None: vae_path = self._get_model_file("vae/diffusion_pytorch_model.safetensors") - vae = MageVAE( - ckpt_path=vae_path, + vae = MageVAE.load_model( + vae_path, sample_posterior=bool( self.model_config.model_kwargs.get("vae_sample_posterior", True) ), diff --git a/extensions_built_in/diffusion_models/mageflow/src/vae.py b/extensions_built_in/diffusion_models/mageflow/src/vae.py index 1995315..aecaf40 100644 --- a/extensions_built_in/diffusion_models/mageflow/src/vae.py +++ b/extensions_built_in/diffusion_models/mageflow/src/vae.py @@ -17,6 +17,8 @@ import os from functools import lru_cache import torch + +from toolkit.models.v2._mixin import OstrisModelMixin import torch.nn as nn import torch.nn.functional as F @@ -591,7 +593,7 @@ def _load_state_dict(ckpt_path: str): return state -class MageVAE(nn.Module): +class MageVAE(nn.Module, OstrisModelMixin): """ Encode: DConvEncoder (one-step diffusion) → latent [B, 128, H/16, W/16] Decode: DConvDenoiser + CoD Decoder → image [B, 3, H, W] in [-1, 1] @@ -600,20 +602,32 @@ class MageVAE(nn.Module): latent_channels = 128 downsample_factor = 16 - def __init__(self, ckpt_path: str, sample_posterior: bool = True): + def __init__(self, ckpt_path: str = None, sample_posterior: bool = True): super().__init__() self.sample_posterior = sample_posterior self.dconv_encoder = _DConvEncoder() self.decoder_model = _DConvDenoiser() - sd = _load_state_dict(ckpt_path) - self._load_encoder(sd, ckpt_path) - self._load_decoder(sd, ckpt_path) + if ckpt_path is not None: + sd = _load_state_dict(ckpt_path) + self._load_encoder(sd, ckpt_path) + self._load_decoder(sd, ckpt_path) - # adaLN modulation depends only on t, and we always run at t=0. - # Precompute and drop the MLPs once at construction (~37M params saved). - self._freeze_adaln_cache() + # adaLN modulation depends only on t, and we always run at t=0. + # Precompute and drop the MLPs once at construction (~37M params saved). + self._freeze_adaln_cache() + + @classmethod + def load_from_state_dict(cls, state_dict, dtype=None, **kwargs): + model = cls( + ckpt_path=None, + sample_posterior=kwargs.pop("sample_posterior", True), + ) + model._load_encoder(state_dict, "") + model._load_decoder(state_dict, "") + model._freeze_adaln_cache() + return model def _load_encoder(self, sd, ckpt_path): prefix = "student.dconv_encoder." diff --git a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py index 05a0d4a..c9ce2ed 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py @@ -52,6 +52,7 @@ from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.memory_management import MemoryManager from toolkit.metadata import get_meta_for_safetensors from toolkit.models.base_model import BaseModel +from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder from toolkit.models.v2.resolver import ( find_file_recursive, repo_id_from_name_or_path, @@ -69,7 +70,7 @@ from optimum.quanto import freeze from .src import packing packing_video_exts = [".mp4", ".avi", ".mov", ".webm", ".mkv", ".wmv", ".m4v", ".flv"] -from .src.audio_vae import MiniMaxH3AudioVAE, fold_audio_vae_weight_norm +from .src.audio_vae import MiniMaxH3AudioVAE from .src.packing import ( KEYFRAME_ENCODE_SEED, KEYFRAME_NOISE_AUG_T, @@ -404,7 +405,7 @@ class MinimaxH3Model(BaseModel): ) config = AutoConfig.from_pretrained(te_path) config.text_config.num_hidden_layers = TEXT_ENCODER_LAYER - text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( + text_encoder = Qwen3VLTextEncoder.from_pretrained( te_path, config=config, torch_dtype=self.te_torch_dtype ) else: @@ -426,7 +427,7 @@ class MinimaxH3Model(BaseModel): config.text_config.num_hidden_layers = TEXT_ENCODER_LAYER config.tie_word_embeddings = False with init_empty_weights(): - text_encoder = Qwen3VLForConditionalGeneration(config) + text_encoder = Qwen3VLTextEncoder(config) text_encoder.lm_head = None state_dict = load_file(te_file) @@ -481,38 +482,9 @@ class MinimaxH3Model(BaseModel): def _load_vaes(self) -> MiniMaxH3VaeBundle: self.print_and_status_update("Loading video VAE") - video_sd = load_file(self._resolve_comfy_file("video_vae")) - # normalization stats ride along in the comfy file; the module holds - # them as non-persistent buffers, keep them float32 - video_stats = { - k: video_sd.pop(k).float() - for k in ("latents_mean", "latents_std") - if k in video_sd - } - video_vae = MiniMaxH3VideoVAE() - video_vae.load_state_dict(video_sd, strict=True, assign=True) - for k, v in video_stats.items(): - getattr(video_vae, k).copy_(v) - video_vae.eval().requires_grad_(False) - del video_sd - + video_vae = MiniMaxH3VideoVAE.load_model(self._resolve_comfy_file("video_vae")) self.print_and_status_update("Loading audio VAE") - audio_sd = load_file(self._resolve_comfy_file("audio_vae")) - audio_stats = { - k: audio_sd.pop(k).float() - for k in ("latents_mean", "latents_std") - if k in audio_sd - } - # comfy repack ships the weight norm already folded; fold only if the - # raw parametrization is present (original-repo file) - if any(k.endswith("weight_g") for k in audio_sd.keys()): - audio_sd = fold_audio_vae_weight_norm(audio_sd) - audio_vae = MiniMaxH3AudioVAE() - audio_vae.load_state_dict(audio_sd, strict=True, assign=True) - for k, v in audio_stats.items(): - getattr(audio_vae, k).copy_(v) - audio_vae.to(torch.float32).eval().requires_grad_(False) - del audio_sd + audio_vae = MiniMaxH3AudioVAE.load_model(self._resolve_comfy_file("audio_vae")) flush() return MiniMaxH3VaeBundle(video_vae, audio_vae) diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/audio_vae.py b/extensions_built_in/diffusion_models/minimax_h3/src/audio_vae.py index a12830a..bdd081a 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/audio_vae.py +++ b/extensions_built_in/diffusion_models/minimax_h3/src/audio_vae.py @@ -17,6 +17,8 @@ weights:: import math import torch + +from toolkit.models.v2._mixin import OstrisModelMixin import torch.nn as nn import torch.nn.functional as F from torch import Tensor @@ -420,7 +422,26 @@ class BigVGANDecoder(nn.Module): return torch.clamp(x, min=-1.0, max=1.0) -class MiniMaxH3AudioVAE(nn.Module): +class MiniMaxH3AudioVAE(nn.Module, OstrisModelMixin): + @classmethod + def load_from_state_dict(cls, state_dict, dtype=None, **kwargs): + """Comfy single-file load: stats buffers ride along; original-repo + files still carry the raw weight-norm parametrization and get folded.""" + state_dict = dict(state_dict) + stats = { + k: state_dict.pop(k).float() + for k in ("latents_mean", "latents_std") + if k in state_dict + } + if any(k.endswith("weight_g") for k in state_dict.keys()): + state_dict = fold_audio_vae_weight_norm(state_dict) + model = cls() + model.load_state_dict(state_dict, strict=True, assign=True) + for k, v in stats.items(): + getattr(model, k).copy_(v) + model.to(torch.float32).eval().requires_grad_(False) + return model + """Frozen MiniMax-H3 audio autoencoder with diagonal latent normalization. - ``encode``: waveform ``(B, 1, samples)`` at 32 kHz -> latents ``(B, 32, T)`` diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/vae.py b/extensions_built_in/diffusion_models/minimax_h3/src/vae.py index a159a9a..d57aa15 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/vae.py +++ b/extensions_built_in/diffusion_models/minimax_h3/src/vae.py @@ -22,6 +22,8 @@ import math from contextlib import nullcontext import torch + +from toolkit.models.v2._mixin import OstrisModelMixin import torch.nn as nn import torch.nn.functional as F from torch.utils.checkpoint import checkpoint @@ -434,7 +436,24 @@ class ViTDecoder3d(nn.Module): # --------------------------------------------------------------------------- -class MiniMaxH3VideoVAE(nn.Module): +class MiniMaxH3VideoVAE(nn.Module, OstrisModelMixin): + @classmethod + def load_from_state_dict(cls, state_dict, dtype=None, **kwargs): + """Comfy single-file load: the normalization stats ride along in the + file and land in the non-persistent fp32 buffers.""" + state_dict = dict(state_dict) + stats = { + k: state_dict.pop(k).float() + for k in ("latents_mean", "latents_std") + if k in state_dict + } + model = cls() + model.load_state_dict(state_dict, strict=True, assign=True) + for k, v in stats.items(): + getattr(model, k).copy_(v) + model.eval().requires_grad_(False) + return model + def __init__( self, in_channels=3, diff --git a/extensions_built_in/diffusion_models/omnigen2/__init__.py b/extensions_built_in/diffusion_models/omnigen2/__init__.py index 09b00f9..8887767 100644 --- a/extensions_built_in/diffusion_models/omnigen2/__init__.py +++ b/extensions_built_in/diffusion_models/omnigen2/__init__.py @@ -1,4 +1,6 @@ import os +from toolkit.models.v2.text_encoders.qwen25_vl import Qwen25VLTextEncoder +from toolkit.models.v2.vae.autoencoder_kl import KLVAE from typing import TYPE_CHECKING, List, Optional import torch @@ -78,8 +80,8 @@ class OmniGen2Model(BaseModel): extras_path, subfolder="processor", use_fast=True ) - mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained( - extras_path, subfolder="mllm", torch_dtype=torch.bfloat16 + mllm = Qwen25VLTextEncoder.load_model( + extras_path, dtype=torch.bfloat16, subfolder="mllm" ) mllm.to(self.device_torch, dtype=dtype) if self.model_config.quantize_te: @@ -117,9 +119,9 @@ class OmniGen2Model(BaseModel): self.print_and_status_update("Loading vae") - vae = AutoencoderKL.from_pretrained( - extras_path, subfolder="vae", torch_dtype=torch.bfloat16 - ).to(self.device_torch, dtype=dtype) + vae = KLVAE.load_model(extras_path, dtype=torch.bfloat16).to( + self.device_torch, dtype=dtype + ) flush() self.print_and_status_update("Loading Qwen2.5 VLProcessor") diff --git a/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py b/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py index ffea929..7358b1b 100644 --- a/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py +++ b/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py @@ -49,6 +49,7 @@ from toolkit.accelerator import unwrap_model from toolkit.basic import flush from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.models.base_model import BaseModel +from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextOnlyEncoder from toolkit.models.FakeVAE import FakeVAE from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import ( @@ -144,9 +145,7 @@ class PRXPixelT2IModel(BaseModel): # --- text encoder + tokenizer (Qwen3-VL text tower) --- self.print_and_status_update("Loading text encoder") tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder="tokenizer") - text_encoder = Qwen3VLTextModel.from_pretrained( - model_path, subfolder="text_encoder", torch_dtype=dtype - ) + text_encoder = Qwen3VLTextOnlyEncoder.load_model(model_path, dtype=dtype) text_encoder.to(self.te_device_torch) text_encoder.eval() text_encoder.requires_grad_(False) diff --git a/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py b/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py index 46ca0d0..289ca1a 100644 --- a/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py +++ b/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py @@ -22,7 +22,7 @@ HF_TOKEN = os.getenv("HF_TOKEN", None) try: from diffusers import ZImagePipeline - from diffusers.models.transformers.transformer_z_image import ( + from toolkit.models.v2.diffusion_models.z_image import ( ZImageTransformer2DModel as ZImageTransformer2DModelOriginal, ) from diffusers.models.modeling_outputs import Transformer2DModelOutput diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index 77ae4ca..56da4f9 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -174,6 +174,19 @@ MODEL_TESTS = { "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 4.0}, "needs_control_image": True, }, + # ---- legacy monolith archs (components adopted into v2 on load) ---- + "sd1": { + "model": {"name_or_path": "stable-diffusion-v1-5/stable-diffusion-v1-5"}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 7.5, "seed": 42}, + }, + "sdxl": { + "model": {"name_or_path": "stabilityai/stable-diffusion-xl-base-1.0"}, + "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 6.0, "seed": 42}, + }, + "ace_step_15": { + "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 4.0, "seed": 42}, + }, "f-lite": { "model": {"name_or_path": "Freepik/F-Lite", "quantize": True, "quantize_te": True}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, @@ -217,9 +230,12 @@ def run_one(arch: str, device: str, allow_download: bool) -> dict: model_config = ModelConfig(arch=arch, dtype="bf16", **entry["model"]) ModelClass = get_model_class(model_config) - # get_model_class silently falls back to the legacy SD class on an - # unknown arch; that is never what a registered test wants - if getattr(ModelClass, "arch", None) not in (arch, model_config.arch): + from toolkit.util.get_model import LEGACY_ARCHS + + if ( + getattr(ModelClass, "arch", None) not in (arch, model_config.arch) + and model_config.arch not in LEGACY_ARCHS + ): raise ValueError( f"arch {arch!r} resolved to {ModelClass.__name__} " f"(arch={getattr(ModelClass, 'arch', None)!r}) — registry mismatch" @@ -228,6 +244,26 @@ def run_one(arch: str, device: str, allow_download: bool) -> dict: sampler = None if hasattr(ModelClass, "get_train_scheduler"): sampler = ModelClass.get_train_scheduler() + else: + # legacy monolith archs build their scheduler the way training does + from toolkit.sampler import get_sampler + + legacy_arch = "sd" + if model_config.is_pixart: + legacy_arch = "pixart" + elif model_config.is_flux: + legacy_arch = "flux" + elif model_config.is_lumina2: + legacy_arch = "lumina2" + sampler = get_sampler( + "ddpm", + { + "prediction_type": "v_prediction" + if model_config.is_v_pred + else "epsilon", + }, + arch=legacy_arch, + ) sd = ModelClass( device=device, @@ -255,7 +291,11 @@ def run_one(arch: str, device: str, allow_download: bool) -> dict: output_ext="png", **sample_kwargs, ) - sd.generate_images([gen]) + gen_kwargs = {} + if not hasattr(ModelClass, "get_train_scheduler"): + # the legacy monolith takes the sampler NAME at generate time + gen_kwargs["sampler"] = "ddpm" + sd.generate_images([gen], **gen_kwargs) produced = [ p diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index 3e943cc..52a509e 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -371,6 +371,43 @@ Decisions: unload/load the difference - [ ] Legacy `stable_diffusion_model.py` archs: grandfather or port last +## 100% mixin coverage (2026-08-28) + +Every component of every non-legacy arch now loads through OstrisModelMixin — +DiTs, text encoders, VAEs, vision encoders, connectors/vocoders, and the +custom cases that previously bypassed it: + +- New wrapper classes: llama, gemma3 + gemma4, mistral3 (×2), qwen3 base, + qwen3-vl base + text-only, wan VAE, diffusers flux2 KL VAE, CLIP vision + (first vision encoder), cosmos DiT, anima text conditioner, the full ltx2 + set (transformer, video/audio VAEs, connectors, vocoders). +- Existing-class swaps: omnigen2 (mllm + VAE), boogu (TE + VAE), klein TE, + hidream llama, ideogram4 TE, prx TE. +- Custom restructures: minimax video/audio VAEs and TE, MageVAE + (deferred-load ctor), the shared flux2_kl AutoEncoder (small-decoder sniff + as a class hook; flux2 + ideogram4 route through it), ace_step's bundle + (per-component class loaders), anima (modular-pipeline load replaced with + component-wise v2 loads + update_components), hidream_o1's Qwen3VL DiT, + z_image_l2p rebased onto the v2 class, ltx2's converter builds v2 classes. +- Mixin: kwargs flow through the single-file chain (component ctor args like + MageVAE's sample_posterior). + +Verified with real weights: ltx2.3 + ltx2.5 (full v2 family), anima, +ace_step_15 (new harness entry), minimax VAEs, plus the standing harness +coverage. The legacy monolith archs joined the system per the inference-engine +goal ("send a job with any base_model, reload/unload components on the fly"): +six new wrapper classes complete the component vocabulary (UNet2DCondition, +SD3/PixArt ×2/AuraFlow/Lumina2 transformers, Gemma2), and +`adopt_component` (in-place class swap, OstrisLinear-style) rebinds +pipeline-loaded components onto their wrappers at the monolith's single +post-load funnel — covering every legacy arch without touching its fifteen +load branches, with all pipeline references staying valid. Verified: sd1 +(UNet/KLVAE/CLIPTextEncoder all mixin instances + generation) and sdxl +(dual-CLIP adoption + 1024² generation); harness gained sd1/sdxl entries, +a legacy scheduler fallback, and the sampler-name pass-through. Full +monolith decomposition (per-arch v2 loading with comfy candidates) remains +future work, but every resident component is now poolable by the engine. + ## Testing - [x] `testing/test_model_loading.py`: per-arch load + one small sample through @@ -418,8 +455,12 @@ Decisions: flux2_klein_9b, prx_pixel, zeta_chroma, zimage_l2p, both qwen edit archs, f-lite. Fixes found by the run: ltx2.5's fp32 scale_shift tables promoted hidden states into bf16 linears under the diffusers - class (ComfyUI casts per-op) — the DiT/connectors now cast to compute - dtype after quantized attach (ConvRot storage immune); harness configs + class — fixed ComfyUI-style (toolkit/util/mixed_precision.py): + per-op input casting hooks on every weighted module + the stored-fp32 + tensors pinned against parent .to(dtype) casts (device moves pass + through), so the tables stay genuinely fp32 at sample time. The + mechanism is general — any mixed-precision comfy checkpoint on a + diffusers-class arch can use it; harness configs for e1 resolution and zeta/l2p extras_name_or_path corrected to match the UI defaults. - [ ] Each newly migrated model adds its test in the same PR as its migration. diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index e4d58de..c2e9699 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -42,6 +42,24 @@ class BasicModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMix 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 @@ -195,6 +213,7 @@ class OstrisModelMixin: config_path=config_path, config=config, subfolder=subfolder, + **kwargs, ) else: if os.path.isdir(name_or_path): @@ -259,6 +278,7 @@ class OstrisModelMixin: config_path: Optional[str] = None, config=None, subfolder: Optional[str] = None, + **kwargs, ): state_dict = load_file(file_path) return cls.load_from_state_dict( @@ -267,6 +287,7 @@ class OstrisModelMixin: config_path=config_path, config=config, subfolder=subfolder, + **kwargs, ) @classmethod @@ -284,6 +305,7 @@ class OstrisModelMixin: 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 diff --git a/toolkit/models/v2/diffusion_models/auraflow.py b/toolkit/models/v2/diffusion_models/auraflow.py new file mode 100644 index 0000000..829faa7 --- /dev/null +++ b/toolkit/models/v2/diffusion_models/auraflow.py @@ -0,0 +1,13 @@ +from diffusers import AuraFlowTransformer2DModel as DiffusersAuraFlowTransformer2DModel + +from .._mixin import OstrisModelMixin + + +class AuraFlowTransformer2DModel( + DiffusersAuraFlowTransformer2DModel, OstrisModelMixin +): + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["joint_transformer_blocks", "single_transformer_blocks"] diff --git a/toolkit/models/v2/diffusion_models/cosmos.py b/toolkit/models/v2/diffusion_models/cosmos.py new file mode 100644 index 0000000..4c2a034 --- /dev/null +++ b/toolkit/models/v2/diffusion_models/cosmos.py @@ -0,0 +1,15 @@ +from diffusers.models import ( + CosmosTransformer3DModel as DiffusersCosmosTransformer3DModel, +) + +from .._mixin import OstrisModelMixin + + +class CosmosTransformer3DModel(DiffusersCosmosTransformer3DModel, OstrisModelMixin): + """Cosmos video/image DiT (anima).""" + + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["transformer_blocks"] diff --git a/toolkit/models/v2/diffusion_models/ltx2.py b/toolkit/models/v2/diffusion_models/ltx2.py new file mode 100644 index 0000000..2657336 --- /dev/null +++ b/toolkit/models/v2/diffusion_models/ltx2.py @@ -0,0 +1,34 @@ +from diffusers.models.transformers import ( + LTX2VideoTransformer3DModel as DiffusersLTX2VideoTransformer3DModel, +) +from diffusers.pipelines.ltx2 import ( + LTX2TextConnectors as DiffusersLTX2TextConnectors, +) +from diffusers.pipelines.ltx2 import LTX2Vocoder as DiffusersLTX2Vocoder +from diffusers.pipelines.ltx2 import ( + LTX2VocoderWithBWE as DiffusersLTX2VocoderWithBWE, +) + +from .._mixin import OstrisModelMixin + + +class LTX2VideoTransformer3DModel( + DiffusersLTX2VideoTransformer3DModel, OstrisModelMixin +): + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["transformer_blocks"] + + +class LTX2TextConnectors(DiffusersLTX2TextConnectors, OstrisModelMixin): + aitk_subfolder = "connectors" + + +class LTX2Vocoder(DiffusersLTX2Vocoder, OstrisModelMixin): + aitk_subfolder = "vocoder" + + +class LTX2VocoderWithBWE(DiffusersLTX2VocoderWithBWE, OstrisModelMixin): + aitk_subfolder = "vocoder" diff --git a/toolkit/models/v2/diffusion_models/lumina2.py b/toolkit/models/v2/diffusion_models/lumina2.py new file mode 100644 index 0000000..410c972 --- /dev/null +++ b/toolkit/models/v2/diffusion_models/lumina2.py @@ -0,0 +1,13 @@ +from diffusers import Lumina2Transformer2DModel as DiffusersLumina2Transformer2DModel + +from .._mixin import OstrisModelMixin + + +class Lumina2Transformer2DModel( + DiffusersLumina2Transformer2DModel, OstrisModelMixin +): + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["layers"] diff --git a/toolkit/models/v2/diffusion_models/pixart.py b/toolkit/models/v2/diffusion_models/pixart.py new file mode 100644 index 0000000..231100a --- /dev/null +++ b/toolkit/models/v2/diffusion_models/pixart.py @@ -0,0 +1,24 @@ +from diffusers import PixArtTransformer2DModel as DiffusersPixArtTransformer2DModel +from diffusers import Transformer2DModel as DiffusersTransformer2DModel + +from .._mixin import OstrisModelMixin + + +class PixArtTransformer2DModel( + DiffusersPixArtTransformer2DModel, OstrisModelMixin +): + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["transformer_blocks"] + + +class Transformer2DModel(DiffusersTransformer2DModel, OstrisModelMixin): + """The generic diffusers DiT the pixart sigma path loads.""" + + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["transformer_blocks"] diff --git a/toolkit/models/v2/diffusion_models/sd3.py b/toolkit/models/v2/diffusion_models/sd3.py new file mode 100644 index 0000000..0c839b8 --- /dev/null +++ b/toolkit/models/v2/diffusion_models/sd3.py @@ -0,0 +1,11 @@ +from diffusers import SD3Transformer2DModel as DiffusersSD3Transformer2DModel + +from .._mixin import OstrisModelMixin + + +class SD3Transformer2DModel(DiffusersSD3Transformer2DModel, OstrisModelMixin): + aitk_subfolder = "transformer" + + @classmethod + def get_transformer_block_names(cls): + return ["transformer_blocks"] diff --git a/toolkit/models/v2/diffusion_models/unet.py b/toolkit/models/v2/diffusion_models/unet.py new file mode 100644 index 0000000..3c664c0 --- /dev/null +++ b/toolkit/models/v2/diffusion_models/unet.py @@ -0,0 +1,13 @@ +from diffusers import UNet2DConditionModel as DiffusersUNet2DConditionModel + +from .._mixin import OstrisModelMixin + + +class UNet2DConditionModel(DiffusersUNet2DConditionModel, OstrisModelMixin): + """The SD1/SD2/SDXL-family UNet.""" + + aitk_subfolder = "unet" + + @classmethod + def get_transformer_block_names(cls): + return ["down_blocks", "up_blocks"] diff --git a/toolkit/models/v2/text_encoders/anima.py b/toolkit/models/v2/text_encoders/anima.py new file mode 100644 index 0000000..7430404 --- /dev/null +++ b/toolkit/models/v2/text_encoders/anima.py @@ -0,0 +1,9 @@ +from diffusers import AnimaTextConditioner as DiffusersAnimaTextConditioner + +from .._mixin import OstrisModelMixin + + +class AnimaTextConditioner(DiffusersAnimaTextConditioner, OstrisModelMixin): + """Anima's learned text conditioner (rides next to the Qwen3 encoder).""" + + aitk_subfolder = "text_conditioner" diff --git a/toolkit/models/v2/text_encoders/gemma2.py b/toolkit/models/v2/text_encoders/gemma2.py new file mode 100644 index 0000000..0fe1e1c --- /dev/null +++ b/toolkit/models/v2/text_encoders/gemma2.py @@ -0,0 +1,14 @@ +from transformers import Gemma2Model + +from .._mixin import OstrisTransformersMixin + + +class Gemma2ModelEncoder(Gemma2Model, OstrisTransformersMixin): + """Gemma2 base model (lumina2's text encoder — what AutoModel resolves).""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + return ["layers"] diff --git a/toolkit/models/v2/text_encoders/gemma3.py b/toolkit/models/v2/text_encoders/gemma3.py new file mode 100644 index 0000000..8391fb3 --- /dev/null +++ b/toolkit/models/v2/text_encoders/gemma3.py @@ -0,0 +1,32 @@ +from transformers import Gemma3ForConditionalGeneration + +from .._mixin import OstrisTransformersMixin + + +class Gemma3TextEncoder(Gemma3ForConditionalGeneration, OstrisTransformersMixin): + """Gemma3 conditioning stack (ltx2 / ltx2.3).""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + # both layouts seen across transformers versions; missing paths skip + return ["model.language_model.layers", "language_model.model.layers"] + + +try: + from transformers.models.gemma4.modeling_gemma4 import Gemma4TextModel + + class Gemma4TextEncoder(Gemma4TextModel, OstrisTransformersMixin): + """Gemma4 text decoder (ltx2.5's conditioning stack).""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + return ["layers"] + +except ImportError: + Gemma4TextEncoder = None diff --git a/toolkit/models/v2/text_encoders/llama.py b/toolkit/models/v2/text_encoders/llama.py new file mode 100644 index 0000000..be2e4ed --- /dev/null +++ b/toolkit/models/v2/text_encoders/llama.py @@ -0,0 +1,14 @@ +from transformers import LlamaForCausalLM + +from .._mixin import OstrisTransformersMixin + + +class LlamaTextEncoder(LlamaForCausalLM, OstrisTransformersMixin): + """Llama causal-LM text encoder (hidream's text_encoder_4).""" + + aitk_subfolder = "text_encoder_4" + aitk_tokenizer_subfolder = "tokenizer_4" + + @classmethod + def get_transformer_block_names(cls): + return ["model.layers"] diff --git a/toolkit/models/v2/text_encoders/mistral3.py b/toolkit/models/v2/text_encoders/mistral3.py new file mode 100644 index 0000000..29be74e --- /dev/null +++ b/toolkit/models/v2/text_encoders/mistral3.py @@ -0,0 +1,25 @@ +from transformers import Mistral3ForConditionalGeneration, Mistral3Model + +from .._mixin import OstrisTransformersMixin + + +class Mistral3TextEncoder(Mistral3ForConditionalGeneration, OstrisTransformersMixin): + """Mistral3 conditioning stack (flux2).""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + return ["model.language_model.layers", "language_model.model.layers"] + + +class Mistral3ModelEncoder(Mistral3Model, OstrisTransformersMixin): + """The inner Mistral3 base model (ernie_image's text encoder).""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + return ["language_model.layers"] diff --git a/toolkit/models/v2/text_encoders/qwen3.py b/toolkit/models/v2/text_encoders/qwen3.py index 328712b..bf87311 100644 --- a/toolkit/models/v2/text_encoders/qwen3.py +++ b/toolkit/models/v2/text_encoders/qwen3.py @@ -1,4 +1,4 @@ -from transformers import Qwen3ForCausalLM +from transformers import Qwen3ForCausalLM, Qwen3Model from .._mixin import OstrisTransformersMixin @@ -15,3 +15,14 @@ class Qwen3TextEncoder(Qwen3ForCausalLM, OstrisTransformersMixin): @classmethod def get_transformer_block_names(cls): return ["model.layers"] + + +class Qwen3ModelEncoder(Qwen3Model, OstrisTransformersMixin): + """The inner Qwen3 base model (anima's text encoder).""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + return ["layers"] diff --git a/toolkit/models/v2/text_encoders/qwen3_vl.py b/toolkit/models/v2/text_encoders/qwen3_vl.py index dd42e22..2657bd1 100644 --- a/toolkit/models/v2/text_encoders/qwen3_vl.py +++ b/toolkit/models/v2/text_encoders/qwen3_vl.py @@ -1,6 +1,10 @@ import torch import torch.nn.functional as F -from transformers import Qwen3VLForConditionalGeneration +from transformers import ( + Qwen3VLForConditionalGeneration, + Qwen3VLModel, + Qwen3VLTextModel, +) from .._mixin import OstrisTransformersMixin @@ -51,3 +55,35 @@ class Qwen3VLTextEncoder(Qwen3VLForConditionalGeneration, OstrisTransformersMixi """Keep the vision tower (reference images ride into the embeddings) but swap its Conv3d patch_embed for an equivalent GEMM.""" return patch_qwen_vl_patch_embed(self) + + +class Qwen3VLModelEncoder(Qwen3VLModel, OstrisTransformersMixin): + """The inner Qwen3-VL base model (what AutoModel resolves): its + last_hidden_state is the instruction feature boogu_image / ideogram4 + consume.""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + return ["language_model.layers"] + + def drop_vision_tower(self): + if getattr(self, "visual", None) is not None: + self.visual = None + return self + + def patch_vision_patch_embed(self) -> int: + return patch_qwen_vl_patch_embed(self) + + +class Qwen3VLTextOnlyEncoder(Qwen3VLTextModel, OstrisTransformersMixin): + """The Qwen3-VL text tower alone (prx_pixel).""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + + @classmethod + def get_transformer_block_names(cls): + return ["layers"] diff --git a/toolkit/models/v2/vae/autoencoder_kl_flux2.py b/toolkit/models/v2/vae/autoencoder_kl_flux2.py new file mode 100644 index 0000000..4357ced --- /dev/null +++ b/toolkit/models/v2/vae/autoencoder_kl_flux2.py @@ -0,0 +1,10 @@ +from diffusers import AutoencoderKLFlux2 + +from .._mixin import OstrisModelMixin + + +class Flux2KLVAE(AutoencoderKLFlux2, OstrisModelMixin): + """The diffusers Flux2 KL VAE (ernie_image; distinct from the hand-rolled + BFL AutoEncoder in vae/flux2_kl.py that flux2/ideogram4 use).""" + + aitk_subfolder = "vae" diff --git a/toolkit/models/v2/vae/flux2_kl.py b/toolkit/models/v2/vae/flux2_kl.py index e7f6b00..fcef0d6 100644 --- a/toolkit/models/v2/vae/flux2_kl.py +++ b/toolkit/models/v2/vae/flux2_kl.py @@ -20,6 +20,8 @@ import torch.utils.checkpoint as ckpt from einops import rearrange from torch import Tensor, nn +from .._mixin import OstrisModelMixin + @dataclass class AutoEncoderParams: @@ -354,7 +356,14 @@ class Decoder(nn.Module): return h -class AutoEncoder(nn.Module): +class AutoEncoder(nn.Module, OstrisModelMixin): + @classmethod + def aitk_config_from_state_dict(cls, state_dict): + # small decoder builds report 96ch at the first decoder block + if state_dict["decoder.up.0.block.0.conv1.bias"].shape[0] == 96: + return AutoEncoderSmallDecoderParams() + return AutoEncoderParams() + def __init__(self, params: AutoEncoderParams): super().__init__() self.params = params diff --git a/toolkit/models/v2/vae/ltx2.py b/toolkit/models/v2/vae/ltx2.py new file mode 100644 index 0000000..f386d2e --- /dev/null +++ b/toolkit/models/v2/vae/ltx2.py @@ -0,0 +1,16 @@ +from diffusers.models.autoencoders import ( + AutoencoderKLLTX2Audio as DiffusersAutoencoderKLLTX2Audio, +) +from diffusers.models.autoencoders import ( + AutoencoderKLLTX2Video as DiffusersAutoencoderKLLTX2Video, +) + +from .._mixin import OstrisModelMixin + + +class LTX2VideoVAE(DiffusersAutoencoderKLLTX2Video, OstrisModelMixin): + aitk_subfolder = "vae" + + +class LTX2AudioVAE(DiffusersAutoencoderKLLTX2Audio, OstrisModelMixin): + aitk_subfolder = "audio_vae" diff --git a/toolkit/models/v2/vae/wan.py b/toolkit/models/v2/vae/wan.py new file mode 100644 index 0000000..2ad242e --- /dev/null +++ b/toolkit/models/v2/vae/wan.py @@ -0,0 +1,9 @@ +from diffusers import AutoencoderKLWan + +from .._mixin import OstrisModelMixin + + +class WanVAE(AutoencoderKLWan, OstrisModelMixin): + """The wan 2.1/2.2 causal video VAE.""" + + aitk_subfolder = "vae" diff --git a/toolkit/models/v2/vision_encoders/clip_vision.py b/toolkit/models/v2/vision_encoders/clip_vision.py new file mode 100644 index 0000000..2a16451 --- /dev/null +++ b/toolkit/models/v2/vision_encoders/clip_vision.py @@ -0,0 +1,13 @@ +from transformers import CLIPVisionModel + +from .._mixin import OstrisTransformersMixin + + +class CLIPVisionEncoder(CLIPVisionModel, OstrisTransformersMixin): + """CLIP vision tower (wan21 i2v image conditioning).""" + + aitk_subfolder = "image_encoder" + + @classmethod + def get_transformer_block_names(cls): + return ["vision_model.encoder.layers"] diff --git a/toolkit/models/wan21/wan21.py b/toolkit/models/wan21/wan21.py index 5e07efc..dc55098 100644 --- a/toolkit/models/wan21/wan21.py +++ b/toolkit/models/wan21/wan21.py @@ -44,6 +44,7 @@ from typing import Any, Callable, Dict, List, Optional, Union from toolkit.models.wan21.wan_lora_convert import convert_to_diffusers, convert_to_original from toolkit.util.quantize import quantize_model from toolkit.models.v2.text_encoders.umt5 import UMT5TextEncoder +from toolkit.models.v2.vae.wan import WanVAE from toolkit.metadata import get_meta_for_safetensors # for generation only? @@ -431,11 +432,9 @@ class Wan21(BaseModel): if self._wan_vae_path is not None: # load the vae from individual repo - vae = AutoencoderKLWan.from_pretrained( - self._wan_vae_path, torch_dtype=dtype).to(dtype=dtype) + vae = WanVAE.load_model(self._wan_vae_path, dtype=dtype, subfolder="") else: - vae = AutoencoderKLWan.from_pretrained( - vae_path, subfolder="vae", torch_dtype=dtype).to(dtype=dtype) + vae = WanVAE.load_model(vae_path, dtype=dtype) flush() self.print_and_status_update("Making pipe") diff --git a/toolkit/models/wan21/wan21_i2v.py b/toolkit/models/wan21/wan21_i2v.py index fbe012f..308820d 100644 --- a/toolkit/models/wan21/wan21_i2v.py +++ b/toolkit/models/wan21/wan21_i2v.py @@ -1,4 +1,5 @@ # WIP, coming soon ish +from toolkit.models.v2.vision_encoders.clip_vision import CLIPVisionEncoder from functools import partial import torch import yaml @@ -314,10 +315,8 @@ class Wan21I2V(Wan21): self.model_config.extras_name_or_path , subfolder="image_processor" ) - self.image_encoder = CLIPVisionModel.from_pretrained( - self.model_config.extras_name_or_path, - subfolder="image_encoder", - torch_dtype=dtype, + self.image_encoder = CLIPVisionEncoder.load_model( + self.model_config.extras_name_or_path, dtype=dtype ) except Exception as e: # load from name_or_path @@ -325,10 +324,8 @@ class Wan21I2V(Wan21): self.model_config.name_or_path_original, subfolder="image_processor" ) - self.image_encoder = CLIPVisionModel.from_pretrained( - self.model_config.name_or_path_original, - subfolder="image_encoder", - torch_dtype=dtype, + self.image_encoder = CLIPVisionEncoder.load_model( + self.model_config.name_or_path_original, dtype=dtype ) self.image_encoder.to(self.device_torch, dtype=dtype) self.image_encoder.eval() diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py index 5421647..8111cd1 100644 --- a/toolkit/stable_diffusion_model.py +++ b/toolkit/stable_diffusion_model.py @@ -101,6 +101,46 @@ DO_NOT_TRAIN_WEIGHTS = [ DeviceStatePreset = Literal['cache_latents', 'generate'] +# diffusers-class name -> v2 wrapper for the legacy archs' components. The +# adoption is an in-place class swap, so pipeline-held references stay valid. +_V2_ADOPTION_MAP = { + "UNet2DConditionModel": ("toolkit.models.v2.diffusion_models.unet", "UNet2DConditionModel"), + "AutoencoderKL": ("toolkit.models.v2.vae.autoencoder_kl", "KLVAE"), + "CLIPTextModel": ("toolkit.models.v2.text_encoders.clip", "CLIPTextEncoder"), + "CLIPTextModelWithProjection": ("toolkit.models.v2.text_encoders.clip", "CLIPTextEncoderWithProjection"), + "T5EncoderModel": ("toolkit.models.v2.text_encoders.t5", "T5TextEncoder"), + "UMT5EncoderModel": ("toolkit.models.v2.text_encoders.umt5", "UMT5TextEncoder"), + "SD3Transformer2DModel": ("toolkit.models.v2.diffusion_models.sd3", "SD3Transformer2DModel"), + "PixArtTransformer2DModel": ("toolkit.models.v2.diffusion_models.pixart", "PixArtTransformer2DModel"), + "Transformer2DModel": ("toolkit.models.v2.diffusion_models.pixart", "Transformer2DModel"), + "AuraFlowTransformer2DModel": ("toolkit.models.v2.diffusion_models.auraflow", "AuraFlowTransformer2DModel"), + "FluxTransformer2DModel": ("toolkit.models.v2.diffusion_models.flux", "FluxTransformer2DModel"), + "Lumina2Transformer2DModel": ("toolkit.models.v2.diffusion_models.lumina2", "Lumina2Transformer2DModel"), + "Gemma2Model": ("toolkit.models.v2.text_encoders.gemma2", "Gemma2ModelEncoder"), +} + + +def _adopt_v2(module): + """Rebind a loaded legacy component onto its v2 wrapper class so every + resident component is an OstrisModelMixin instance the inference engine + can pool and hot-swap. No-op for unknown or already-adopted classes.""" + import importlib + + from toolkit.models.v2._mixin import OstrisModelMixin, adopt_component + + if module is None or isinstance(module, OstrisModelMixin): + return module + entry = _V2_ADOPTION_MAP.get(type(module).__name__) + if entry is None: + return module + try: + wrapper = getattr(importlib.import_module(entry[0]), entry[1]) + return adopt_component(module, wrapper) + except (ImportError, TypeError): + return module + + + class BlankNetwork: def __init__(self): @@ -1048,6 +1088,16 @@ class StableDiffusion: self.unet.requires_grad_(False) self.unet.eval() + # every resident component joins the v2 mixin system (in-place class + # adoption for components the pipeline loaders built directly) + _adopt_v2(self.unet) + _adopt_v2(self.vae) + if isinstance(text_encoder, list): + for te in text_encoder: + _adopt_v2(te) + elif text_encoder is not None: + _adopt_v2(text_encoder) + # load any loras we have if self.model_config.lora_path is not None and not self.is_flux and not self.is_lumina2: pipe.load_lora_weights(self.model_config.lora_path, adapter_name="lora1") diff --git a/toolkit/util/mixed_precision.py b/toolkit/util/mixed_precision.py new file mode 100644 index 0000000..ae59618 --- /dev/null +++ b/toolkit/util/mixed_precision.py @@ -0,0 +1,93 @@ +"""ComfyUI-style handling for mixed-precision checkpoints (e.g. fp32 +scale_shift tables next to bf16 linears in one file). + +Two pieces, used together on a loaded module tree: + +- attach_per_op_casting(root): every weighted module casts its floating-point + inputs to its own weight dtype at forward (what comfy's manual-cast ops do), + so activations promoted to fp32 by a stored-fp32 tensor drop back to the + layer's dtype at the next op instead of erroring in a bf16 matmul. +- pin_stored_fp32(root): tensors that are fp32 after load stay fp32 through + any parent-level dtype cast (`module.to(dtype=...)`, `.half()`, ...) — + device moves still apply. Without this, a holder's blanket `.to(dtype)` + would silently downcast the deliberately-fp32 pieces. +""" + +import types + +import torch +from torch import nn + +from toolkit.util.ostris_quant import OstrisLinear + + +def _cast_args(args, kwargs, dtype): + def cast(t): + if torch.is_tensor(t) and t.is_floating_point() and t.dtype != dtype: + return t.to(dtype) + return t + + return tuple(cast(a) for a in args), {k: cast(v) for k, v in kwargs.items()} + + +def attach_per_op_casting(root: nn.Module) -> int: + """Register forward-pre-hooks casting inputs to each weighted module's own + weight dtype. Covers Linear/Conv/Norm-style modules that own a floating + ``weight`` parameter, and OstrisLinear via its stored orig dtype (its + ``weight`` property would materialize the dequantized tensor). Returns the + number of modules hooked.""" + hooked = 0 + for module in root.modules(): + if isinstance(module, OstrisLinear): + def hook(mod, args, kwargs): + return _cast_args(args, kwargs, mod.ostris_orig_dtype) + + else: + weight = module._parameters.get("weight", None) + if weight is None or not weight.is_floating_point(): + continue + + def hook(mod, args, kwargs): + return _cast_args(args, kwargs, mod._parameters["weight"].dtype) + + module.register_forward_pre_hook(hook, with_kwargs=True) + hooked += 1 + return hooked + + +def _tag_fp32(root: nn.Module): + for t in list(root.parameters()) + list(root.buffers()): + if t.is_floating_point() and t.dtype == torch.float32: + t._pin_dtype = True + if isinstance(t, nn.Parameter): + t.data._pin_dtype = True + + +def pin_stored_fp32(root: nn.Module): + """Mark every currently-fp32 param/buffer in ``root`` and wrap the root's + ``_apply`` so parent dtype casts skip them (device moves still apply). + The wrapper probes the cast function with a scalar to detect whether it + changes float dtypes, so ``.to(device)`` passes through untouched.""" + _tag_fp32(root) + orig_apply = root._apply + + def _apply(self, fn, *args, **kwargs): + probe = fn(torch.zeros((), dtype=torch.float32)) + if probe.dtype == torch.float32: + return orig_apply(fn, *args, **kwargs) + + def fn_pinned(t): + if getattr(t, "_pin_dtype", False): + out = t if t.device == probe.device else t.to(probe.device) + out._pin_dtype = True + return out + return fn(t) + + result = orig_apply(fn_pinned, *args, **kwargs) + # _apply may rewrap tensors (dropping attributes); the invariant is + # simple — the pinned set IS the fp32 set — so re-tag + _tag_fp32(root) + return result + + root._apply = types.MethodType(_apply, root) + return root From 85a6880643eac40e594230a93544ea492b34f3bd Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 28 Aug 2026 08:37:56 -0600 Subject: [PATCH 09/14] Move the offload and quantization to the base model --- .../ace_step/ace_step_15_model.py | 32 +- .../captioner/prompts/ideogram4_prompt.py | 4 +- .../diffusion_models/anima/anima.py | 68 +--- .../boogu_image/boogu_image.py | 30 +- .../diffusion_models/chroma/chroma_model.py | 17 +- .../chroma/chroma_radiance_model.py | 17 +- .../ernie_image/ernie_image.py | 52 +-- .../ernie_image/transformer.py | 3 + .../example_model/example_model.py | 23 +- .../example_model/src/model.py | 24 +- .../diffusion_models/f_light/f_light.py | 22 +- .../flux2/flux2_klein_model.py | 27 +- .../diffusion_models/flux2/flux2_model.py | 54 +-- .../diffusion_models/flux2/src/model.py | 4 + .../flux_kontext/flux_kontext.py | 23 +- .../diffusion_models/hidream/hidream_model.py | 73 +--- .../hidream/hidream_o1_model.py | 23 +- .../diffusion_models/ideogram4/ideogram4.py | 57 +-- .../ideogram4/src/transformer.py | 11 + .../diffusion_models/krea2/krea2.py | 50 +-- .../diffusion_models/krea2/src/mmdit.py | 8 + .../diffusion_models/ltx2/ltx2.py | 121 +----- .../diffusion_models/mageflow/mageflow.py | 48 +-- .../mageflow/src/transformer.py | 20 + .../diffusion_models/minimax_h3/minimax_h3.py | 55 +-- .../nucleus_image/nucleus_image_model.py | 29 +- .../diffusion_models/omnigen2/__init__.py | 34 +- .../prx_pixel_t2i/prx_pixel_t2i.py | 32 +- .../diffusion_models/qwen_image/qwen_image.py | 35 +- .../qwen_image/qwen_image_edit.py | 1 - .../qwen_image/qwen_image_edit_plus.py | 1 - .../diffusion_models/wan22/wan22_14b_model.py | 78 +--- .../diffusion_models/z_image/z_image.py | 43 +- .../z_image/z_image_l2p_model.py | 31 +- .../zeta_chroma/zeta_chroma_model.py | 33 +- .../zeta_chroma/zeta_chroma_transformer.py | 7 + toolkit/models/base_model.py | 66 ++-- toolkit/models/v2/PLANNING.md | 34 ++ toolkit/models/v2/_mixin.py | 130 ++++++ toolkit/models/v2/diffusion_models/ltx2.py | 12 + toolkit/models/v2/diffusion_models/wan.py | 6 + toolkit/models/v2/diffusion_models/z_image.py | 3 + toolkit/models/v2/text_encoders/gemma3.py | 10 + toolkit/models/wan21/wan21.py | 41 +- toolkit/util/quantize.py | 373 +++++++++++------- 45 files changed, 725 insertions(+), 1140 deletions(-) diff --git a/extensions_built_in/audio_models/ace_step/ace_step_15_model.py b/extensions_built_in/audio_models/ace_step/ace_step_15_model.py index 430f82b..e91c444 100644 --- a/extensions_built_in/audio_models/ace_step/ace_step_15_model.py +++ b/extensions_built_in/audio_models/ace_step/ace_step_15_model.py @@ -11,9 +11,7 @@ from toolkit.prompt_utils import PromptEmbeds, concat_prompt_embeds from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) -from toolkit.util.quantize import get_qtype, quantize, quantize_model -from optimum.quanto import freeze from .src.model import ( AceStep15, OobleckVAE, @@ -116,32 +114,22 @@ class AceStep15Model(BaseAudioModel): models = load_models(model_path, device=load_device, dtype=dtype) self.model = models["model"] - - if self.model_config.quantize: - self.print_and_status_update("Quantizing Transformer") - # quantize_model(self, self.model.decoder) - quantize(self.model, weights=get_qtype(self.model_config.qtype)) - freeze(self.model) - flush() - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - self.model.to("cpu") - - + if ( self.model_config.layer_offloading and self.model_config.layer_offloading_transformer_percent > 0 ): raise NotImplementedError("Layer offloading not yet implemented for AceStep15Model") - + + # quantize + offload + placement, all driven by model_config + self.model.aitk_post_load(**self.component_load_kwargs("transformer")) + flush() + self.text_encoder = models["text_encoder"] - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Text Encoder") - quantize(self.text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(self.text_encoder) - flush() + + # quantize + offload + placement, all driven by model_config + self.text_encoder.aitk_post_load(**self.component_load_kwargs("te")) + flush() self.vae = models["vae"] diff --git a/extensions_built_in/captioner/prompts/ideogram4_prompt.py b/extensions_built_in/captioner/prompts/ideogram4_prompt.py index 6cd814b..976b363 100644 --- a/extensions_built_in/captioner/prompts/ideogram4_prompt.py +++ b/extensions_built_in/captioner/prompts/ideogram4_prompt.py @@ -1,4 +1,4 @@ -ideogram4_prompt = """ +ideogram4_prompt = r""" [META] frozen: false description: Slim single-shot magic prompt — splatter planning + v15 output discipline, deduped for faster inference. Thinking off. @@ -15,7 +15,7 @@ You convert a natural-language user idea into a structured JSON caption an image - Emit a SINGLE-LINE MINIFIED JSON object — no markdown fences, no commentary, no other top-level keys. - Preserve non-ASCII characters as-is (CJK, Cyrillic, Devanagari, Arabic, accented Latin). Never escape with `\uNNNN`, transliterate, or replace `café` with `cafe`. -- Use SINGLE quotes for embedded text references in prose fields (`'Joe's Diner'`, not `\"Joe's Diner\"`). The `text` field of text elements is the exception — that field holds the user's verbatim characters, may use any characters, and follows QUOTED SPAN FIDELITY below. +- Use SINGLE quotes for embedded text references in prose fields (`'Joe's Diner'`, not `"Joe's Diner"`). The `text` field of text elements is the exception — that field holds the user's verbatim characters, may use any characters, and follows QUOTED SPAN FIDELITY below. ### Target aspect ratio (input only — never emit it) diff --git a/extensions_built_in/diffusion_models/anima/anima.py b/extensions_built_in/diffusion_models/anima/anima.py index 0bd6c99..c9db405 100644 --- a/extensions_built_in/diffusion_models/anima/anima.py +++ b/extensions_built_in/diffusion_models/anima/anima.py @@ -3,13 +3,11 @@ from typing import List, Optional import torch import yaml -from optimum.quanto import freeze from safetensors.torch import load_file, save_file from toolkit.accelerator import unwrap_model from toolkit.basic import flush from toolkit.config_modules import GenerateImageConfig, ModelConfig -from toolkit.memory_management import MemoryManager from toolkit.models.base_model import BaseModel from toolkit.models.v2.diffusion_models.cosmos import CosmosTransformer3DModel from toolkit.models.v2.text_encoders.anima import AnimaTextConditioner @@ -17,7 +15,6 @@ from toolkit.models.v2.text_encoders.qwen3 import Qwen3ModelEncoder from toolkit.models.v2.vae.qwen_image import QwenImageVAE from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler -from toolkit.util.quantize import get_qtype, quantize, quantize_model try: from diffusers import AnimaAutoBlocks, AnimaModularPipeline @@ -287,62 +284,23 @@ class AnimaModel(BaseModel): transformer = pipe.transformer text_conditioner = pipe.text_conditioner - if self.model_config.quantize: - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) - self.print_and_status_update("Quantizing Text Conditioner") - quantize(text_conditioner, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_conditioner) - flush() + # the text conditioner rides the transformer quantize flag (at qtype_te) + # but takes the text-encoder offload/placement policy + tc_kwargs = self.component_load_kwargs("te") + tc_kwargs["qtype"] = ( + self.model_config.qtype_te if self.model_config.quantize else None + ) + text_conditioner.aitk_post_load(**tc_kwargs) + flush() - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ) - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - pipe.text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - MemoryManager.attach( - text_conditioner, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") - text_conditioner.to("cpu") - else: - transformer.to(self.device_torch, dtype=dtype) - text_conditioner.to(self.device_torch, dtype=dtype) - - pipe.text_encoder.to(self.device_torch, dtype=dtype) + # quantize + offload + placement, all driven by model_config + pipe.text_encoder.aitk_post_load(**self.component_load_kwargs("te")) pipe.text_encoder.requires_grad_(False) pipe.text_encoder.eval() - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Text Encoder") - quantize(pipe.text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(pipe.text_encoder) - flush() - - if self.model_config.low_vram: - pipe.text_encoder.to("cpu") - flush() + flush() self.noise_scheduler = pipe.scheduler self.vae = pipe.vae diff --git a/extensions_built_in/diffusion_models/boogu_image/boogu_image.py b/extensions_built_in/diffusion_models/boogu_image/boogu_image.py index 2cb4fde..24d4d60 100644 --- a/extensions_built_in/diffusion_models/boogu_image/boogu_image.py +++ b/extensions_built_in/diffusion_models/boogu_image/boogu_image.py @@ -36,9 +36,9 @@ from toolkit.models.v2.vae.autoencoder_kl import KLVAE from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) -from toolkit.util.quantize import quantize, get_qtype, quantize_model -from optimum.quanto import freeze, QTensor + +from optimum.quanto import QTensor from diffusers import AutoencoderKL from .src.transformer import BooguImageTransformer2DModel @@ -164,16 +164,8 @@ class BooguImageModel(BaseModel): if attention_backend != "native": transformer.set_attention_backend(attention_backend) - if self.model_config.quantize: - self.print_and_status_update("Quantizing transformer") - quantize_model(self, transformer) - flush() - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") - else: - transformer.to(self.device_torch, dtype=dtype) + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() # --- instruction encoder (Qwen3-VL) + processor --- @@ -203,18 +195,8 @@ class BooguImageModel(BaseModel): ) flush() - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing instruction encoder") - text_encoder.to(self.device_torch) - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - - if self.model_config.low_vram: - self.print_and_status_update("Moving instruction encoder to CPU") - text_encoder.to("cpu") - else: - text_encoder.to(self.device_torch) + # quantize + offload + placement, all driven by model_config + text_encoder.aitk_post_load(**self.component_load_kwargs("te")) flush() # --- VAE (FLUX AutoencoderKL) --- diff --git a/extensions_built_in/diffusion_models/chroma/chroma_model.py b/extensions_built_in/diffusion_models/chroma/chroma_model.py index 6a3aa9a..8d5a11a 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_model.py @@ -13,7 +13,6 @@ from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler from toolkit.accelerator import unwrap_model from optimum.quanto import QTensor -from toolkit.util.quantize import quantize, quantize_model from .pipeline import ChromaPipeline, prepare_latent_image_ids from einops import rearrange, repeat import random @@ -158,27 +157,21 @@ class ChromaModel(BaseModel): transformer = Chroma.load_from_state_dict(load_file(model_path, "cpu"), dtype) # add dtype, not sure why it doesnt have it transformer.dtype = dtype - transformer.to(self.quantize_device, dtype=dtype) transformer.config = FakeConfig() transformer.config.num_layers = transformer.params.depth transformer.config.num_single_layers = transformer.params.depth_single_blocks - if self.model_config.quantize: - # block-streaming quantize (handles dequant-on-save patching, - # excludes, ARA, and quantize_kwargs) - self.print_and_status_update("Quantizing transformer") - quantize_model(self, transformer) - transformer.to(self.device_torch) - else: - transformer.to(self.device_torch, dtype=dtype) + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() self.print_and_status_update("Loading T5") tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path) - text_encoder_2 = T5TextEncoder.load_model(extras_path, dtype=dtype) - self.prepare_text_encoder(text_encoder_2, dtype=dtype) + text_encoder_2 = T5TextEncoder.load( + extras_path, **self.component_load_kwargs("te") + ) # self.print_and_status_update("Loading CLIP") text_encoder = FakeCLIP(device=self.device_torch) diff --git a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py index 2fcc390..6aacd39 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py @@ -12,7 +12,6 @@ from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler from toolkit.accelerator import unwrap_model from optimum.quanto import QTensor -from toolkit.util.quantize import quantize, quantize_model from .pipeline import ChromaPipeline, prepare_latent_image_ids from einops import rearrange, repeat import random @@ -137,27 +136,21 @@ class ChromaRadianceModel(BaseModel): transformer = Chroma.load_model(model_path, dtype=dtype) # add dtype, not sure why it doesnt have it transformer.dtype = dtype - transformer.to(self.quantize_device, dtype=dtype) transformer.config = FakeConfig() transformer.config.num_layers = transformer.params.depth transformer.config.num_single_layers = transformer.params.depth_single_blocks - if self.model_config.quantize: - # block-streaming quantize (handles dequant-on-save patching, - # excludes, ARA, and quantize_kwargs) - self.print_and_status_update("Quantizing transformer") - quantize_model(self, transformer) - transformer.to(self.device_torch) - else: - transformer.to(self.device_torch, dtype=dtype) + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() self.print_and_status_update("Loading T5") tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path) - text_encoder_2 = T5TextEncoder.load_model(extras_path, dtype=dtype) - self.prepare_text_encoder(text_encoder_2, dtype=dtype) + text_encoder_2 = T5TextEncoder.load( + extras_path, **self.component_load_kwargs("te") + ) # self.print_and_status_update("Loading CLIP") text_encoder = FakeCLIP(device=self.device_torch) diff --git a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py index 497207a..393cb0a 100644 --- a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py +++ b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py @@ -13,9 +13,6 @@ from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.accelerator import unwrap_model -from optimum.quanto import freeze -from toolkit.util.quantize import quantize, get_qtype, quantize_model -from toolkit.memory_management import MemoryManager from transformers import AutoTokenizer, AutoModel @@ -88,29 +85,10 @@ class ErnieImageModel(BaseModel): if os.path.exists(te_folder_path): base_model_path = model_path - transformer = ErnieImageTransformer2DModel.load_model(model_path, dtype=dtype) - - if self.model_config.quantize: - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=[ - transformer.x_embedder, - ], - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") + # load + quantize + offload + placement, all driven by model_config + transformer = ErnieImageTransformer2DModel.load( + model_path, **self.component_load_kwargs("transformer") + ) flush() @@ -118,29 +96,11 @@ class ErnieImageModel(BaseModel): tokenizer = AutoTokenizer.from_pretrained( base_model_path, subfolder="tokenizer", torch_dtype=dtype ) - text_encoder = Mistral3ModelEncoder.load_model( - base_model_path, subfolder="text_encoder", dtype=dtype + text_encoder = Mistral3ModelEncoder.load( + base_model_path, subfolder="text_encoder", **self.component_load_kwargs("te") ) - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - - text_encoder.to(self.device_torch, dtype=dtype) flush() - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Text Encoder") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - self.print_and_status_update("Loading VAE") vae = Flux2KLVAE.load_model( base_model_path, dtype=dtype ).to(self.device_torch, dtype=dtype) diff --git a/extensions_built_in/diffusion_models/ernie_image/transformer.py b/extensions_built_in/diffusion_models/ernie_image/transformer.py index ca3a2ae..642de6d 100644 --- a/extensions_built_in/diffusion_models/ernie_image/transformer.py +++ b/extensions_built_in/diffusion_models/ernie_image/transformer.py @@ -298,6 +298,9 @@ class ErnieImageTransformer2DModel(ModelMixin, ConfigMixin, OstrisModelMixin): def get_transformer_block_names(cls): return ["layers"] + def get_offload_ignore_modules(self): + return [self.x_embedder] + @register_to_config def __init__( diff --git a/extensions_built_in/diffusion_models/example_model/example_model.py b/extensions_built_in/diffusion_models/example_model/example_model.py index 51b973e..c13bf75 100644 --- a/extensions_built_in/diffusion_models/example_model/example_model.py +++ b/extensions_built_in/diffusion_models/example_model/example_model.py @@ -32,7 +32,7 @@ from toolkit.models.base_model import BaseModel from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) -from toolkit.util.quantize import quantize, get_qtype, quantize_model +from toolkit.util.quantize import quantize, get_qtype from .src.model import ExampleTransformer2DModel from .src.pipeline import ExamplePipeline, pad_prompt_embeds @@ -165,21 +165,14 @@ class ExampleModel(BaseModel): del state_dict flush() # gc + empty cuda cache; call it after dropping anything big - if self.model_config.quantize: - # quantize_model handles qtype selection, exclusions and device - # juggling, and leaves the model on CPU - self.print_and_status_update("Quantizing transformer") - quantize_model(self, transformer) - flush() - - if self.model_config.low_vram: - # leave it on CPU; get_noise_prediction moves it over when needed - transformer.to("cpu") - else: - transformer.to(self.device_torch, dtype=dtype) + # quantize + offload + placement, all driven by model_config: + # component_load_kwargs derives qtype (incl. an accuracy recovery + # adapter), the layer-offload fraction and the target device + # (low_vram parks on CPU); aitk_post_load applies them. Models whose + # checkpoint sourcing is standard can collapse the build + this into + # one call: ExampleTransformer2DModel.load(path, **kwargs). + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() - # For partial layer offloading support see MemoryManager.attach usage - # in ../ideogram4/ideogram4.py or ../z_image/z_image.py. # --- text encoder + tokenizer (stock transformers model) --- self.print_and_status_update("Loading text encoder") diff --git a/extensions_built_in/diffusion_models/example_model/src/model.py b/extensions_built_in/diffusion_models/example_model/src/model.py index 30fc348..80db21a 100644 --- a/extensions_built_in/diffusion_models/example_model/src/model.py +++ b/extensions_built_in/diffusion_models/example_model/src/model.py @@ -51,6 +51,8 @@ import torch.nn.functional as F from torch import nn from torch.utils.checkpoint import checkpoint +from toolkit.models.v2._mixin import OstrisModelMixin + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: """Standard sinusoidal embedding. @@ -135,8 +137,16 @@ class ExampleTransformerBlock(nn.Module): return x -class ExampleTransformer2DModel(nn.Module): - """The denoiser. Plain ``nn.Module`` on purpose. +class ExampleTransformer2DModel(nn.Module, OstrisModelMixin): + """The denoiser. Plain ``nn.Module`` plus ``OstrisModelMixin``. + + The mixin supplies the universal component API every v2 module shares: + ``load()`` / ``load_model()`` / ``aitk_post_load()`` (quantize, layer + offloading and device placement driven by the holder's model_config via + ``BaseModel.component_load_kwargs(role)``), plus comfy-format save/load. + Override its class hooks (``get_transformer_block_names``, + ``get_quantization_exclude_modules``, ``get_offload_ignore_modules``, + ``convert_state_dict_on_load/save``) as needed. You could also subclass ``diffusers.ModelMixin``/``ConfigMixin`` (see ../../ernie_image/transformer.py) to get ``save_pretrained``, @@ -153,6 +163,12 @@ class ExampleTransformer2DModel(nn.Module): names when deciding where to attach LoRA layers. """ + @classmethod + def get_transformer_block_names(cls): + # attribute name(s) of the repeated-block ModuleList(s); the quantizer + # streams these blocks through the GPU one at a time + return ["blocks"] + def __init__( self, in_channels: int = 16, # VAE latent channels @@ -183,8 +199,8 @@ class ExampleTransformer2DModel(nn.Module): ) # ``blocks`` is the repeated-layer ModuleList. The attribute name is - # what ExampleModel.get_transformer_block_names() returns, which the - # LoRA code uses for block targeting / "transformer only" training. + # what get_transformer_block_names() returns, which the LoRA code uses + # for block targeting and the quantizer uses for block streaming. self.blocks = nn.ModuleList( [ ExampleTransformerBlock(hidden_size, num_heads) diff --git a/extensions_built_in/diffusion_models/f_light/f_light.py b/extensions_built_in/diffusion_models/f_light/f_light.py index 018b97b..647cd88 100644 --- a/extensions_built_in/diffusion_models/f_light/f_light.py +++ b/extensions_built_in/diffusion_models/f_light/f_light.py @@ -13,7 +13,7 @@ from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler from toolkit.accelerator import unwrap_model from optimum.quanto import QTensor -from toolkit.util.quantize import quantize, quantize_model + from .src import FLitePipeline, DiT if TYPE_CHECKING: @@ -76,27 +76,17 @@ class FLiteModel(BaseModel): self.print_and_status_update("Loading transformer") - transformer = DiT.load_model(model_path, dtype=dtype) - - transformer.to(self.quantize_device, dtype=dtype) - - if self.model_config.quantize: - # block-streaming quantize (handles dequant-on-save patching, - # excludes, ARA, and quantize_kwargs) - self.print_and_status_update("Quantizing transformer") - quantize_model(self, transformer) - transformer.to(self.device_torch) - else: - transformer.to(self.device_torch, dtype=dtype) + transformer = DiT.load( + model_path, **self.component_load_kwargs("transformer") + ) flush() self.print_and_status_update("Loading T5") tokenizer = T5TextEncoder.load_tokenizer(extras_path, subfolder="tokenizer") - text_encoder = T5TextEncoder.load_model( - extras_path, dtype=dtype, subfolder="text_encoder" + text_encoder = T5TextEncoder.load( + extras_path, subfolder="text_encoder", **self.component_load_kwargs("te") ) - self.prepare_text_encoder(text_encoder, dtype=dtype) self.noise_scheduler = FLiteModel.get_train_scheduler() diff --git a/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py b/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py index f565abe..505d65e 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py @@ -1,10 +1,7 @@ from .flux2_model import Flux2Model from transformers import Qwen3ForCausalLM, Qwen2Tokenizer from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder -from optimum.quanto import freeze -from toolkit.util.quantize import quantize, get_qtype from toolkit.config_modules import ModelConfig -from toolkit.memory_management.manager import MemoryManager from toolkit.basic import flush from .src.model import Klein9BParams, Klein4BParams @@ -41,27 +38,11 @@ class Flux2KleinModel(Flux2Model): dtype = self.torch_dtype self.print_and_status_update("Loading Qwen3") - text_encoder = Qwen3TextEncoder.load_model( - self.flux2_klein_te_path, dtype=dtype, subfolder="" + # load + quantize + offload + placement, all driven by model_config + text_encoder = Qwen3TextEncoder.load( + self.flux2_klein_te_path, subfolder="", **self.component_load_kwargs("te") ) - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Qwen3") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - elif not self.model_config.low_vram: - text_encoder.to(self.device_torch, dtype=dtype) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) + flush() tokenizer = Qwen2Tokenizer.from_pretrained(self.flux2_klein_te_path) return text_encoder, tokenizer diff --git a/extensions_built_in/diffusion_models/flux2/flux2_model.py b/extensions_built_in/diffusion_models/flux2/flux2_model.py index ec4c106..cc023b4 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_model.py @@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, List, Optional import huggingface_hub import torch from toolkit.config_modules import GenerateImageConfig, ModelConfig -from toolkit.memory_management.manager import MemoryManager from toolkit.metadata import get_meta_for_safetensors from toolkit.models.base_model import BaseModel from toolkit.basic import flush @@ -13,10 +12,8 @@ from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) -from toolkit.dequantize import patch_dequantization_on_save from toolkit.accelerator import unwrap_model -from optimum.quanto import freeze, QTensor -from toolkit.util.quantize import quantize, get_qtype, quantize_model +from optimum.quanto import QTensor from transformers import AutoProcessor, Mistral3ForConditionalGeneration from toolkit.models.v2.text_encoders.mistral3 import Mistral3TextEncoder @@ -100,29 +97,12 @@ class Flux2Model(BaseModel): dtype = self.torch_dtype self.print_and_status_update("Loading Mistral") - text_encoder = Mistral3TextEncoder.load_model( - MISTRAL_PATH, dtype=dtype, subfolder="" + # load + quantize + offload + placement, all driven by model_config + text_encoder = Mistral3TextEncoder.load( + MISTRAL_PATH, subfolder="", **self.component_load_kwargs("te") ) - text_encoder.to(self.device_torch, dtype=dtype) - flush() - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Mistral") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - tokenizer = AutoProcessor.from_pretrained(MISTRAL_PATH) return text_encoder, tokenizer @@ -151,32 +131,10 @@ class Flux2Model(BaseModel): transformer_state_dict, dtype, config=self.get_flux2_params() ) - if self.model_config.quantize: - # patch the state dict method - patch_dequantization_on_save(transformer) - # Avoid full-model peak VRAM allocation before quantization. - self.print_and_status_update("Keeping transformer on CPU for quantization") - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() - else: - transformer.to(self.device_torch, dtype=dtype) + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") - text_encoder, tokenizer = self.load_te() self.print_and_status_update("Loading VAE") diff --git a/extensions_built_in/diffusion_models/flux2/src/model.py b/extensions_built_in/diffusion_models/flux2/src/model.py index 5c04503..ee7dd58 100644 --- a/extensions_built_in/diffusion_models/flux2/src/model.py +++ b/extensions_built_in/diffusion_models/flux2/src/model.py @@ -57,6 +57,10 @@ class FakeConfig: class Flux2(nn.Module, OstrisModelMixin): + @classmethod + def get_transformer_block_names(cls): + return ["double_blocks", "single_blocks"] + def __init__(self, params: Flux2Params): super().__init__() self.config = FakeConfig() diff --git a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py index 535e9fb..a67ba7f 100644 --- a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py +++ b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py @@ -20,7 +20,7 @@ from toolkit.models.flux import add_model_gpu_splitter_to_flux, bypass_flux_guid from toolkit.accelerator import get_accelerator, unwrap_model from optimum.quanto import QTensor from toolkit.util.mask import generate_random_mask, random_dialate_mask -from toolkit.util.quantize import quantize, quantize_model + from einops import rearrange, repeat import random import torch.nn.functional as F @@ -38,8 +38,6 @@ scheduler_config = { "use_dynamic_shifting": True } - - class FluxKontextModel(BaseModel): arch = "flux_kontext" @@ -93,24 +91,17 @@ class FluxKontextModel(BaseModel): base_model_path = model_path self.print_and_status_update("Loading transformer") - transformer = FluxTransformer2DModel.load_model(model_path, dtype=dtype) - transformer.to(self.quantize_device, dtype=dtype) - - if self.model_config.quantize: - # block-streaming quantize (handles dequant-on-save patching, - # excludes, ARA, and quantize_kwargs) - self.print_and_status_update("Quantizing transformer") - quantize_model(self, transformer) - transformer.to(self.device_torch) - else: - transformer.to(self.device_torch, dtype=dtype) + transformer = FluxTransformer2DModel.load( + model_path, **self.component_load_kwargs("transformer") + ) flush() self.print_and_status_update("Loading T5") tokenizer_2 = T5TextEncoder.load_tokenizer(base_model_path) - text_encoder_2 = T5TextEncoder.load_model(base_model_path, dtype=dtype) - self.prepare_text_encoder(text_encoder_2, dtype=dtype) + text_encoder_2 = T5TextEncoder.load( + base_model_path, **self.component_load_kwargs("te") + ) self.print_and_status_update("Loading CLIP") text_encoder = CLIPTextEncoder.load_model( diff --git a/extensions_built_in/diffusion_models/hidream/hidream_model.py b/extensions_built_in/diffusion_models/hidream/hidream_model.py index 9a1925d..01a378c 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_model.py @@ -22,7 +22,6 @@ from toolkit.dequantize import patch_dequantization_on_save from toolkit.accelerator import get_accelerator, unwrap_model from optimum.quanto import freeze, QTensor from toolkit.util.mask import generate_random_mask, random_dialate_mask -from toolkit.util.quantize import quantize, get_qtype from .src.pipelines.hidream_image.pipeline_hidream_image import HiDreamImagePipeline from .src.models.transformers.transformer_hidream_image import HiDreamImageTransformer2DModel from .src.schedulers.fm_solvers_unipc import FlowUniPCMultistepScheduler @@ -97,62 +96,23 @@ class HidreamModel(BaseModel): use_fast=False ) - text_encoder_4 = LlamaTextEncoder.load_model( + # load + quantize + offload + placement, all driven by model_config + text_encoder_4 = LlamaTextEncoder.load( llama_model_path, - dtype=torch.bfloat16, subfolder="", output_hidden_states=True, output_attentions=True, + **self.component_load_kwargs("te"), ) - text_encoder_4.to(self.device_torch, dtype=dtype) - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing llama 8b model") - quantization_type = get_qtype(self.model_config.qtype_te) - quantize(text_encoder_4, weights=quantization_type) - freeze(text_encoder_4) - - if self.low_vram: - # unload it for now - text_encoder_4.to('cpu') - + flush() - + self.print_and_status_update("Loading transformer") - - transformer = self.hidream_transformer_class.load_model( - model_path, dtype=torch.bfloat16 + + transformer = self.hidream_transformer_class.load( + model_path, **self.component_load_kwargs("transformer") ) - - if not self.low_vram: - transformer.to(self.device_torch, dtype=dtype) - - if self.model_config.quantize: - self.print_and_status_update("Quantizing transformer") - quantization_type = get_qtype(self.model_config.qtype) - if self.low_vram: - # move and quantize only certain pieces at a time. - all_blocks = list(transformer.double_stream_blocks) + list(transformer.single_stream_blocks) - self.print_and_status_update(" - quantizing transformer blocks") - for block in tqdm(all_blocks): - block.to(self.device_torch, dtype=dtype) - quantize(block, weights=quantization_type) - freeze(block) - block.to('cpu') - # flush() - - self.print_and_status_update(" - quantizing extras") - transformer.to(self.device_torch, dtype=dtype) - quantize(transformer, weights=quantization_type) - freeze(transformer) - else: - quantize(transformer, weights=quantization_type) - freeze(transformer) - - if self.low_vram: - # unload it for now - transformer.to('cpu') - + flush() self.print_and_status_update("Loading vae") @@ -183,16 +143,11 @@ class HidreamModel(BaseModel): flush() self.print_and_status_update("Loading T5 encoders") - text_encoder_3 = T5TextEncoder.load_model( - extras_path, dtype=torch.bfloat16, subfolder="text_encoder_3" - ).to(self.device_torch, dtype=dtype) - - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing T5") - quantization_type = get_qtype(self.model_config.qtype_te) - quantize(text_encoder_3, weights=quantization_type) - freeze(text_encoder_3) - flush() + # load + quantize + offload + placement, all driven by model_config + text_encoder_3 = T5TextEncoder.load( + extras_path, subfolder="text_encoder_3", **self.component_load_kwargs("te") + ) + flush() tokenizer_3 = T5TextEncoder.load_tokenizer( extras_path, subfolder="tokenizer_3", use_fast=False diff --git a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py index 1b374ec..41bcf0c 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py @@ -16,8 +16,6 @@ from toolkit.samplers.custom_flowmatch_sampler import ( from safetensors.torch import load_file, save_file from toolkit.accelerator import unwrap_model from optimum.quanto import freeze -from toolkit.util.quantize import quantize, get_qtype, quantize_model -from toolkit.memory_management import MemoryManager from transformers import AutoProcessor from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig @@ -203,25 +201,8 @@ class HidreamO1Model(BaseModel): torch_dtype=self.torch_dtype, ) flush() - if not self.model_config.low_vram: - transformer.to(self.device_torch) - - if self.model_config.quantize: - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=[], - ) - + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() # move over to device now if low vram diff --git a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py index 308892d..c57844c 100644 --- a/extensions_built_in/diffusion_models/ideogram4/ideogram4.py +++ b/extensions_built_in/diffusion_models/ideogram4/ideogram4.py @@ -17,9 +17,7 @@ from toolkit.samplers.custom_flowmatch_sampler import ( ) from toolkit.accelerator import unwrap_model from toolkit.metadata import get_meta_for_safetensors -from toolkit.memory_management import MemoryManager -from toolkit.util.quantize import quantize, get_qtype, quantize_model -from optimum.quanto import freeze, QTensor +from optimum.quanto import QTensor import huggingface_hub from huggingface_hub.errors import EntryNotFoundError @@ -365,58 +363,13 @@ class Ideogram4Model(BaseModel): transformer = self._load_transformer(base) - if self.model_config.quantize: - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() - else: - transformer.to(self.device_torch, dtype=dtype) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=[ - transformer.rotary_emb.inv_freq, - transformer.input_proj, - transformer.llm_cond_proj, - ], - ) - elif self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") - else: - # quantize_model leaves the model on CPU; make sure it lands on device. - transformer.to(self.device_torch) + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() tokenizer, text_encoder = self._load_text_encoder(base) - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Text Encoder") - text_encoder.to(self.device_torch) - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - elif self.model_config.low_vram: - self.print_and_status_update("Moving text encoder to CPU") - text_encoder.to("cpu") - else: - self.print_and_status_update("Moving text encoder to device") - text_encoder.to(self.device_torch) + # quantize + offload + placement, all driven by model_config + text_encoder.aitk_post_load(**self.component_load_kwargs("te")) flush() vae = self._load_vae(base) diff --git a/extensions_built_in/diffusion_models/ideogram4/src/transformer.py b/extensions_built_in/diffusion_models/ideogram4/src/transformer.py index 4abc5c8..fce6c43 100644 --- a/extensions_built_in/diffusion_models/ideogram4/src/transformer.py +++ b/extensions_built_in/diffusion_models/ideogram4/src/transformer.py @@ -360,6 +360,17 @@ class Ideogram4FinalLayer(nn.Module): class Ideogram4Transformer2DModel(nn.Module, OstrisModelMixin): """Ideogram 4 flow-matching transformer.""" + @classmethod + def get_transformer_block_names(cls): + return ["layers"] + + def get_offload_ignore_modules(self): + return [ + self.rotary_emb.inv_freq, + self.input_proj, + self.llm_cond_proj, + ] + def __init__(self, config: Ideogram4Config) -> None: super().__init__() self.config = config diff --git a/extensions_built_in/diffusion_models/krea2/krea2.py b/extensions_built_in/diffusion_models/krea2/krea2.py index 90462d9..705744d 100644 --- a/extensions_built_in/diffusion_models/krea2/krea2.py +++ b/extensions_built_in/diffusion_models/krea2/krea2.py @@ -47,7 +47,6 @@ from toolkit.samplers.custom_flowmatch_sampler import ( ) from toolkit.accelerator import unwrap_model from toolkit.metadata import get_meta_for_safetensors -from toolkit.util.quantize import quantize, get_qtype, quantize_model from toolkit.memory_management import MemoryManager from .src.mmdit import ( @@ -383,55 +382,12 @@ class Krea2Model(QwenImageVAEHolderMixin, BaseModel): if self.model_config.qtype == "qfloat8": self.model_config.qtype = "float8" - if self.model_config.quantize: - self.print_and_status_update("Quantizing transformer") - quantize_model(self, transformer) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=[ - module - for module in transformer.modules() - if isinstance(module, (SimpleModulation, DoubleSharedModulation)) - ], - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") - else: - transformer.to(self.device_torch, dtype=dtype) + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() tokenizer, processor, vl_processor, text_encoder = self._load_text_encoder() - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing text encoder") - text_encoder.to(self.device_torch) - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving text encoder to CPU") - text_encoder.to("cpu") - else: - text_encoder.to(self.device_torch) + text_encoder.aitk_post_load(**self.component_load_kwargs("te")) flush() vae = self._load_vae() diff --git a/extensions_built_in/diffusion_models/krea2/src/mmdit.py b/extensions_built_in/diffusion_models/krea2/src/mmdit.py index 5e7cfcb..f4036be 100644 --- a/extensions_built_in/diffusion_models/krea2/src/mmdit.py +++ b/extensions_built_in/diffusion_models/krea2/src/mmdit.py @@ -411,6 +411,14 @@ class SingleStreamBlock(nn.Module): class SingleStreamDiT(nn.Module, OstrisModelMixin): + def get_offload_ignore_modules(self): + # modulation modules hold tiny live state the offloader must not page + return [ + module + for module in self.modules() + if isinstance(module, (SimpleModulation, DoubleSharedModulation)) + ] + def __init__(self, config: SingleMMDiTConfig): super().__init__() self.config = config diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index dc561a3..7e947d0 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -17,9 +17,6 @@ from toolkit.samplers.custom_flowmatch_sampler import ( ) from accelerate import init_empty_weights from toolkit.accelerator import unwrap_model -from optimum.quanto import freeze -from toolkit.util.quantize import quantize, get_qtype, quantize_model -from toolkit.memory_management import MemoryManager from toolkit.paths import MODELS_PATH from toolkit.util.mixed_precision import attach_per_op_casting, pin_stored_fp32 from safetensors.torch import load_file @@ -307,33 +304,8 @@ class LTX2Model(BaseModel): transformer = LTX2VideoTransformer3DModel.load_model(model_path, dtype=dtype) - if self.model_config.quantize: - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - ignore_modules = [] - for block in transformer.transformer_blocks: - ignore_modules.append(block.scale_shift_table) - ignore_modules.append(block.audio_scale_shift_table) - ignore_modules.append(block.video_a2v_cross_attn_scale_shift_table) - ignore_modules.append(block.audio_a2v_cross_attn_scale_shift_table) - ignore_modules.append(transformer.scale_shift_table) - ignore_modules.append(transformer.audio_scale_shift_table) - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=ignore_modules, - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() @@ -441,25 +413,8 @@ class LTX2Model(BaseModel): text_encoder.model.vision_tower = None flush() - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Text Encoder") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ignore_modules=[ - text_encoder.model.language_model.base_model.embed_tokens - ], - ) - + # quantize + offload + placement, all driven by model_config + text_encoder.aitk_post_load(**self.component_load_kwargs("te")) text_encoder.to(self.device_torch, dtype=dtype) flush() @@ -1375,38 +1330,11 @@ class LTX25Model(LTX2Model): pin_stored_fp32(transformer) flush() - if self.model_config.quantize: - if num_quantized_dit: - self.print_and_status_update( - "Transformer is pre-quantized (ConvRot); skipping quantize" - ) - else: - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - ignore_modules = [] - for block in transformer.transformer_blocks: - ignore_modules.append(block.scale_shift_table) - ignore_modules.append(block.audio_scale_shift_table) - ignore_modules.append(block.video_a2v_cross_attn_scale_shift_table) - ignore_modules.append(block.audio_a2v_cross_attn_scale_shift_table) - ignore_modules.append(transformer.scale_shift_table) - ignore_modules.append(transformer.audio_scale_shift_table) - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=ignore_modules, - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") + if num_quantized_dit: + # pre-quantized (ConvRot) comfy file; aitk_post_load skips quantize + transformer.aitk_is_quantized = True + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() self.print_and_status_update("Loading connectors") @@ -1433,32 +1361,11 @@ class LTX25Model(LTX2Model): del te_state_dict flush() - if self.model_config.quantize_te: - if num_quantized_te: - self.print_and_status_update( - "Text encoder is pre-quantized (ConvRot); skipping quantize" - ) - else: - self.print_and_status_update("Quantizing Text Encoder") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - # layer_scalar is a bare tensor buffer on each decoder layer; the - # manager never enumerates it, so it must ride along explicitly - ignore_modules = [text_encoder.embed_tokens] - for layer in text_encoder.layers: - ignore_modules.append(layer.layer_scalar) - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ignore_modules=ignore_modules, - ) + if num_quantized_te: + # pre-quantized (ConvRot) comfy file; aitk_post_load skips quantize + text_encoder.aitk_is_quantized = True + # quantize + offload + placement, all driven by model_config + text_encoder.aitk_post_load(**self.component_load_kwargs("te")) text_encoder.to(self.device_torch) text_encoder.requires_grad_(False) diff --git a/extensions_built_in/diffusion_models/mageflow/mageflow.py b/extensions_built_in/diffusion_models/mageflow/mageflow.py index d54d628..53429b5 100644 --- a/extensions_built_in/diffusion_models/mageflow/mageflow.py +++ b/extensions_built_in/diffusion_models/mageflow/mageflow.py @@ -36,7 +36,6 @@ from safetensors.torch import load_file, save_file import huggingface_hub from transformers import AutoProcessor, AutoTokenizer from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder -from optimum.quanto import freeze from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.models.base_model import BaseModel @@ -47,8 +46,6 @@ from toolkit.samplers.custom_flowmatch_sampler import ( ) from toolkit.accelerator import unwrap_model from toolkit.metadata import get_meta_for_safetensors -from toolkit.util.quantize import quantize, get_qtype, quantize_model -from toolkit.memory_management import MemoryManager from .src.transformer import MageFlow, MageFlowParams from .src.vae import MageVAE @@ -262,50 +259,13 @@ class MageFlowModel(BaseModel): transformer = self._load_transformer() - if self.model_config.quantize: - self.print_and_status_update("Quantizing transformer") - quantize_model(self, transformer) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") - else: - transformer.to(self.device_torch, dtype=dtype) + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() tokenizer, vl_processor, text_encoder = self._load_text_encoder() - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing text encoder") - text_encoder.to(self.device_torch) - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving text encoder to CPU") - text_encoder.to("cpu") - else: - text_encoder.to(self.device_torch) + # quantize + offload + placement, all driven by model_config + text_encoder.aitk_post_load(**self.component_load_kwargs("te")) flush() vae = self._load_vae() diff --git a/extensions_built_in/diffusion_models/mageflow/src/transformer.py b/extensions_built_in/diffusion_models/mageflow/src/transformer.py index 7ce1a76..d2ae0e1 100644 --- a/extensions_built_in/diffusion_models/mageflow/src/transformer.py +++ b/extensions_built_in/diffusion_models/mageflow/src/transformer.py @@ -684,6 +684,26 @@ class MageFlowParams: class MageFlow(nn.Module, OstrisModelMixin): + @classmethod + def get_transformer_block_names(cls): + return ["transformer_blocks"] + + @classmethod + def get_quantization_exclude_modules(cls): + # sensitive modules kept in full precision (fnmatch patterns): + # img_in / txt_in / txt_norm - input projections + # time_text_embed* - timestep embedder feeding every + # block's modulation + # norm_out* / proj_out - final adaptive norm / output projection + return [ + "img_in", + "txt_in", + "txt_norm", + "time_text_embed*", + "norm_out*", + "proj_out", + ] + def __init__(self, params: MageFlowParams): super().__init__() self.params = params diff --git a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py index c9ce2ed..0c7fb71 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py @@ -49,7 +49,6 @@ from toolkit.accelerator import unwrap_model from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds from toolkit.basic import flush from toolkit.config_modules import GenerateImageConfig, ModelConfig -from toolkit.memory_management import MemoryManager from toolkit.metadata import get_meta_for_safetensors from toolkit.models.base_model import BaseModel from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder @@ -64,8 +63,6 @@ from toolkit.util.ostris_quant import OstrisLinear from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) -from toolkit.util.quantize import get_qtype, quantize, quantize_model -from optimum.quanto import freeze from .src import packing @@ -498,54 +495,16 @@ class MinimaxH3Model(BaseModel): if self.model_config.assistant_lora_path is not None: self.load_training_adapter(transformer) - if self.model_config.quantize: - self.print_and_status_update("Quantizing transformer") - quantize_model(self, transformer) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ) - - if self.model_config.low_vram: - self.print_and_status_update("Keeping transformer on CPU") - transformer.to("cpu") - else: - transformer.to(self.device_torch) + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() tokenizer, processor, text_encoder = self._load_text_encoder() - te_prequantized = any( - isinstance(m, OstrisLinear) for m in text_encoder.modules() - ) - if self.model_config.quantize_te and not te_prequantized: - self.print_and_status_update("Quantizing text encoder") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - elif self.model_config.quantize_te: - self.print_and_status_update( - "Text encoder is already nvfp4/int8 quantized; skipping quantize_te" - ) - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - if self.model_config.low_vram: - text_encoder.to("cpu") - else: - text_encoder.to(self.device_torch) + if any(isinstance(m, OstrisLinear) for m in text_encoder.modules()): + # already nvfp4/int8 quantized; aitk_post_load skips quantize_te + text_encoder.aitk_is_quantized = True + # quantize + offload + placement, all driven by model_config + text_encoder.aitk_post_load(**self.component_load_kwargs("te")) flush() vae_bundle = self._load_vaes() diff --git a/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py b/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py index ee80c39..1a39802 100644 --- a/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py +++ b/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py @@ -14,8 +14,6 @@ from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.accelerator import unwrap_model -from toolkit.util.quantize import quantize_model -from toolkit.memory_management import MemoryManager from transformers import Qwen3VLProcessor import torch.nn.functional as F @@ -100,26 +98,8 @@ class NucleusImageModel(QwenImageVAEHolderMixin, BaseModel): if isinstance(m, SwiGLUExperts): m.use_grouped_mm = False - if self.model_config.quantize: - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=[ - ], - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() @@ -127,8 +107,9 @@ class NucleusImageModel(QwenImageVAEHolderMixin, BaseModel): tokenizer = Qwen3VLProcessor.from_pretrained( base_model_path, subfolder="processor", torch_dtype=dtype ) - text_encoder = Qwen3VLTextEncoder.load_model(base_model_path, dtype=dtype) - self.prepare_text_encoder(text_encoder, dtype=dtype) + text_encoder = Qwen3VLTextEncoder.load( + base_model_path, **self.component_load_kwargs("te") + ) self.print_and_status_update("Loading VAE") vae = QwenImageVAE.load_model( diff --git a/extensions_built_in/diffusion_models/omnigen2/__init__.py b/extensions_built_in/diffusion_models/omnigen2/__init__.py index 8887767..58b6be8 100644 --- a/extensions_built_in/diffusion_models/omnigen2/__init__.py +++ b/extensions_built_in/diffusion_models/omnigen2/__init__.py @@ -14,8 +14,6 @@ from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.accelerator import unwrap_model -from optimum.quanto import freeze -from toolkit.util.quantize import quantize, get_qtype, quantize_model from .src.pipelines.omnigen2.pipeline_omnigen2 import OmniGen2Pipeline from .src.models.transformers import OmniGen2Transformer2DModel from .src.models.transformers.repo import OmniGen2RotaryPosEmbed @@ -80,41 +78,19 @@ class OmniGen2Model(BaseModel): extras_path, subfolder="processor", use_fast=True ) - mllm = Qwen25VLTextEncoder.load_model( - extras_path, dtype=torch.bfloat16, subfolder="mllm" + # load + quantize + offload + placement, all driven by model_config + mllm = Qwen25VLTextEncoder.load( + extras_path, subfolder="mllm", **self.component_load_kwargs("te") ) - mllm.to(self.device_torch, dtype=dtype) - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing Qwen2.5 VL model") - quantization_type = get_qtype(self.model_config.qtype_te) - quantize(mllm, weights=quantization_type) - freeze(mllm) - - if self.low_vram: - # unload it for now - mllm.to("cpu") flush() self.print_and_status_update("Loading transformer") - transformer = OmniGen2Transformer2DModel.load_model( - model_path, dtype=torch.bfloat16 + transformer = OmniGen2Transformer2DModel.load( + model_path, **self.component_load_kwargs("transformer") ) - if not self.low_vram: - transformer.to(self.device_torch, dtype=dtype) - - if self.model_config.quantize: - self.print_and_status_update("Quantizing transformer") - quantize_model(self, transformer) - if not self.low_vram: - transformer.to(self.device_torch) - - if self.low_vram: - # unload it for now - transformer.to("cpu") - flush() self.print_and_status_update("Loading vae") diff --git a/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py b/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py index 7358b1b..120d7e1 100644 --- a/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py +++ b/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py @@ -43,7 +43,7 @@ import torch import yaml from transformers import AutoTokenizer, Qwen3VLTextModel -from optimum.quanto import freeze + from toolkit.accelerator import unwrap_model from toolkit.basic import flush @@ -56,7 +56,7 @@ from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.train_tools import apply_noise_offset -from toolkit.util.quantize import quantize, get_qtype, quantize_model + from .src.transformer_prx import PRXTransformer2DModel from .src.pipeline import PRXPixelPipeline @@ -127,36 +127,22 @@ class PRXPixelT2IModel(BaseModel): self.print_and_status_update("Loading transformer") # from_pretrained reads config.json (bottleneck_size, resolution_embeds, # in_channels=3, ...) and the safetensors in one shot. - transformer = PRXTransformer2DModel.load_model(model_path, dtype=dtype) - transformer.to(dtype=dtype) - flush() - - if self.model_config.quantize: - self.print_and_status_update("Quantizing transformer") - quantize_model(self, transformer) - flush() - - if self.model_config.low_vram: - transformer.to("cpu") - else: - transformer.to(self.device_torch, dtype=dtype) + # load + quantize + offload + placement, all driven by model_config + transformer = PRXTransformer2DModel.load( + model_path, **self.component_load_kwargs("transformer") + ) flush() # --- text encoder + tokenizer (Qwen3-VL text tower) --- self.print_and_status_update("Loading text encoder") tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder="tokenizer") - text_encoder = Qwen3VLTextOnlyEncoder.load_model(model_path, dtype=dtype) - text_encoder.to(self.te_device_torch) + text_encoder = Qwen3VLTextOnlyEncoder.load( + model_path, **self.component_load_kwargs("te") + ) text_encoder.eval() text_encoder.requires_grad_(False) flush() - if self.model_config.quantize_te: - self.print_and_status_update("Quantizing text encoder") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - # --- "VAE": identity, since PRXPixel is pixel space --- self.print_and_status_update("Preparing pixel-space VAE (identity)") vae = FakeVAE(scaling_factor=1.0) diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py index f1ad25e..4c4be72 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py @@ -14,7 +14,6 @@ from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.accelerator import get_accelerator, unwrap_model -from toolkit.util.quantize import quantize_model import torch.nn.functional as F from toolkit.memory_management import MemoryManager from toolkit.metadata import get_meta_for_safetensors @@ -102,35 +101,10 @@ class QwenImageModel(QwenImageVAEHolderMixin, BaseModel): if os.path.exists(te_folder_path): base_model_path = model_path - transformer = QwenImageTransformer2DModel.load_model( - model_path, - dtype=model_dtype, - use_comfy_weights=self.model_config.model_kwargs.get( - "use_comfy_weights", True - ), + # load + quantize + offload + placement, all driven by model_config + transformer = QwenImageTransformer2DModel.load( + model_path, **self.component_load_kwargs("transformer") ) - - if self.model_config.quantize and not getattr( - transformer, "aitk_is_quantized", False - ): - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") - flush() self.print_and_status_update("Text Encoder") @@ -138,11 +112,12 @@ class QwenImageModel(QwenImageVAEHolderMixin, BaseModel): text_encoder = Qwen25VLTextEncoder.load_model(base_model_path, dtype=dtype) # remove the visual model as it is not needed for image generation + # (before quantization, so the dead tower is never quantized) self.processor = None if not self._qwen_image_keep_visual: text_encoder.drop_vision_tower() - self.prepare_text_encoder(text_encoder, dtype=dtype) + text_encoder.aitk_post_load(**self.component_load_kwargs("te")) self.print_and_status_update("Loading VAE") vae = QwenImageVAE.load_model(base_model_path, dtype=dtype) diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py index efab9be..3e15cc4 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit.py @@ -15,7 +15,6 @@ from toolkit.samplers.custom_flowmatch_sampler import ( ) from toolkit.accelerator import get_accelerator, unwrap_model from optimum.quanto import freeze, QTensor -from toolkit.util.quantize import quantize, get_qtype, quantize_model import torch.nn.functional as F from diffusers import ( diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py index bb3a9c2..871159c 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image_edit_plus.py @@ -15,7 +15,6 @@ from toolkit.samplers.custom_flowmatch_sampler import ( ) from toolkit.accelerator import get_accelerator, unwrap_model from optimum.quanto import freeze, QTensor -from toolkit.util.quantize import quantize, get_qtype, quantize_model import torch.nn.functional as F from diffusers import ( diff --git a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py index 97ccd52..fa43a58 100644 --- a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py +++ b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py @@ -286,62 +286,25 @@ class Wan2214bModel(Wan21): # we have a hf path, replace it with transformer_2 subfolder subfolder_2 = "transformer_2" + # per-transformer load kwargs; the ARA (if any) applies to the combined + # dual model below, offload attaches per transformer after that + per_kwargs = self.component_load_kwargs("transformer") + per_kwargs["offload"] = 0.0 + if self.model_config.accuracy_recovery_adapter is not None: + per_kwargs["qtype"] = None + self.print_and_status_update("Loading transformer 1") - dtype = self.torch_dtype - transformer_1 = WanTransformer3DModel.load_model( - transformer_path_1, dtype=dtype, subfolder=subfolder_1 + transformer_1 = WanTransformer3DModel.load( + transformer_path_1, subfolder=subfolder_1, **per_kwargs ) - flush() - if self.model_config.low_vram: - # quantize on the device - transformer_1.to('cpu', dtype=dtype) - flush() - else: - transformer_1.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize and self.model_config.accuracy_recovery_adapter is None: - # todo handle two ARAs - self.print_and_status_update("Quantizing Transformer 1") - quantize_model(self, transformer_1) - flush() - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer 1 to CPU") - transformer_1.to("cpu") - else: - transformer_1.to(self.device_torch) - self.print_and_status_update("Loading transformer 2") - dtype = self.torch_dtype - transformer_2 = WanTransformer3DModel.load_model( - transformer_path_2, dtype=dtype, subfolder=subfolder_2 + transformer_2 = WanTransformer3DModel.load( + transformer_path_2, subfolder=subfolder_2, **per_kwargs ) - flush() - if self.model_config.low_vram: - # quantize on the device - transformer_2.to('cpu', dtype=dtype) - flush() - else: - transformer_2.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize and self.model_config.accuracy_recovery_adapter is None: - # todo handle two ARAs - self.print_and_status_update("Quantizing Transformer 2") - quantize_model(self, transformer_2) - flush() - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer 2 to CPU") - transformer_2.to("cpu") - else: - transformer_2.to(self.device_torch) - layer_offloading_transformer = self.model_config.layer_offloading and self.model_config.layer_offloading_transformer_percent > 0 # make the combined model self.print_and_status_update("Creating DualWanTransformer3DModel") @@ -362,18 +325,13 @@ class Wan2214bModel(Wan21): if layer_offloading_transformer: - MemoryManager.attach( - transformer_1, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=[transformer_1.scale_shift_table] + [block.scale_shift_table for block in transformer_1.blocks] - ) - MemoryManager.attach( - transformer_2, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=[transformer_2.scale_shift_table] + [block.scale_shift_table for block in transformer_2.blocks] - ) + for t in (transformer_1, transformer_2): + MemoryManager.attach( + t, + self.device_torch, + offload_percent=self.model_config.layer_offloading_transformer_percent, + ignore_modules=t.get_offload_ignore_modules(), + ) return transformer 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 bc2a30e..7ffc823 100644 --- a/extensions_built_in/diffusion_models/z_image/z_image.py +++ b/extensions_built_in/diffusion_models/z_image/z_image.py @@ -12,7 +12,6 @@ from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.accelerator import unwrap_model -from toolkit.util.quantize import quantize_model from toolkit.memory_management import MemoryManager from toolkit.metadata import get_meta_for_safetensors from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder @@ -182,19 +181,9 @@ class ZImageModel(BaseModel): # 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 - 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, use_comfy_weights=self.model_config.model_kwargs.get( "use_comfy_weights", True @@ -223,37 +212,15 @@ class ZImageModel(BaseModel): if self.model_config.qtype == "qfloat8": self.model_config.qtype = "float8" - # 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() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=[ - transformer.x_pad_token, - transformer.cap_pad_token, - ], - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") - + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() self.print_and_status_update("Text Encoder") tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path) - text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype) - self.prepare_text_encoder(text_encoder, dtype=dtype) + text_encoder = Qwen3TextEncoder.load( + base_model_path, **self.component_load_kwargs("te") + ) self.print_and_status_update("Loading VAE") vae = KLVAE.load_model(base_model_path, dtype=dtype) diff --git a/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py b/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py index 289ca1a..b642935 100644 --- a/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py +++ b/extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py @@ -9,8 +9,6 @@ import torch.nn.functional as F import yaml from toolkit.basic import flush from toolkit.accelerator import unwrap_model -from toolkit.util.quantize import quantize_model -from toolkit.memory_management import MemoryManager from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder from toolkit.models.FakeVAE import FakeVAE @@ -441,35 +439,16 @@ class ZImageL2PModel(ZImageModel): if self.model_config.qtype == "qfloat8": self.model_config.qtype = "float8" - if self.model_config.quantize: - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=[ - transformer.x_pad_token, - transformer.cap_pad_token, - ], - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() self.print_and_status_update("Text Encoder") tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path) - text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype) - self.prepare_text_encoder(text_encoder, dtype=dtype) + text_encoder = Qwen3TextEncoder.load( + base_model_path, **self.component_load_kwargs("te") + ) self.print_and_status_update("Loading VAE") # vae = AutoencoderKL.from_pretrained( diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py index d9af932..fcc2ce3 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py @@ -13,8 +13,6 @@ from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.accelerator import unwrap_model -from toolkit.util.quantize import quantize_model -from toolkit.memory_management import MemoryManager from safetensors.torch import load_file from optimum.quanto import QTensor from toolkit.metadata import get_meta_for_safetensors @@ -108,37 +106,16 @@ class ZetaChromaModel(BaseModel): ) del transformer_state_dict - transformer.to(self.quantize_device, dtype=dtype) - - if self.model_config.quantize: - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() - - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_transformer_percent > 0 - ): - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent, - ignore_modules=[ - transformer.x_pad_token, - transformer.cap_pad_token, - ], - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to("cpu") + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) flush() self.print_and_status_update("Text Encoder") tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path) - text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype) - self.prepare_text_encoder(text_encoder, dtype=dtype) + text_encoder = Qwen3TextEncoder.load( + base_model_path, **self.component_load_kwargs("te") + ) self.print_and_status_update("Loading VAE") vae = FakeVAE(scaling_factor=1.0) diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py index 03995b1..8a0a5f2 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py @@ -452,6 +452,13 @@ class SimpleMLPAdaLN(nn.Module): class ZImageDCT(nn.Module, OstrisModelMixin): + @classmethod + def get_transformer_block_names(cls): + return ["layers"] + + def get_offload_ignore_modules(self): + return [self.x_pad_token, self.cap_pad_token] + def __init__(self, params: ZImageDCTParams): super().__init__() self.config = FakeConfig() diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index 8f06dfb..131647b 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -1628,37 +1628,41 @@ class BaseModel: else: self.text_encoder.to(*args, **kwargs) - def prepare_text_encoder(self, text_encoder, dtype=None): - """Standard post-load text-encoder policy: layer offloading, device - placement, then quantize_te. Skips quantization when the checkpoint - loaded pre-quantized.""" - from optimum.quanto import freeze - from toolkit.memory_management import MemoryManager - from toolkit.util.quantize import get_qtype, quantize - - dtype = dtype if dtype is not None else self.torch_dtype - if ( - self.model_config.layer_offloading - and self.model_config.layer_offloading_text_encoder_percent > 0 - ): - MemoryManager.attach( - text_encoder, - self.device_torch, - offload_percent=self.model_config.layer_offloading_text_encoder_percent, - ) - - text_encoder.to(self.device_torch, dtype=dtype) - flush() - - if self.model_config.quantize_te and not getattr( - text_encoder, "aitk_is_quantized", False - ): - self.print_and_status_update("Quantizing Text Encoder") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) - freeze(text_encoder) - flush() - return text_encoder - + def component_load_kwargs(self, role: str = "transformer", dtype=None): + """kwargs for a v2 module's .load()/.aitk_post_load(), derived from + model_config: qtype (with the accuracy recovery adapter recombined), + offload fraction, devices, low_vram placement. Roles: "transformer", + "te", "vae".""" + mc = self.model_config + qtype, offload = None, 0.0 + if role == "transformer": + if mc.quantize: + qtype = mc.qtype + if mc.accuracy_recovery_adapter and "|" not in (qtype or ""): + qtype = f"{qtype}|{mc.accuracy_recovery_adapter}" + if mc.layer_offloading: + offload = mc.layer_offloading_transformer_percent + elif role == "te": + if mc.quantize_te: + qtype = mc.qtype_te + if mc.layer_offloading: + offload = mc.layer_offloading_text_encoder_percent + if dtype is None: + dtype = self.vae_torch_dtype if role == "vae" else self.torch_dtype + device = self.te_device_torch if role == "te" else self.device_torch + if mc.low_vram and role in ("transformer", "te"): + device = "cpu" + elif role == "vae": + device = self.vae_device_torch + return dict( + qtype=qtype, + offload=offload, + dtype=dtype, + device=device, + quantize_device=self.device_torch, + base_model=self, + use_comfy_weights=mc.model_kwargs.get("use_comfy_weights", True), + ) def convert_lora_weights_before_save(self, state_dict): # can be overridden in child classes to convert weights before saving diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index 52a509e..ca0d6da 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -408,6 +408,40 @@ a legacy scheduler fallback, and the sampler-name pass-through. Full monolith decomposition (per-arch v2 loading with comfy candidates) remains future work, but every resident component is now poolable by the engine. +## Unified load API (2026-08-28) + +Loading policy moved out of the holders into the mixin. The surface: + +- `ModelClassName.load(name_or_path, qtype=..., offload=..., dtype=..., + device=..., ...)` — sourcing (`load_model`) + `aitk_post_load` in one call. +- `module.aitk_post_load(**kwargs)` — the post-load half alone, for holders + whose checkpoint sourcing is custom (state-dict surgery, combined files). + Handles: qtype (incl. `"qtype|ara_path"` accuracy recovery adapters, with + prequantized-skip via `aitk_is_quantized`), block-streamed quantization + (`quantize_module` in toolkit/util/quantize.py), MemoryManager layer + offloading with per-class `get_offload_ignore_modules()` hooks, and device + placement (low_vram parks on cpu). +- `BaseModel.component_load_kwargs(role)` derives the kwargs from + model_config for roles "transformer" / "te" / "vae" (qtype/qtype_te, ARA + recombination, offload percents, low_vram, te/vae devices). + +All holders converted (transformer + TE quantize/offload/placement blocks +deleted): krea2, z_image, qwen_image, chroma ×2, flux_kontext, f_light, +nucleus, zeta_chroma, z_image_l2p, wan21 (+ subclasses), wan22_14b (dual: +per-transformer load, dual-ARA branch kept), boogu, ernie, ideogram4, +mageflow, prx_pixel, omnigen2, flux2 + klein, hidream + hidream_o1, ltx2 + +ltx2.5 (prequantized ConvRot flagged via aitk_is_quantized), minimax_h3, +anima (conditioner rides the transformer quantize flag at qtype_te), +ace_step_15, example_model (template now teaches the new API). +New `get_offload_ignore_modules` hooks: wan (scale_shift tables — also fixes +wan21 offload which previously offloaded them), ltx2 (all four per-block +tables), gemma3 (embed_tokens), gemma4 (embed_tokens + layer_scalar), +ideogram4 (rotary inv_freq + input/cond projections), ernie (x_embedder), +zeta ZImageDCT (pad tokens). New block-name hooks where only the holder had +them: Flux2, MageFlow, ZImageDCT, Ideogram4Transformer2DModel (+ its exclude +list on MageFlow). `prepare_text_encoder` remains only as the legacy-monolith +path. + ## Testing - [x] `testing/test_model_loading.py`: per-arch load + one small sample through diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index c2e9699..def4cbb 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -122,6 +122,11 @@ class OstrisModelMixin: """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. @@ -152,6 +157,131 @@ class OstrisModelMixin: # 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, + **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) + + 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 [] + ) + 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) + else: + status_fn(f"Quantizing ({qtype})") + quantize_module( + self, + qtype, + device=quantize_device if quantize_device is not None else 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, + ) + self.aitk_is_quantized = True + self.aitk_qtype = qtype + elif qtype and getattr(self, "aitk_is_quantized", False): + status_fn("Checkpoint is pre-quantized; skipping quantization") + + if offload and offload > 0: + from toolkit.memory_management import MemoryManager + + MemoryManager.attach( + self, + 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, diff --git a/toolkit/models/v2/diffusion_models/ltx2.py b/toolkit/models/v2/diffusion_models/ltx2.py index 2657336..9cd4aa3 100644 --- a/toolkit/models/v2/diffusion_models/ltx2.py +++ b/toolkit/models/v2/diffusion_models/ltx2.py @@ -21,6 +21,18 @@ class LTX2VideoTransformer3DModel( def get_transformer_block_names(cls): return ["transformer_blocks"] + def get_offload_ignore_modules(self): + # fp32 scale/shift tables (bare tensors) must stay resident + mods = [] + for block in self.transformer_blocks: + mods += [ + block.scale_shift_table, + block.audio_scale_shift_table, + block.video_a2v_cross_attn_scale_shift_table, + block.audio_a2v_cross_attn_scale_shift_table, + ] + return mods + [self.scale_shift_table, self.audio_scale_shift_table] + class LTX2TextConnectors(DiffusersLTX2TextConnectors, OstrisModelMixin): aitk_subfolder = "connectors" diff --git a/toolkit/models/v2/diffusion_models/wan.py b/toolkit/models/v2/diffusion_models/wan.py index 67f62b6..d71cdaf 100644 --- a/toolkit/models/v2/diffusion_models/wan.py +++ b/toolkit/models/v2/diffusion_models/wan.py @@ -80,6 +80,12 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin): def get_transformer_block_names(cls): return ["blocks"] + def get_offload_ignore_modules(self): + # fp32 modulation tables must stay resident on the compute device + return [self.scale_shift_table] + [ + block.scale_shift_table for block in self.blocks + ] + @classmethod def convert_state_dict_on_load(cls, state_dict): # original/comfy wan keys -> diffusers layout via diffusers' own diff --git a/toolkit/models/v2/diffusion_models/z_image.py b/toolkit/models/v2/diffusion_models/z_image.py index 794c3ec..e77990b 100644 --- a/toolkit/models/v2/diffusion_models/z_image.py +++ b/toolkit/models/v2/diffusion_models/z_image.py @@ -26,6 +26,9 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix def get_transformer_block_names(cls): return ["layers"] + def get_offload_ignore_modules(self): + return [self.x_pad_token, self.cap_pad_token] + @classmethod def get_quantization_exclude_modules(cls): # sensitive modules kept in full precision (fnmatch patterns on module diff --git a/toolkit/models/v2/text_encoders/gemma3.py b/toolkit/models/v2/text_encoders/gemma3.py index 8391fb3..8e6bd1f 100644 --- a/toolkit/models/v2/text_encoders/gemma3.py +++ b/toolkit/models/v2/text_encoders/gemma3.py @@ -14,6 +14,9 @@ class Gemma3TextEncoder(Gemma3ForConditionalGeneration, OstrisTransformersMixin) # both layouts seen across transformers versions; missing paths skip return ["model.language_model.layers", "language_model.model.layers"] + def get_offload_ignore_modules(self): + return [self.model.language_model.base_model.embed_tokens] + try: from transformers.models.gemma4.modeling_gemma4 import Gemma4TextModel @@ -28,5 +31,12 @@ try: def get_transformer_block_names(cls): return ["layers"] + def get_offload_ignore_modules(self): + # layer_scalar is a bare tensor buffer on each decoder layer; the + # manager never enumerates it, so it must ride along explicitly + return [self.embed_tokens] + [ + layer.layer_scalar for layer in self.layers + ] + except ImportError: Gemma4TextEncoder = None diff --git a/toolkit/models/wan21/wan21.py b/toolkit/models/wan21/wan21.py index dc55098..73fde4c 100644 --- a/toolkit/models/wan21/wan21.py +++ b/toolkit/models/wan21/wan21.py @@ -6,7 +6,6 @@ from toolkit.accelerator import unwrap_model from toolkit.basic import flush from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.dequantize import patch_dequantization_on_save -from toolkit.memory_management.manager import MemoryManager from toolkit.models.base_model import BaseModel from toolkit.prompt_utils import PromptEmbeds from transformers import AutoTokenizer, UMT5EncoderModel @@ -42,7 +41,6 @@ from diffusers.pipelines.wan.pipeline_wan import XLA_AVAILABLE from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback from typing import Any, Callable, Dict, List, Optional, Union from toolkit.models.wan21.wan_lora_convert import convert_to_diffusers, convert_to_original -from toolkit.util.quantize import quantize_model from toolkit.models.v2.text_encoders.umt5 import UMT5TextEncoder from toolkit.models.v2.vae.wan import WanVAE from toolkit.metadata import get_meta_for_safetensors @@ -344,23 +342,11 @@ class Wan21(BaseModel): def load_wan_transformer(self, transformer_path, subfolder=None): self.print_and_status_update("Loading transformer") - dtype = self.torch_dtype - transformer = WanTransformer3DModel.load_model( - transformer_path, dtype=dtype, subfolder=subfolder - ) if self.model_config.split_model_over_gpus: raise ValueError( "Splitting model over gpus is not supported for Wan2.1 models") - if self.model_config.low_vram: - # quantize on the device - transformer.to('cpu', dtype=dtype) - flush() - else: - transformer.to(self.device_torch, dtype=dtype) - flush() - if self.model_config.assistant_lora_path is not None or self.model_config.inference_lora_path is not None: raise ValueError( "Assistant LoRA is not supported for Wan2.1 models currently") @@ -369,23 +355,13 @@ class Wan21(BaseModel): raise ValueError( "Loading LoRA is not supported for Wan2.1 models currently") + # load + quantize + offload + placement, all driven by model_config + transformer = WanTransformer3DModel.load( + transformer_path, + subfolder=subfolder, + **self.component_load_kwargs("transformer"), + ) flush() - - if self.model_config.quantize: - self.print_and_status_update("Quantizing Transformer") - quantize_model(self, transformer) - flush() - - if self.model_config.layer_offloading and self.model_config.layer_offloading_transformer_percent > 0: - MemoryManager.attach( - transformer, - self.device_torch, - offload_percent=self.model_config.layer_offloading_transformer_percent - ) - - if self.model_config.low_vram: - self.print_and_status_update("Moving transformer to CPU") - transformer.to('cpu') return transformer @@ -418,8 +394,9 @@ class Wan21(BaseModel): self.print_and_status_update("Loading UMT5EncoderModel") tokenizer = UMT5TextEncoder.load_tokenizer(te_path) - text_encoder = UMT5TextEncoder.load_model(te_path, dtype=dtype) - self.prepare_text_encoder(text_encoder, dtype=dtype) + text_encoder = UMT5TextEncoder.load( + te_path, **self.component_load_kwargs("te") + ) if self.model_config.low_vram: print("Moving transformer back to GPU") diff --git a/toolkit/util/quantize.py b/toolkit/util/quantize.py index db54f58..1f820a1 100644 --- a/toolkit/util/quantize.py +++ b/toolkit/util/quantize.py @@ -261,6 +261,68 @@ def _has_quantizable_linear(module: torch.nn.Module, weights, exclude=None) -> b return False +@torch.no_grad() +def quantize_module( + module: torch.nn.Module, + qtype: str, + device=None, + dtype: torch.dtype = torch.bfloat16, + block_names: Optional[List[str]] = None, + exclude: Optional[List[str]] = None, + quantize_kwargs: Optional[dict] = None, + status_fn=print_acc, +): + """Module-centric quantization: block-streamed (each repeated block moves + to ``device`` for the math and returns to cpu) with a whole-module pass + for the extras. This is the core the per-model loaders call; holders' + quantize_model wraps it with model_config plumbing.""" + from toolkit.dequantize import patch_dequantization_on_save + + patch_dequantization_on_save(module) + quantization_type = get_qtype(qtype) + exclude = list(exclude or []) + quantize_kwargs = quantize_kwargs or {} + + all_blocks: List[torch.nn.Module] = [] + for name in block_names or []: + # name may be a dotted path for models that nest their blocks + block_list = module + for part in name.split("."): + block_list = getattr(block_list, part, None) + if block_list is None: + break + if block_list is not None: + all_blocks += list(block_list) + if all_blocks: + status_fn(f" - quantizing {len(all_blocks)} blocks") + already_quantized = 0 + for block in tqdm(all_blocks): + if not _has_quantizable_linear(block, quantization_type, exclude): + # pre-quantized checkpoint with a matching qtype: nothing in this + # block would change — skip the device round-trip and the dtype + # cast entirely so the load stays byte-identical + already_quantized += 1 + continue + if device is not None: + block.to(device, dtype=dtype, non_blocking=True) + quantize(block, weights=quantization_type, exclude=exclude, **quantize_kwargs) + freeze(block) + # NOT non_blocking: an async D2H allocates the cpu destination in pinned + # memory, which the caching host allocator keeps forever — that silently + # retained a model-sized chunk of host ram + if device is not None: + block.to("cpu") + if already_quantized: + status_fn( + f" - {already_quantized} blocks already quantized with a matching qtype; left untouched" + ) + + status_fn(" - quantizing extras") + quantize(module, weights=quantization_type, exclude=exclude, **quantize_kwargs) + freeze(module) + return module + + def quantize_model( base_model: "BaseModel", model_to_quantize: torch.nn.Module, @@ -279,153 +341,176 @@ def quantize_model( exclude_modules = base_model.get_quantization_exclude_modules() or [] if base_model.model_config.accuracy_recovery_adapter is not None: - from toolkit.config_modules import NetworkConfig - from toolkit.lora_special import LoRASpecialNetwork - - # we need to load and quantize with an accuracy recovery adapter - # todo handle hf repos - load_lora_path = base_model.model_config.accuracy_recovery_adapter - - if not os.path.exists(load_lora_path): - # not local file, grab from the hub - - path_split = load_lora_path.split("/") - if len(path_split) > 3: - raise ValueError( - "The accuracy recovery adapter path must be a local path or for a hf repo, 'username/repo_name/filename.safetensors'." - ) - repo_id = f"{path_split[0]}/{path_split[1]}" - print_acc(f"Grabbing lora from the hub: {load_lora_path}") - new_lora_path = hf_hub_download( - repo_id, - filename=path_split[-1], - ) - # replace the path - load_lora_path = new_lora_path - - # build the lora config based on the lora weights - lora_state_dict = load_file(load_lora_path) - - if hasattr(base_model, "convert_lora_weights_before_load"): - lora_state_dict = base_model.convert_lora_weights_before_load(lora_state_dict) - - network_config = { - "type": "lora", - "network_kwargs": {"only_if_contains": []}, - "transformer_only": False, - } - first_key = list(lora_state_dict.keys())[0] - first_weight = lora_state_dict[first_key] - # if it starts with lycoris and includes lokr - if first_key.startswith("lycoris") and any( - "lokr" in key for key in lora_state_dict.keys() - ): - network_config["type"] = "lokr" - - network_kwargs = {} - - # find firse loraA weight - if network_config["type"] == "lora": - linear_dim = None - for key, value in lora_state_dict.items(): - if "lora_A" in key: - linear_dim = int(value.shape[0]) - break - linear_alpha = linear_dim - network_config["linear"] = linear_dim - network_config["linear_alpha"] = linear_alpha - - # we build the keys to match every key - only_if_contains = [] - for key in lora_state_dict.keys(): - contains_key = key.split(".lora_")[0] - if contains_key not in only_if_contains: - only_if_contains.append(contains_key) - - network_kwargs["only_if_contains"] = only_if_contains - elif network_config["type"] == "lokr": - # find the factor - largest_factor = 0 - for key, value in lora_state_dict.items(): - if "lokr_w1" in key: - factor = int(value.shape[0]) - if factor > largest_factor: - largest_factor = factor - network_config["lokr_full_rank"] = True - network_config["lokr_factor"] = largest_factor - - only_if_contains = [] - for key in lora_state_dict.keys(): - if "lokr_w1" in key: - contains_key = key.split(".lokr_w1")[0] - contains_key = contains_key.replace("lycoris_", "") - if contains_key not in only_if_contains: - only_if_contains.append(contains_key) - network_kwargs["only_if_contains"] = only_if_contains - - if hasattr(base_model, 'target_lora_modules'): - network_kwargs['target_lin_modules'] = base_model.target_lora_modules - - # todo auto grab these - # get dim and scale - network_config = NetworkConfig(**network_config) - - network = LoRASpecialNetwork( - text_encoder=None, - unet=model_to_quantize, - lora_dim=network_config.linear, - multiplier=1.0, - alpha=network_config.linear_alpha, - # conv_lora_dim=self.network_config.conv, - # conv_alpha=self.network_config.conv_alpha, - train_unet=True, - train_text_encoder=False, - network_config=network_config, - network_type=network_config.type, - transformer_only=network_config.transformer_only, - is_transformer=base_model.is_transformer, - base_model=base_model, - is_ara=True, - **network_kwargs - ) - network.apply_to( - None, model_to_quantize, apply_text_encoder=False, apply_unet=True - ) - network.force_to(base_model.device_torch, dtype=base_model.torch_dtype) - network._update_torch_multiplier() - network.load_weights(lora_state_dict) - network.eval() - network.is_active = True - network.can_merge_in = False - base_model.accuracy_recovery_adapter = network - - # quantize it - lora_exclude_modules = [] - quantization_type = get_qtype(base_model.model_config.qtype) - for lora_module in tqdm(network.unet_loras, desc="Attaching quantization"): - # the lora has already hijacked the original module - orig_module = lora_module.org_module[0] - orig_module.to(base_model.torch_dtype) - # make the params not require gradients - for param in orig_module.parameters(): - param.requires_grad = False - quantize(orig_module, weights=quantization_type) - freeze(orig_module) - module_name = lora_module.lora_name.replace('$$', '.').replace('transformer.', '') - lora_exclude_modules.append(module_name) - if base_model.model_config.low_vram: - # move it back to cpu - orig_module.to("cpu") - pass - # quantize additional layers - print_acc(" - quantizing additional layers") - quantization_type = get_qtype('uint8') - quantize( + attach_ara_and_quantize( + base_model, model_to_quantize, - weights=quantization_type, - exclude=lora_exclude_modules + exclude_modules + ara_path=base_model.model_config.accuracy_recovery_adapter, + exclude=exclude_modules, ) else: + _quantize_model_blocks(base_model, model_to_quantize, exclude_modules) + + +@torch.no_grad() +def attach_ara_and_quantize( + base_model: "BaseModel", + model_to_quantize: torch.nn.Module, + ara_path: str, + exclude: Optional[List[str]] = None, +): + """Load an accuracy recovery adapter as a live network on the module and + quantize around it (adapter-hijacked linears at the configured qtype, + everything else uint8). The network lands on + base_model.accuracy_recovery_adapter.""" + from toolkit.config_modules import NetworkConfig + from toolkit.lora_special import LoRASpecialNetwork + + exclude_modules = list(exclude or []) + load_lora_path = ara_path + + if not os.path.exists(load_lora_path): + # not local file, grab from the hub + + path_split = load_lora_path.split("/") + if len(path_split) > 3: + raise ValueError( + "The accuracy recovery adapter path must be a local path or for a hf repo, 'username/repo_name/filename.safetensors'." + ) + repo_id = f"{path_split[0]}/{path_split[1]}" + print_acc(f"Grabbing lora from the hub: {load_lora_path}") + new_lora_path = hf_hub_download( + repo_id, + filename=path_split[-1], + ) + # replace the path + load_lora_path = new_lora_path + + # build the lora config based on the lora weights + lora_state_dict = load_file(load_lora_path) + + if hasattr(base_model, "convert_lora_weights_before_load"): + lora_state_dict = base_model.convert_lora_weights_before_load(lora_state_dict) + + network_config = { + "type": "lora", + "network_kwargs": {"only_if_contains": []}, + "transformer_only": False, + } + first_key = list(lora_state_dict.keys())[0] + first_weight = lora_state_dict[first_key] + # if it starts with lycoris and includes lokr + if first_key.startswith("lycoris") and any( + "lokr" in key for key in lora_state_dict.keys() + ): + network_config["type"] = "lokr" + + network_kwargs = {} + + # find firse loraA weight + if network_config["type"] == "lora": + linear_dim = None + for key, value in lora_state_dict.items(): + if "lora_A" in key: + linear_dim = int(value.shape[0]) + break + linear_alpha = linear_dim + network_config["linear"] = linear_dim + network_config["linear_alpha"] = linear_alpha + + # we build the keys to match every key + only_if_contains = [] + for key in lora_state_dict.keys(): + contains_key = key.split(".lora_")[0] + if contains_key not in only_if_contains: + only_if_contains.append(contains_key) + + network_kwargs["only_if_contains"] = only_if_contains + elif network_config["type"] == "lokr": + # find the factor + largest_factor = 0 + for key, value in lora_state_dict.items(): + if "lokr_w1" in key: + factor = int(value.shape[0]) + if factor > largest_factor: + largest_factor = factor + network_config["lokr_full_rank"] = True + network_config["lokr_factor"] = largest_factor + + only_if_contains = [] + for key in lora_state_dict.keys(): + if "lokr_w1" in key: + contains_key = key.split(".lokr_w1")[0] + contains_key = contains_key.replace("lycoris_", "") + if contains_key not in only_if_contains: + only_if_contains.append(contains_key) + network_kwargs["only_if_contains"] = only_if_contains + + if hasattr(base_model, 'target_lora_modules'): + network_kwargs['target_lin_modules'] = base_model.target_lora_modules + + # todo auto grab these + # get dim and scale + network_config = NetworkConfig(**network_config) + + network = LoRASpecialNetwork( + text_encoder=None, + unet=model_to_quantize, + lora_dim=network_config.linear, + multiplier=1.0, + alpha=network_config.linear_alpha, + # conv_lora_dim=self.network_config.conv, + # conv_alpha=self.network_config.conv_alpha, + train_unet=True, + train_text_encoder=False, + network_config=network_config, + network_type=network_config.type, + transformer_only=network_config.transformer_only, + is_transformer=base_model.is_transformer, + base_model=base_model, + is_ara=True, + **network_kwargs + ) + network.apply_to( + None, model_to_quantize, apply_text_encoder=False, apply_unet=True + ) + network.force_to(base_model.device_torch, dtype=base_model.torch_dtype) + network._update_torch_multiplier() + network.load_weights(lora_state_dict) + network.eval() + network.is_active = True + network.can_merge_in = False + base_model.accuracy_recovery_adapter = network + + # quantize it + lora_exclude_modules = [] + quantization_type = get_qtype(base_model.model_config.qtype) + for lora_module in tqdm(network.unet_loras, desc="Attaching quantization"): + # the lora has already hijacked the original module + orig_module = lora_module.org_module[0] + orig_module.to(base_model.torch_dtype) + # make the params not require gradients + for param in orig_module.parameters(): + param.requires_grad = False + quantize(orig_module, weights=quantization_type) + freeze(orig_module) + module_name = lora_module.lora_name.replace('$$', '.').replace('transformer.', '') + lora_exclude_modules.append(module_name) + if base_model.model_config.low_vram: + # move it back to cpu + orig_module.to("cpu") + pass + # quantize additional layers + print_acc(" - quantizing additional layers") + quantization_type = get_qtype('uint8') + quantize( + model_to_quantize, + weights=quantization_type, + exclude=lora_exclude_modules + exclude_modules + ) + + +def _quantize_model_blocks(base_model, model_to_quantize, exclude_modules): + if True: # quantize model the original way without an accuracy recovery adapter # move and quantize only certain pieces at a time. quantization_type = get_qtype(base_model.model_config.qtype) From 92df2899312052599dca33a85763ddaa11186962 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 28 Aug 2026 10:48:34 -0600 Subject: [PATCH 10/14] Fix quantization issues --- .../hidream/hidream_o1_model.py | 22 +- testing/stage_profiler.py | 157 +++++++++++ testing/test_model_loading.py | 266 +++++++++++++++++- toolkit/models/v2/_mixin.py | 22 ++ toolkit/models/v2/diffusion_models/wan.py | 6 + toolkit/models/v2/text_encoders/umt5.py | 3 + toolkit/util/convrot_quant.py | 13 + 7 files changed, 470 insertions(+), 19 deletions(-) create mode 100644 testing/stage_profiler.py diff --git a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py index 41bcf0c..1011808 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py @@ -1,5 +1,5 @@ import os -from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder +from toolkit.models.v2._mixin import OstrisTransformersMixin from typing import List, Optional import torch @@ -28,6 +28,17 @@ from .src.hidream_o1.model_config import model_config if TYPE_CHECKING: from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO + +class HidreamO1Transformer(Qwen3VLForConditionalGeneration, OstrisTransformersMixin): + """The o1 DiT-in-LLM: the vendored Qwen3VL with the image-diffusion heads + (x_embedder / t_embedder1 / final_layer2). The generic Qwen3VLTextEncoder + must NOT be used here — it drops those keys as unexpected.""" + + @classmethod + def get_transformer_block_names(cls): + return ["model.language_model.layers"] + + scheduler_config = { "num_train_timesteps": 1000, "shift": 3.0, @@ -129,7 +140,10 @@ class HidreamO1Model(BaseModel): self.use_old_lokr_format = False self.is_flow_matching = True self.is_transformer = True - self.target_lora_modules = ["Qwen3VLForConditionalGeneration"] + self.target_lora_modules = [ + "Qwen3VLForConditionalGeneration", + "HidreamO1Transformer", + ] self.noise_scale = self.model_config.model_kwargs.get( "noise_scale", DEFAULT_NOISE_SCALE ) @@ -188,7 +202,7 @@ class HidreamO1Model(BaseModel): ) # transformer.load_state_dict(state_dict, assign=True) - transformer = Qwen3VLTextEncoder.from_pretrained( + transformer = HidreamO1Transformer.from_pretrained( None, config=Qwen3VLConfig(**model_config), state_dict=state_dict, @@ -196,7 +210,7 @@ class HidreamO1Model(BaseModel): ) del state_dict # free memory else: - transformer = Qwen3VLTextEncoder.from_pretrained( + transformer = HidreamO1Transformer.from_pretrained( model_path, torch_dtype=self.torch_dtype, ) diff --git a/testing/stage_profiler.py b/testing/stage_profiler.py new file mode 100644 index 0000000..ab12020 --- /dev/null +++ b/testing/stage_profiler.py @@ -0,0 +1,157 @@ +"""Per-stage speed / memory / CPU profiler for the model loading test harness. + +A StageProfiler splits a run into named stages (the harness hooks +BaseModel.print_and_status_update so every holder status line — "Loading +transformer", "Quantizing (convrot8)", "Loading VAE", ... — opens a new +stage). For every stage it records: + + - seconds wall time + - vram_peak_gb torch.cuda.max_memory_allocated during the stage + - vram_reserved_gb torch.cuda.max_memory_reserved during the stage + - vram_end_gb memory_allocated at the stage boundary (resident model) + - rss_peak_gb peak process RSS (sampled at 50ms) + - cpu_avg_pct / process CPU utilization (100 = one full core), sampled + cpu_max_pct at 50ms — this is what shows where quantization burns CPU + - threads_max peak thread count + +A background sampler thread does the RSS/CPU sampling so short spikes inside +a stage are caught, not just the boundary values. +""" + +import threading +import time + +import psutil +import torch + +_GB = 1024**3 + + +class _Sampler: + """Background thread sampling RSS / CPU% / thread count at 50ms.""" + + def __init__(self): + self.proc = psutil.Process() + self._lock = threading.Lock() + self._stop = False + self.reset() + self.proc.cpu_percent(None) # prime the cpu_percent window + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + + def reset(self): + with self._lock: + self.rss_peak = self.proc.memory_info().rss + self.cpu_samples = [] + self.threads_max = self.proc.num_threads() + + def _run(self): + while not self._stop: + try: + rss = self.proc.memory_info().rss + cpu = self.proc.cpu_percent(None) + threads = self.proc.num_threads() + with self._lock: + self.rss_peak = max(self.rss_peak, rss) + self.cpu_samples.append(cpu) + self.threads_max = max(self.threads_max, threads) + except Exception: + pass + time.sleep(0.05) + + def stats(self): + with self._lock: + samples = self.cpu_samples or [0.0] + return { + "rss_peak_gb": round(self.rss_peak / _GB, 3), + "cpu_avg_pct": round(sum(samples) / len(samples), 1), + "cpu_max_pct": round(max(samples), 1), + "threads_max": self.threads_max, + } + + def stop(self): + self._stop = True + + +class StageProfiler: + def __init__(self, device): + self.device = torch.device(device) if torch.cuda.is_available() else None + self.stages = [] + self._current = None + self._sampler = _Sampler() + + def _cuda_sync_reset(self): + if self.device is not None: + torch.cuda.synchronize(self.device) + torch.cuda.reset_peak_memory_stats(self.device) + + def stage(self, name: str): + """Close the current stage (if any) and open a new one.""" + self._close() + self._cuda_sync_reset() + self._sampler.reset() + self._current = {"name": str(name), "t0": time.perf_counter()} + + def _close(self): + if self._current is None: + return + cur = self._current + self._current = None + seconds = time.perf_counter() - cur["t0"] + entry = {"name": cur["name"], "seconds": round(seconds, 3)} + if self.device is not None: + torch.cuda.synchronize(self.device) + entry["vram_peak_gb"] = round( + torch.cuda.max_memory_allocated(self.device) / _GB, 3 + ) + entry["vram_reserved_gb"] = round( + torch.cuda.max_memory_reserved(self.device) / _GB, 3 + ) + entry["vram_end_gb"] = round( + torch.cuda.memory_allocated(self.device) / _GB, 3 + ) + entry.update(self._sampler.stats()) + self.stages.append(entry) + + def finish(self): + self._close() + self._sampler.stop() + return self.stages + + +def profile_top(profile, limit=40): + """Top functions from a cProfile.Profile by cumulative time, as dicts.""" + import pstats + + stats = pstats.Stats(profile) + stats.sort_stats("cumulative") + rows = [] + for func in stats.fcn_list[: limit * 3]: + cc, nc, tt, ct, _ = stats.stats[func] + filename, line, name = func + # drop the profiler/exec wrappers and trivial rows + if name in ("",) or ct < 0.05: + continue + # cProfile catches this module's own 50ms sampling thread; its idle + # sleep loop and psutil polling otherwise show up as a giant fake + # "time.sleep" hotspot spanning the whole stage + if "stage_profiler" in filename or "psutil" in filename or "_pslinux" in filename: + continue + if name == "": + continue + # shorten site-packages / repo paths to keep the report readable + for marker in ("site-packages/", "ai-toolkit/"): + if marker in filename: + filename = filename.split(marker, 1)[1] + break + rows.append( + { + "func": f"{filename}:{line}({name})", + "ncalls": nc, + "tottime": round(tt, 3), + "cumtime": round(ct, 3), + } + ) + if len(rows) >= limit: + break + return rows diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index 56da4f9..a1f45b1 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -23,6 +23,7 @@ import json import os import subprocess import sys +import time TOOLKIT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, TOOLKIT_ROOT) @@ -57,11 +58,11 @@ MODEL_TESTS = { "boogu_image": { # native ~1024; 512/low-step/high-CFG degenerates to a black frame "model": {"name_or_path": "Boogu/Boogu-Image-0.1-Base", "quantize": True, "quantize_te": True}, - "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, }, "ernie_image": { "model": {"name_or_path": "baidu/ERNIE-Image", "quantize": True, "quantize_te": True}, - "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, }, "mageflow": { "model": {"name_or_path": "microsoft/Mage-Flow-Base", "quantize": True, "quantize_te": True}, @@ -69,15 +70,15 @@ MODEL_TESTS = { }, "ideogram4": { "model": {"name_or_path": "ideogram-ai/ideogram-4-fp8", "quantize": True, "quantize_te": True}, - "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, }, "hidream_o1": { "model": {"name_or_path": "HiDream-ai/HiDream-O1-Image", "quantize": True, "quantize_te": True}, - "sample": {"width": 1024, "height": 1024, "num_inference_steps": 28, "guidance_scale": 5.0, "seed": 42}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 28, "guidance_scale": 5.0, "seed": 42}, }, "anima": { "model": {"name_or_path": "circlestone-labs/Anima-Base-v1.0-Diffusers"}, - "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.5, "seed": 42}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.5, "seed": 42}, }, "wan21": { "model": {"name_or_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"}, @@ -124,7 +125,7 @@ MODEL_TESTS = { }, "hidream": { "model": {"name_or_path": "HiDream-ai/HiDream-I1-Full", "quantize": True, "quantize_te": True}, - "sample": {"width": 1024, "height": 1024, "num_inference_steps": 28, "guidance_scale": 5.0, "seed": 42}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 28, "guidance_scale": 5.0, "seed": 42}, }, "hidream_e1": { # editing seq budget fits 768x768 (source+target concat) @@ -134,11 +135,11 @@ MODEL_TESTS = { }, "nucleus_image": { "model": {"name_or_path": "NucleusAI/Nucleus-Image", "quantize": True, "quantize_te": True}, - "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, }, "omnigen2": { "model": {"name_or_path": "OmniGen2/OmniGen2", "quantize": True, "quantize_te": True}, - "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, }, "ltx2.5": { "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, @@ -181,7 +182,7 @@ MODEL_TESTS = { }, "sdxl": { "model": {"name_or_path": "stabilityai/stable-diffusion-xl-base-1.0"}, - "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 6.0, "seed": 42}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 6.0, "seed": 42}, }, "ace_step_15": { "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True}, @@ -189,7 +190,7 @@ MODEL_TESTS = { }, "f-lite": { "model": {"name_or_path": "Freepik/F-Lite", "quantize": True, "quantize_te": True}, - "sample": {"width": 1024, "height": 1024, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, }, } @@ -228,7 +229,14 @@ def run_one(arch: str, device: str, allow_download: bool) -> dict: from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.util.get_model import get_model_class - model_config = ModelConfig(arch=arch, dtype="bf16", **entry["model"]) + # quantization always tests the convrot8 backend (per-entry override wins) + model_kwargs = dict(entry["model"]) + if model_kwargs.get("quantize"): + model_kwargs.setdefault("qtype", "convrot8") + if model_kwargs.get("quantize_te"): + model_kwargs.setdefault("qtype_te", "convrot8") + + model_config = ModelConfig(arch=arch, dtype="bf16", **model_kwargs) ModelClass = get_model_class(model_config) from toolkit.util.get_model import LEGACY_ARCHS @@ -271,7 +279,32 @@ def run_one(arch: str, device: str, allow_download: bool) -> dict: dtype="bf16", noise_scheduler=sampler, ) + + # ---- per-stage speed / VRAM / CPU instrumentation + cProfile ---- + import cProfile + + from testing.stage_profiler import StageProfiler, profile_top + + prof = StageProfiler(device) + if hasattr(sd, "print_and_status_update"): + # every holder announces its stages through print_and_status_update; + # each status line opens a new profiler stage + _orig_status = sd.print_and_status_update + + def _status_hook(msg, *a, **k): + prof.stage(str(msg)) + return _orig_status(msg, *a, **k) + + sd.print_and_status_update = _status_hook + + prof.stage("load: init") + load_profile = cProfile.Profile() + load_profile.enable() + t_load0 = time.perf_counter() sd.load_model() + load_seconds = time.perf_counter() - t_load0 + load_profile.disable() + load_profile.dump_stats(os.path.join(out_dir, "load.prof")) sample_kwargs = dict(entry["sample"]) if entry.get("needs_control_image"): @@ -295,16 +328,175 @@ def run_one(arch: str, device: str, allow_download: bool) -> dict: if not hasattr(ModelClass, "get_train_scheduler"): # the legacy monolith takes the sampler NAME at generate time gen_kwargs["sampler"] = "ddpm" + + prof.stage("generate") + gen_profile = cProfile.Profile() + gen_profile.enable() + t_gen0 = time.perf_counter() sd.generate_images([gen], **gen_kwargs) + gen_seconds = time.perf_counter() - t_gen0 + gen_profile.disable() + gen_profile.dump_stats(os.path.join(out_dir, "gen.prof")) + + stages = prof.finish() produced = [ p for p in glob.glob(os.path.join(out_dir, "*")) - if os.path.isfile(p) and os.path.getsize(p) > 1024 and not p.endswith(".txt") + if os.path.isfile(p) + and os.path.getsize(p) > 1024 + and not p.endswith((".txt", ".prof", ".json")) ] if not produced: raise RuntimeError(f"no output file produced in {out_dir}") - return {"arch": arch, "status": "PASS", "outputs": produced} + + import torch + + result = { + "arch": arch, + "status": "PASS", + "outputs": produced, + "qtype": model_config.qtype if model_kwargs.get("quantize") else None, + "qtype_te": model_config.qtype_te if model_kwargs.get("quantize_te") else None, + "load_seconds": round(load_seconds, 3), + "gen_seconds": round(gen_seconds, 3), + "stages": stages, + "profile": { + "load_top": profile_top(load_profile), + "gen_top": profile_top(gen_profile, limit=15), + "load_prof": os.path.join(out_dir, "load.prof"), + "gen_prof": os.path.join(out_dir, "gen.prof"), + }, + "env": { + "torch_num_threads": torch.get_num_threads(), + "cpu_count": os.cpu_count(), + "omp_num_threads": os.environ.get("OMP_NUM_THREADS"), + }, + } + with open(os.path.join(out_dir, "metrics.json"), "w") as f: + json.dump(result, f, indent=2) + return result + + +def format_stage_table(stages) -> str: + """Aligned markdown table: stage column left-aligned, numbers right-aligned + with fixed decimals so they line up in the terminal too.""" + headers = [ + "stage", "seconds", "vram peak GB", "vram resv GB", "vram end GB", + "rss peak GB", "cpu avg %", "cpu max %", "threads", + ] + + def fmt(value, decimals): + if value is None or value == "-": + return "-" + return f"{value:.{decimals}f}" + + rows = [ + [ + s["name"], + fmt(s.get("seconds"), 2), + fmt(s.get("vram_peak_gb", "-"), 2), + fmt(s.get("vram_reserved_gb", "-"), 2), + fmt(s.get("vram_end_gb", "-"), 2), + fmt(s.get("rss_peak_gb"), 2), + fmt(s.get("cpu_avg_pct"), 1), + fmt(s.get("cpu_max_pct"), 1), + str(s.get("threads_max", "-")), + ] + for s in stages + ] + + widths = [ + max(len(headers[i]), *(len(r[i]) for r in rows)) if rows else len(headers[i]) + for i in range(len(headers)) + ] + + def line(cells): + out = [] + for i, cell in enumerate(cells): + # stage name left-aligned, everything else right-aligned + out.append(cell.ljust(widths[i]) if i == 0 else cell.rjust(widths[i])) + return "| " + " | ".join(out) + " |" + + sep = "|" + "|".join( + (":" + "-" * (w + 1)) if i == 0 else ("-" * (w + 1) + ":") + for i, w in enumerate(widths) + ) + "|" + return "\n".join([line(headers), sep] + [line(r) for r in rows]) + + +def write_report(results, path): + """Aggregated markdown report: per-arch stage tables + cross-arch summary + + load-profile hotspots, for diagnosing load/quantization speedups and + CPU usage.""" + + def _sum(stages, pred): + return round(sum(s["seconds"] for s in stages if pred(s)), 1) + + lines = ["# Model loading test report", ""] + lines.append( + "Per-arch stage timings, peak VRAM (torch allocated/reserved), peak " + "process RSS, and CPU utilization (100 = one core, sampled at 50ms). " + "Quantization is convrot8 across the board. Full cProfile dumps sit " + "next to each arch's outputs (load.prof / gen.prof; inspect with " + "`python -m pstats ` or snakeviz)." + ) + lines.append("") + + # ---- summary ---- + lines += ["## Summary", ""] + lines += [ + "| arch | status | load s | quantize s | generate s | vram peak GB | rss peak GB | quant cpu avg % |", + "|---|---|---|---|---|---|---|---|", + ] + for r in results: + if r["status"] != "PASS" or "stages" not in r: + lines.append( + f"| {r['arch']} | {r['status']} | - | - | - | - | - | - | " + ) + continue + stages = r["stages"] + quant = [s for s in stages if "uantiz" in s["name"]] + vram_peak = max((s.get("vram_peak_gb", 0) for s in stages), default=0) + rss_peak = max((s.get("rss_peak_gb", 0) for s in stages), default=0) + quant_cpu = ( + round(sum(s["cpu_avg_pct"] for s in quant) / len(quant), 1) + if quant + else "-" + ) + lines.append( + f"| {r['arch']} | PASS | {r['load_seconds']} " + f"| {_sum(quant, lambda s: True)} | {r['gen_seconds']} " + f"| {vram_peak} | {rss_peak} | {quant_cpu} |" + ) + lines.append("") + + # ---- per-arch detail ---- + for r in results: + if r["status"] != "PASS" or "stages" not in r: + lines += [f"## {r['arch']} — {r['status']}", "", r.get("error", ""), ""] + continue + lines += [f"## {r['arch']}", ""] + lines += [ + f"load {r['load_seconds']}s, generate {r['gen_seconds']}s, " + f"qtype {r.get('qtype')}, qtype_te {r.get('qtype_te')}, " + f"torch threads {r['env']['torch_num_threads']}/{r['env']['cpu_count']} cores", + "", + ] + lines += [format_stage_table(r["stages"]), ""] + top = r.get("profile", {}).get("load_top", [])[:12] + if top: + lines += ["Load hotspots (cumulative):", "", "```"] + for row in top: + lines.append( + f"{row['cumtime']:>9.2f}s tot {row['tottime']:>8.2f}s " + f"x{row['ncalls']:<8} {row['func']}" + ) + lines += ["```", ""] + + with open(path, "w") as f: + f.write("\n".join(lines)) + return path def main(): @@ -315,6 +507,11 @@ def main(): parser.add_argument("--device", type=str, default="cuda:0") parser.add_argument("--allow-download", action="store_true") parser.add_argument("--json-result", type=str, default=None) + parser.add_argument( + "--report-only", + action="store_true", + help="rebuild report.md from each arch's saved metrics.json, no runs", + ) args = parser.parse_args() if args.list: @@ -322,6 +519,24 @@ def main(): print(arch) return + if args.report_only: + results = [] + for arch in MODEL_TESTS: + out_dir = os.path.join( + OUTPUT_ROOT, arch.replace("/", "_").replace(":", "_") + ) + metrics_path = os.path.join(out_dir, "metrics.json") + if os.path.exists(metrics_path): + with open(metrics_path) as f: + results.append(json.load(f)) + else: + results.append( + {"arch": arch, "status": "SKIP", "error": "no metrics.json saved"} + ) + report_path = write_report(results, os.path.join(OUTPUT_ROOT, "report.md")) + print(f"report: {report_path}") + return + os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") if not args.allow_download: os.environ.setdefault("HF_HUB_OFFLINE", "1") @@ -343,6 +558,10 @@ def main(): if args.json_result: with open(args.json_result, "w") as f: json.dump(result, f) + if result["status"] == "PASS" and "stages" in result: + print() + print(format_stage_table(result["stages"])) + print() print(f"[{result['status']}] {args.arch}" + (f" — {result.get('error', '')}" if result["status"] != "PASS" else "")) if result["status"] == "FAIL": sys.exit(1) @@ -355,8 +574,10 @@ def main(): # --all: one subprocess per arch so every model fully unloads (clean CUDA # teardown) before the next loads results = [] - for arch in MODEL_TESTS: - print(f"\n===== {arch} =====") + total = len(MODEL_TESTS) + sweep_t0 = time.perf_counter() + for i, arch in enumerate(MODEL_TESTS, start=1): + print(f"\n===== {arch} ({i}/{total}) =====") result_path = os.path.join(OUTPUT_ROOT, f".{arch.replace('/', '_')}.result.json") cmd = [ sys.executable, @@ -380,6 +601,18 @@ def main(): {"arch": arch, "status": "FAIL", "error": f"subprocess died (exit {proc.returncode})"} ) + # running tally so a long sweep always shows how far along it is + tally = {"PASS": 0, "FAIL": 0, "SKIP": 0} + for r in results: + tally[r["status"]] = tally.get(r["status"], 0) + 1 + elapsed = time.perf_counter() - sweep_t0 + eta = (elapsed / i) * (total - i) + print( + f">>> progress: {i}/{total} ({i * 100 // total}%) — " + f"{tally['PASS']} pass, {tally['FAIL']} fail, {tally['SKIP']} skip — " + f"elapsed {elapsed / 60:.0f}m, eta ~{eta / 60:.0f}m" + ) + print("\n===== summary =====") counts = {"PASS": 0, "FAIL": 0, "SKIP": 0} for r in results: @@ -389,6 +622,9 @@ def main(): line += f" — {r.get('error', '')[:160]}" print(line) print(f"\n{counts['PASS']} passed, {counts['FAIL']} failed, {counts['SKIP']} skipped") + + report_path = write_report(results, os.path.join(OUTPUT_ROOT, "report.md")) + print(f"report: {report_path}") if counts["FAIL"]: sys.exit(1) diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index def4cbb..5a2e71e 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -91,6 +91,13 @@ class OstrisModelMixin: # 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 @@ -461,7 +468,22 @@ class OstrisModelMixin: 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(): diff --git a/toolkit/models/v2/diffusion_models/wan.py b/toolkit/models/v2/diffusion_models/wan.py index d71cdaf..fd0367a 100644 --- a/toolkit/models/v2/diffusion_models/wan.py +++ b/toolkit/models/v2/diffusion_models/wan.py @@ -76,6 +76,12 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin): }, } + # the comfy fp8/scaled_fp8 wan files mix fp16 conv biases with fp32 + # tables; wan's forward assumes one uniform dtype, so cast the + # non-quantized tensors to the load dtype (the old holder did this with a + # blanket .to(device, dtype) after load) + aitk_cast_quantized_load = True + @classmethod def get_transformer_block_names(cls): return ["blocks"] diff --git a/toolkit/models/v2/text_encoders/umt5.py b/toolkit/models/v2/text_encoders/umt5.py index 8bd02e0..2b86dff 100644 --- a/toolkit/models/v2/text_encoders/umt5.py +++ b/toolkit/models/v2/text_encoders/umt5.py @@ -34,6 +34,9 @@ class UMT5TextEncoder(UMT5EncoderModel, OstrisTransformersMixin): aitk_subfolder = "text_encoder" aitk_tokenizer_subfolder = "tokenizer" aitk_config_repo = "ai-toolkit/umt5_xxl_encoder" + # comfy fp8 umt5 files store non-quantized tensors in a dtype mix; T5 + # assumes a uniform dtype, so follow the load dtype + aitk_cast_quantized_load = True aitk_comfy_repo = "Comfy-Org/Wan_2.1_ComfyUI_repackaged" # comfy umt5 files already use the transformers key layout; the diff --git a/toolkit/util/convrot_quant.py b/toolkit/util/convrot_quant.py index 0411281..ec0645a 100644 --- a/toolkit/util/convrot_quant.py +++ b/toolkit/util/convrot_quant.py @@ -1319,6 +1319,19 @@ class ConvRotInt8Quantizer(OstrisQuantizer): f"(needs in divisible by 16, out by 8, and a power-of-4 block >= 16)" ) return False + if d < 128 and module.out_features % 16 != 0: + # cublasLt has no int8 kernel for K < 128 with N % 16 != 0 + # (torch._int_mm raises CUBLAS_STATUS_NOT_SUPPORTED), e.g. + # omnigen2's 64 -> 2520 x_embedder + key = (d, module.out_features) + if key not in _skip_warned: + _skip_warned.add(key) + print_acc( + f"ConvRot: skipping linears with in_features={d}, " + f"out_features={module.out_features} (int8 gemm needs out " + f"divisible by 16 when in < 128)" + ) + return False return True def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None: From ce3df321013d116c7ee31d8ebe2a0988b04fc368 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 28 Aug 2026 10:54:26 -0600 Subject: [PATCH 11/14] Fix speed on wan conv patch --- toolkit/models/v2/diffusion_models/wan.py | 42 +++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/toolkit/models/v2/diffusion_models/wan.py b/toolkit/models/v2/diffusion_models/wan.py index fd0367a..b6011c1 100644 --- a/toolkit/models/v2/diffusion_models/wan.py +++ b/toolkit/models/v2/diffusion_models/wan.py @@ -1,8 +1,45 @@ +import torch +import torch.nn.functional as F from diffusers import WanTransformer3DModel as DiffusersWanTransformer3DModel from .._mixin import OstrisModelMixin +def patch_wan_patch_embedding(transformer) -> bool: + """Wan's patch_embedding is a Conv3d whose kernel == stride (non-overlapping + patchify), i.e. a plain linear projection of each flattened patch. bf16 + Conv3d has no fast cuDNN kernel and falls back to a slow, GPU-underutilizing + path — swap the forward for the exactly-equivalent F.linear (a GEMM), + returning the conv's own [b, out, t', h', w'] layout so it stays a drop-in. + The weight is read lazily so this survives later .to(device)/dtype moves.""" + conv = getattr(transformer, "patch_embedding", None) + if not isinstance(conv, torch.nn.Conv3d): + return False + if ( + tuple(conv.kernel_size) != tuple(conv.stride) + or any(p != 0 for p in conv.padding) + or any(d != 1 for d in conv.dilation) + or conv.groups != 1 + ): + return False + + def fast_forward(x, _conv=conv): + weight = _conv.weight + out_c = weight.shape[0] + kt, kh, kw = _conv.kernel_size + b, c, t, h, w = x.shape + tp, hp, wp = t // kt, h // kh, w // kw + # patchify with (c, kt, kh, kw) fastest — matches weight.reshape below + x = x.reshape(b, c, tp, kt, hp, kh, wp, kw) + x = x.permute(0, 2, 4, 6, 1, 3, 5, 7).reshape(b, tp * hp * wp, -1) + w2d = weight.reshape(out_c, -1) + x = F.linear(x.to(w2d.dtype), w2d, _conv.bias) + return x.transpose(1, 2).reshape(b, out_c, tp, hp, wp) + + conv.forward = fast_forward + return True + + class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin): """Wan 2.1/2.2 video DiT (wan22 loads two of these into its dual wrapper).""" @@ -86,6 +123,11 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin): def get_transformer_block_names(cls): return ["blocks"] + def aitk_post_load(self, **kwargs): + # slow bf16 Conv3d patchify -> equivalent GEMM, before quantize/offload + patch_wan_patch_embedding(self) + return super().aitk_post_load(**kwargs) + def get_offload_ignore_modules(self): # fp32 modulation tables must stay resident on the compute device return [self.scale_shift_table] + [ From 5ddc5f8ca717dbfaaa2a564156daca736a4535c6 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 28 Aug 2026 11:47:29 -0600 Subject: [PATCH 12/14] Optimize quantization path --- toolkit/models/v2/PLANNING.md | 34 +++++++++++++++++ toolkit/models/v2/_mixin.py | 31 ++++++++++++++- toolkit/util/quantize.py | 71 +++++++++++++++++++++++++++++++---- 3 files changed, 127 insertions(+), 9 deletions(-) diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index ca0d6da..773e26b 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -442,6 +442,40 @@ them: Flux2, MageFlow, ZImageDCT, Ideogram4Transformer2DModel (+ its exclude list on MageFlow). `prepare_text_encoder` remains only as the legacy-monolith path. +## Load/quantize optimization round 1 (2026-08-28) + +Baseline: testing/.model_test_outputs/report_baseline.md (+ per-arch +metrics_baseline.json). Profiling findings and fixes: + +- The "slow quantization" in the baseline was mostly NOT quantize math + (convrot8 on chroma's 57 blocks: 0.5s of GPU kernels). It was data + movement: mmap-backed safetensors faulting pages in per-tensor cudaMemcpy + order (cold cache on the 92%-full NVMe reads at ~0.32GB/s; the drive's + sequential ceiling is ~435MB/s — cold loads are storage-bound), plus the + old stream-quantize choreography crossing the bus three times (up, back + into a fresh host copy, up again at final placement). +- `quantize_module(keep_on_device=True)`: when the final home IS the + quantize gpu (no offload/low_vram — aitk_post_load detects this), blocks + stay on the gpu after quantizing and the extras pass runs there too. + Weights cross the bus once; the model-sized host-RAM spike of the + intermediate copy is gone (chroma RSS peak 25.3 -> 17.3GB, the rest is + reclaimable page cache; flux2's 93GB spike class eliminated). Warm-cache + chroma blocks: 46s -> 3.9s; wan22_5b quantize 10.6s -> 3.8s. +- Offload/low_vram path keeps cpu residency but quantizes extras layer-by- + layer via quantize_device instead of the all-cores cpu burst (was 350-650% + cpu avg with 1600% spikes). +- `quantize()` skips the quantize_device round-trip for same-qtype + OstrisLinears (guaranteed no-op) and, for ostris backends, for non-Linear + leaves (norms/embeddings were being ferried across the bus for nothing). +- posix_fadvise(WILLNEED) readahead on single-file loads (mixin._readahead) + — ~10-15% on cold reads, free when warm, no-op off POSIX. + +Ideas not yet done: pinned-buffer staged h2d (~+60% warm-upload bandwidth, +9.5 vs 5.8GB/s measured), overlapping h2d with quantize kernels (small — the +kernels are ~2% of the pass), wan22 low_vram generate choreography (165s of +unload/reload per sample), extending keep_on_device to the legacy +quantize_model/trainer path. + ## Testing - [x] `testing/test_model_loading.py`: per-arch load + one small sample through diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index 5a2e71e..35641e2 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -257,10 +257,21 @@ class OstrisModelMixin: attach_ara_and_quantize(base_model, self, ara_path, exclude=exclude) else: status_fn(f"Quantizing ({qtype})") + 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" + ) quantize_module( self, qtype, - device=quantize_device if quantize_device is not None else device, + device=q_device, dtype=dtype, block_names=self.get_transformer_block_names(), exclude=exclude, @@ -270,6 +281,7 @@ class OstrisModelMixin: else None ), status_fn=status_fn, + keep_on_device=keep_on_device, ) self.aitk_is_quantized = True self.aitk_qtype = qtype @@ -407,6 +419,22 @@ class OstrisModelMixin: 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, @@ -417,6 +445,7 @@ class OstrisModelMixin: subfolder: Optional[str] = None, **kwargs, ): + cls._readahead(file_path) state_dict = load_file(file_path) return cls.load_from_state_dict( state_dict, diff --git a/toolkit/util/quantize.py b/toolkit/util/quantize.py index 1f820a1..0b55cbe 100644 --- a/toolkit/util/quantize.py +++ b/toolkit/util/quantize.py @@ -1,3 +1,5 @@ +import os +import time from fnmatch import fnmatch from typing import List, Optional, Union, TYPE_CHECKING import torch @@ -186,6 +188,16 @@ def quantize( and m.__class__.__name__ == "OstrisLinear" ): continue + if getattr(m.ostris_quantizer, "qtype", None) == weights.quantizer.qtype: + # guaranteed per-layer no-op: skip before any + # quantize_device round-trip + continue + if isinstance(weights, ostristype) and not isinstance(m, torch.nn.Linear): + # ostris backends only quantize nn.Linear; don't ferry norms/ + # embeddings across the bus for nothing when quantize_device + # is set (containers fall through so children are visited) + if quantize_device is not None and next(m.children(), None) is None: + continue if ( isinstance(weights, aotype) and not isinstance(m, torch.nn.Linear) @@ -271,17 +283,29 @@ def quantize_module( exclude: Optional[List[str]] = None, quantize_kwargs: Optional[dict] = None, status_fn=print_acc, + keep_on_device: bool = False, ): """Module-centric quantization: block-streamed (each repeated block moves - to ``device`` for the math and returns to cpu) with a whole-module pass - for the extras. This is the core the per-model loaders call; holders' - quantize_model wraps it with model_config plumbing.""" + to ``device`` for the math) with a whole-module pass for the extras. This + is the core the per-model loaders call; holders' quantize_model wraps it + with model_config plumbing. + + ``keep_on_device``: when the model's final home IS ``device`` (no layer + offloading / low_vram), quantized blocks stay there instead of round- + tripping back to cpu — the weights then cross the bus exactly once + (mmap/page-cache -> gpu) instead of three times (up, back into a fresh + host copy, up again at final placement), and the model-sized host RAM + spike of that intermediate copy never happens. The extras pass runs on + the gpu too. With it off (offload paths), blocks return to cpu as before + and the extras quantize layer-by-layer on ``device`` via quantize_device + rather than burning every cpu core.""" from toolkit.dequantize import patch_dequantization_on_save patch_dequantization_on_save(module) quantization_type = get_qtype(qtype) exclude = list(exclude or []) quantize_kwargs = quantize_kwargs or {} + keep_on_device = keep_on_device and device is not None all_blocks: List[torch.nn.Module] = [] for name in block_names or []: @@ -296,29 +320,60 @@ def quantize_module( if all_blocks: status_fn(f" - quantizing {len(all_blocks)} blocks") already_quantized = 0 + debug_phases = os.environ.get("AITK_QUANT_DEBUG") == "1" + t_check = t_h2d = t_quant = t_d2h = 0.0 for block in tqdm(all_blocks): - if not _has_quantizable_linear(block, quantization_type, exclude): + t = time.perf_counter() + skip = not _has_quantizable_linear(block, quantization_type, exclude) + t_check += time.perf_counter() - t + if skip: # pre-quantized checkpoint with a matching qtype: nothing in this - # block would change — skip the device round-trip and the dtype - # cast entirely so the load stays byte-identical + # block would change — skip the dtype cast entirely so the load + # stays byte-identical (placement still honors keep_on_device) already_quantized += 1 + if keep_on_device: + block.to(device) continue + t = time.perf_counter() if device is not None: block.to(device, dtype=dtype, non_blocking=True) + t_h2d += time.perf_counter() - t + t = time.perf_counter() quantize(block, weights=quantization_type, exclude=exclude, **quantize_kwargs) freeze(block) + t_quant += time.perf_counter() - t # NOT non_blocking: an async D2H allocates the cpu destination in pinned # memory, which the caching host allocator keeps forever — that silently # retained a model-sized chunk of host ram - if device is not None: + t = time.perf_counter() + if device is not None and not keep_on_device: block.to("cpu") + t_d2h += time.perf_counter() - t + if debug_phases and all_blocks: + status_fn( + f" - [debug] check {t_check:.1f}s h2d {t_h2d:.1f}s " + f"quant {t_quant:.1f}s d2h {t_d2h:.1f}s" + ) if already_quantized: status_fn( f" - {already_quantized} blocks already quantized with a matching qtype; left untouched" ) status_fn(" - quantizing extras") - quantize(module, weights=quantization_type, exclude=exclude, **quantize_kwargs) + if keep_on_device: + # blocks already live on the gpu; bring the remainder over and run the + # extras pass there (fast kernels instead of an all-cores cpu burst) + module.to(device) + quantize(module, weights=quantization_type, exclude=exclude, **quantize_kwargs) + else: + # cpu-resident model: quantize each extra layer with a gpu round-trip + quantize( + module, + weights=quantization_type, + exclude=exclude, + quantize_device=device, + **quantize_kwargs, + ) freeze(module) return module From 7195abc32b7d1360f9b0f613a89e48dacbb763f3 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 28 Aug 2026 22:21:09 -0600 Subject: [PATCH 13/14] Performance fixes and bug fixes --- .../diffusion_models/chroma/chroma_model.py | 8 +- .../chroma/chroma_radiance_model.py | 8 +- .../ernie_image/ernie_image.py | 6 +- .../diffusion_models/f_light/f_light.py | 6 +- .../diffusion_models/flux2/flux2_model.py | 13 +- .../diffusion_models/flux2/src/pipeline.py | 10 +- .../flux_kontext/flux_kontext.py | 8 +- .../diffusion_models/ltx2/ltx2.py | 6 +- .../nucleus_image/nucleus_image_model.py | 6 +- .../diffusion_models/qwen_image/qwen_image.py | 6 +- .../diffusion_models/wan22/wan22_14b_model.py | 20 ++- .../zeta_chroma/zeta_chroma_model.py | 6 +- testing/test_model_loading.py | 77 +++++++++- toolkit/dequantize.py | 4 + toolkit/memory_management/manager.py | 92 ++++++++++- toolkit/memory_management/manager_modules.py | 52 +++++++ toolkit/models/v2/PLANNING.md | 68 ++++++++- toolkit/models/v2/_mixin.py | 69 ++++++--- toolkit/models/wan21/wan21.py | 5 +- toolkit/util/convrot_quant.py | 35 ++++- toolkit/util/quantize.py | 144 +++++++++--------- 21 files changed, 511 insertions(+), 138 deletions(-) diff --git a/extensions_built_in/diffusion_models/chroma/chroma_model.py b/extensions_built_in/diffusion_models/chroma/chroma_model.py index 8d5a11a..91bc886 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_model.py @@ -206,11 +206,13 @@ class ChromaModel(BaseModel): pipe.transformer = pipe.transformer.to(self.device_torch) flush() - # just to make sure everything is on the right device and dtype - text_encoder[0].to(self.device_torch) + # low_vram: text encoders stay on cpu; get_prompt_embeds moves them + # to the gpu on demand + if not self.low_vram: + text_encoder[0].to(self.device_torch) + text_encoder[1].to(self.device_torch) text_encoder[0].requires_grad_(False) text_encoder[0].eval() - text_encoder[1].to(self.device_torch) text_encoder[1].requires_grad_(False) text_encoder[1].eval() pipe.transformer = pipe.transformer.to(self.device_torch) diff --git a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py index 6aacd39..fc84461 100644 --- a/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py +++ b/extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py @@ -192,11 +192,13 @@ class ChromaRadianceModel(BaseModel): pipe.transformer = pipe.transformer.to(self.device_torch) flush() - # just to make sure everything is on the right device and dtype - text_encoder[0].to(self.device_torch) + # low_vram: text encoders stay on cpu; get_prompt_embeds moves them + # to the gpu on demand + if not self.low_vram: + text_encoder[0].to(self.device_torch) + text_encoder[1].to(self.device_torch) text_encoder[0].requires_grad_(False) text_encoder[0].eval() - text_encoder[1].to(self.device_torch) text_encoder[1].requires_grad_(False) text_encoder[1].eval() pipe.transformer = pipe.transformer.to(self.device_torch) diff --git a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py index 393cb0a..2c175b2 100644 --- a/extensions_built_in/diffusion_models/ernie_image/ernie_image.py +++ b/extensions_built_in/diffusion_models/ernie_image/ernie_image.py @@ -133,8 +133,10 @@ class ErnieImageModel(BaseModel): pipe.transformer = pipe.transformer.to(self.device_torch) flush() - # just to make sure everything is on the right device and dtype - text_encoder[0].to(self.device_torch) + # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it + # to the gpu on demand + if not self.low_vram: + text_encoder[0].to(self.device_torch) text_encoder[0].requires_grad_(False) text_encoder[0].eval() flush() diff --git a/extensions_built_in/diffusion_models/f_light/f_light.py b/extensions_built_in/diffusion_models/f_light/f_light.py index 647cd88..f85ca11 100644 --- a/extensions_built_in/diffusion_models/f_light/f_light.py +++ b/extensions_built_in/diffusion_models/f_light/f_light.py @@ -115,8 +115,10 @@ class FLiteModel(BaseModel): pipe.transformer = pipe.transformer.to(self.device_torch) flush() - # just to make sure everything is on the right device and dtype - text_encoder[0].to(self.device_torch) + # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it + # to the gpu on demand + if not self.low_vram: + text_encoder[0].to(self.device_torch) text_encoder[0].requires_grad_(False) text_encoder[0].eval() pipe.transformer = pipe.transformer.to(self.device_torch) diff --git a/extensions_built_in/diffusion_models/flux2/flux2_model.py b/extensions_built_in/diffusion_models/flux2/flux2_model.py index cc023b4..1f770fb 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_model.py @@ -98,12 +98,21 @@ class Flux2Model(BaseModel): self.print_and_status_update("Loading Mistral") # load + quantize + offload + placement, all driven by model_config + # tie_word_embeddings=False: the checkpoint carries both embed_tokens + # and lm_head with different values; the config's tie claim is wrong text_encoder = Mistral3TextEncoder.load( - MISTRAL_PATH, subfolder="", **self.component_load_kwargs("te") + MISTRAL_PATH, + subfolder="", + tie_word_embeddings=False, + **self.component_load_kwargs("te"), ) flush() - tokenizer = AutoProcessor.from_pretrained(MISTRAL_PATH) + # fix_mistral_regex=False: keep the exact tokenization flux2 has always + # used (True would change the pre-tokenizer and shift conditioning) + tokenizer = AutoProcessor.from_pretrained( + MISTRAL_PATH, fix_mistral_regex=False + ) return text_encoder, tokenizer def load_model(self): diff --git a/extensions_built_in/diffusion_models/flux2/src/pipeline.py b/extensions_built_in/diffusion_models/flux2/src/pipeline.py index aec4dd0..fd7e619 100644 --- a/extensions_built_in/diffusion_models/flux2/src/pipeline.py +++ b/extensions_built_in/diffusion_models/flux2/src/pipeline.py @@ -109,15 +109,19 @@ class Flux2Pipeline(DiffusionPipeline): # Process all messages at once # with image processing a too short max length can throw an error in here. try: + # tokenization kwargs ride in processor_kwargs (same values end up + # in the same place; loose **kwargs just warn on new transformers) inputs = self.tokenizer.apply_chat_template( messages_batch, add_generation_prompt=False, tokenize=True, return_dict=True, return_tensors="pt", - padding="max_length", - truncation=True, - max_length=max_sequence_length, + processor_kwargs={ + "padding": "max_length", + "truncation": True, + "max_length": max_sequence_length, + }, ) except ValueError as e: print( diff --git a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py index a67ba7f..ef053ff 100644 --- a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py +++ b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py @@ -137,11 +137,13 @@ class FluxKontextModel(BaseModel): pipe.transformer = pipe.transformer.to(self.device_torch) flush() - # just to make sure everything is on the right device and dtype - text_encoder[0].to(self.device_torch) + # low_vram: text encoders stay on cpu; get_prompt_embeds moves them + # to the gpu on demand + if not self.low_vram: + text_encoder[0].to(self.device_torch) + text_encoder[1].to(self.device_torch) text_encoder[0].requires_grad_(False) text_encoder[0].eval() - text_encoder[1].to(self.device_torch) text_encoder[1].requires_grad_(False) text_encoder[1].eval() pipe.transformer = pipe.transformer.to(self.device_torch) diff --git a/extensions_built_in/diffusion_models/ltx2/ltx2.py b/extensions_built_in/diffusion_models/ltx2/ltx2.py index 7e947d0..d2b6944 100644 --- a/extensions_built_in/diffusion_models/ltx2/ltx2.py +++ b/extensions_built_in/diffusion_models/ltx2/ltx2.py @@ -490,8 +490,10 @@ class LTX2Model(BaseModel): pipe.transformer = pipe.transformer.to(self.device_torch) flush() - # just to make sure everything is on the right device and dtype - text_encoder[0].to(self.device_torch) + # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it + # to the gpu on demand + if not self.low_vram: + text_encoder[0].to(self.device_torch) text_encoder[0].requires_grad_(False) text_encoder[0].eval() flush() diff --git a/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py b/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py index 1a39802..9694420 100644 --- a/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py +++ b/extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py @@ -144,8 +144,10 @@ class NucleusImageModel(QwenImageVAEHolderMixin, BaseModel): pipe.transformer = pipe.transformer.to(self.device_torch) flush() - # just to make sure everything is on the right device and dtype - text_encoder[0].to(self.device_torch) + # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it + # to the gpu on demand + if not self.low_vram: + text_encoder[0].to(self.device_torch) text_encoder[0].requires_grad_(False) text_encoder[0].eval() flush() diff --git a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py index 4c4be72..3abfe91 100644 --- a/extensions_built_in/diffusion_models/qwen_image/qwen_image.py +++ b/extensions_built_in/diffusion_models/qwen_image/qwen_image.py @@ -161,8 +161,10 @@ class QwenImageModel(QwenImageVAEHolderMixin, BaseModel): pipe.transformer = pipe.transformer.to(self.device_torch) flush() - # just to make sure everything is on the right device and dtype - text_encoder[0].to(self.device_torch) + # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it + # to the gpu on demand + if not self.low_vram: + text_encoder[0].to(self.device_torch) text_encoder[0].requires_grad_(False) text_encoder[0].eval() flush() diff --git a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py index fa43a58..b6e1463 100644 --- a/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py +++ b/extensions_built_in/diffusion_models/wan22/wan22_14b_model.py @@ -95,6 +95,13 @@ class DualWanTransformer3DModel(torch.nn.Module): def dtype(self) -> torch.dtype: return self.torch_dtype + def get_offload_ignore_modules(self): + # both DiTs' fp32 modulation tables stay resident on the compute device + return ( + self.transformer_1.get_offload_ignore_modules() + + self.transformer_2.get_offload_ignore_modules() + ) + @property def config(self): return self.transformer_1.config @@ -128,17 +135,26 @@ class DualWanTransformer3DModel(torch.nn.Module): else: t_name = "transformer_2" + # memory-managed transformers own their placement (layers bounce + # from cpu per forward); whole-model swaps would haul the full + # 14B up and defeat the offloading + managed = ( + hasattr(self, "_memory_manager") + or hasattr(self.transformer_1, "_memory_manager") + or hasattr(self.transformer_2, "_memory_manager") + ) + # check if we are changing the active transformer, if so, we need to swap the one in # vram if low_vram is enabled # todo swap the loras as well if t_name != self._active_transformer_name: - if self.low_vram: + if self.low_vram and not managed: getattr(self, self._active_transformer_name).to("cpu") getattr(self, t_name).to(self.device_torch) torch.cuda.empty_cache() self._active_transformer_name = t_name - if self.transformer.device != hidden_states.device: + if not managed and self.transformer.device != hidden_states.device: if self.low_vram: # move other transformer to cpu other_tname = ( diff --git a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py index fcc2ce3..0cc7c86 100644 --- a/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py +++ b/extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py @@ -149,8 +149,10 @@ class ZetaChromaModel(BaseModel): pipe.transformer = pipe.transformer.to(self.device_torch) flush() - # just to make sure everything is on the right device and dtype - text_encoder[0].to(self.device_torch) + # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it + # to the gpu on demand + if not self.low_vram: + text_encoder[0].to(self.device_torch) text_encoder[0].requires_grad_(False) text_encoder[0].eval() flush() diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index a1f45b1..ef6ebb9 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -338,6 +338,71 @@ def run_one(arch: str, device: str, allow_download: bool) -> dict: gen_profile.disable() gen_profile.dump_stats(os.path.join(out_dir, "gen.prof")) + # ---- 100% offload forward: attach full layer offloading to the DiT and + # text encoder(s), then run a 1-step sample so the offload path (TE + # encode + denoise + decode with every layer staged from cpu) is proven + # for every arch ---- + import torch + + from toolkit.memory_management import MemoryManager + + prof.stage("offload 100%: attach") + offload_targets = [] + model_module = getattr(sd, "model", None) + if model_module is not None: + subs = [ + getattr(model_module, name, None) + for name in ("transformer_1", "transformer_2") + ] + if all(s is not None for s in subs): + # dual-DiT wrappers (wan22 14b): pipelines hold and .to() the + # sub-transformers directly, so each needs its own manager + offload_targets.extend(subs) + else: + offload_targets.append(model_module) + tes = getattr(sd, "text_encoder", None) + for te in tes if isinstance(tes, list) else [tes]: + offload_targets.append(te) + # auxiliary conditioning stacks where archs have them (ltx connectors, + # i2v vision towers) + offload_targets.append(getattr(sd, "image_encoder", None)) + offload_targets.append(getattr(getattr(sd, "pipeline", None), "connectors", None)) + attached = 0 + for m in offload_targets: + if not isinstance(m, torch.nn.Module): + continue + if type(m).__name__.startswith("Fake"): + continue + if next(m.parameters(), None) is None: + continue + get_ignore = getattr(m, "get_offload_ignore_modules", None) + ignore = get_ignore() if callable(get_ignore) else None + MemoryManager.attach( + m, + torch.device(device), + offload_percent=1.0, + ignore_modules=list(ignore or []), + ) + attached += 1 + + prof.stage("offload 100%: generate") + offload_profile = cProfile.Profile() + offload_profile.enable() + t_off0 = time.perf_counter() + # 2 steps, not 1: diffusers' shift_terminal stretch divides by + # one_minus_z[-1], which is zero for a single step — NaN timesteps and a + # NaN image would mask real numerical breakage in the offload path + gen_offload = GenerateImageConfig( + prompt="a photo of a cat sitting on a wooden table", + output_folder=out_dir, + output_ext="png", + **{**sample_kwargs, "num_inference_steps": 2}, + ) + sd.generate_images([gen_offload], **gen_kwargs) + offload_seconds = time.perf_counter() - t_off0 + offload_profile.disable() + offload_profile.dump_stats(os.path.join(out_dir, "offload.prof")) + stages = prof.finish() produced = [ @@ -350,8 +415,6 @@ def run_one(arch: str, device: str, allow_download: bool) -> dict: if not produced: raise RuntimeError(f"no output file produced in {out_dir}") - import torch - result = { "arch": arch, "status": "PASS", @@ -360,12 +423,16 @@ def run_one(arch: str, device: str, allow_download: bool) -> dict: "qtype_te": model_config.qtype_te if model_kwargs.get("quantize_te") else None, "load_seconds": round(load_seconds, 3), "gen_seconds": round(gen_seconds, 3), + "offload_seconds": round(offload_seconds, 3), + "offload_modules_attached": attached, "stages": stages, "profile": { "load_top": profile_top(load_profile), "gen_top": profile_top(gen_profile, limit=15), + "offload_top": profile_top(offload_profile, limit=15), "load_prof": os.path.join(out_dir, "load.prof"), "gen_prof": os.path.join(out_dir, "gen.prof"), + "offload_prof": os.path.join(out_dir, "offload.prof"), }, "env": { "torch_num_threads": torch.get_num_threads(), @@ -577,7 +644,8 @@ def main(): total = len(MODEL_TESTS) sweep_t0 = time.perf_counter() for i, arch in enumerate(MODEL_TESTS, start=1): - print(f"\n===== {arch} ({i}/{total}) =====") + # flush so tail -f shows headers/progress live when stdout is a file + print(f"\n===== {arch} ({i}/{total}) =====", flush=True) result_path = os.path.join(OUTPUT_ROOT, f".{arch.replace('/', '_')}.result.json") cmd = [ sys.executable, @@ -610,7 +678,8 @@ def main(): print( f">>> progress: {i}/{total} ({i * 100 // total}%) — " f"{tally['PASS']} pass, {tally['FAIL']} fail, {tally['SKIP']} skip — " - f"elapsed {elapsed / 60:.0f}m, eta ~{eta / 60:.0f}m" + f"elapsed {elapsed / 60:.0f}m, eta ~{eta / 60:.0f}m", + flush=True, ) print("\n===== summary =====") diff --git a/toolkit/dequantize.py b/toolkit/dequantize.py index 54c8ec7..8587833 100644 --- a/toolkit/dequantize.py +++ b/toolkit/dequantize.py @@ -46,6 +46,10 @@ def hacked_state_dict(self, *args, **kwargs): # hacks the state dict so we can dequantize before saving def patch_dequantization_on_save(model): + # idempotent: re-patching would capture the hacked state_dict as + # orig_state_dict and recurse on save + if getattr(model, "orig_state_dict", None) is not None: + return model.orig_state_dict = model.state_dict model.state_dict = partial(hacked_state_dict, model) diff --git a/toolkit/memory_management/manager.py b/toolkit/memory_management/manager.py index 910c426..19faf7f 100644 --- a/toolkit/memory_management/manager.py +++ b/toolkit/memory_management/manager.py @@ -6,6 +6,7 @@ from .manager_modules import ( LinearLayerMemoryManager, ConvLayerMemoryManager, OstrisLinearLayerMemoryManager, + EmbeddingLayerMemoryManager, _DEVICE_STATE, ) import random @@ -92,12 +93,59 @@ class MemoryManager: module._mm_to = module.to module.to = module._memory_manager.memory_managed_to - # add ignore modules to unmanaged list + # a fully offloaded module's parameters all live on cpu, which makes + # ModelMixin.device (and pipelines deriving their execution device + # from it) report "cpu" and plant latents/timesteps there. Report the + # compute device instead via an in-place subclass (same class-swap + # pattern as OstrisLinear/adopt_component); detach() restores it. + try: + module._mm_orig_class = module.__class__ + managed_cls = type( + module.__class__.__name__, + (module.__class__,), + { + "device": property( + lambda self: self._memory_manager.process_device, + lambda self, value: self.__dict__.__setitem__( + "_mm_device_shadow", value + ), + ) + }, + ) + # transformers keys per-class registries (e.g. the hidden-states + # capture specs) by str(model.__class__); make the subclass + # stringify identically so those lookups still hit + managed_cls.__module__ = module.__class__.__module__ + managed_cls.__qualname__ = module.__class__.__qualname__ + module.__class__ = managed_cls + except TypeError: + # exotic class layouts (__slots__ etc.): keep the original class + module._mm_orig_class = None + + # add ignore modules to unmanaged list; they must stay RESIDENT on the + # compute device (fp32 tables, pad tokens) — a model attached while + # parked on cpu would otherwise feed cpu tensors into gpu math for im in ignore_modules: module._memory_manager.unmanaged_modules.append(im) + try: + if isinstance(im, torch.Tensor): + im.data = im.data.to(device) + elif isinstance(im, torch.nn.Module): + im.to(device) + except Exception: + pass # count ignore modules as processed modules_processed = [x for x in ignore_modules] + + # weights tied to an embedding (lm_head <-> embed_tokens) must not be + # managed: pinning the linear's weight to cpu strands the (unmanaged) + # embedding that shares the same tensor + embedding_weight_ptrs = { + m.weight.data_ptr() + for m in module.modules() + if isinstance(m, torch.nn.Embedding) + } # attach to all modules for name, sub_module in module.named_modules(): for child_name, child_module in sub_module.named_modules(): @@ -110,6 +158,12 @@ class MemoryManager: # randomly skip some modules if random.random() > offload_percent: skip = True + if ( + not getattr(child_module, "is_ostris_quantized", False) + and isinstance(getattr(child_module, "weight", None), torch.Tensor) + and child_module.weight.data_ptr() in embedding_weight_ptrs + ): + skip = True if skip: module._memory_manager.unmanaged_modules.append(child_module) else: @@ -158,12 +212,40 @@ class MemoryManager: ) modules_processed.append(ara) modules_processed.append(child_module) + elif ( + isinstance(child_module, torch.nn.Embedding) + and child_module not in modules_processed + and child_module.weight.numel() + * child_module.weight.element_size() + > 64 * 1024 * 1024 + ): + # (the cpu gather is autograd-transparent, so a trainable + # embedding still gets grads — they just land on cpu) + # large frozen vocab table: cpu-resident, rows gathered on + # cpu (only the looked-up tokens cross the bus) + EmbeddingLayerMemoryManager.attach( + child_module, module._memory_manager + ) + # the cpu move replaced weight.data, so a tied lm_head no + # longer matches the (original-ptr) tied guard and gets + # managed normally — correct: the shared weight now lives + # on cpu, so the linear must bounce it + modules_processed.append(child_module) elif child_module.__class__.__name__ in UNMANAGED_MODULES or any( inc in child_module.__class__.__name__ for inc in UNMANAGED_MODULES_INCLUDES ): - # unmanaged - module._memory_manager.unmanaged_modules.append(child_module) + # unmanaged — but never re-list a module the nested walk + # already managed (a managed Embedding landing here would + # get hauled back to the gpu by memory_managed_to), and + # don't append duplicates + if ( + child_module not in modules_processed + and not hasattr(child_module, "_layer_memory_manager") + and child_module + not in module._memory_manager.unmanaged_modules + ): + module._memory_manager.unmanaged_modules.append(child_module) else: continue @@ -179,6 +261,10 @@ class MemoryManager: if not hasattr(module, "_memory_manager"): return + if getattr(module, "_mm_orig_class", None) is not None: + module.__class__ = module._mm_orig_class + del module._mm_orig_class + for unmanaged in module._memory_manager.unmanaged_modules: try: if isinstance(unmanaged, torch.Tensor): diff --git a/toolkit/memory_management/manager_modules.py b/toolkit/memory_management/manager_modules.py index d559f18..402981c 100644 --- a/toolkit/memory_management/manager_modules.py +++ b/toolkit/memory_management/manager_modules.py @@ -312,6 +312,12 @@ class _BouncingLinearFn(torch.autograd.Function): ctx.device = torch.device("cpu") return out.to(x.device) + if x.device != device: + # a pipeline that derives its execution device from the module's + # (fully offloaded -> "cpu") parameters can feed cpu activations; + # compute happens on the staged device and the output stays there + x = x.to(device, non_blocking=True) + state = _get_device_state(device) # the guard makes current_stream() (used by the staging helpers' event # waits/records) resolve to the process device; without it they hit @@ -461,6 +467,11 @@ class _BouncingConv2dFn(torch.autograd.Function): ctx.meta = ("cpu", stride, padding, dilation, groups, target_dtype) return out.to(x.device) + if x.device != device: + # cpu activations from a fully-offloaded pipeline: see + # _BouncingLinearFn.forward + x = x.to(device, non_blocking=True) + state = _get_device_state(device) # device guard: see _BouncingLinearFn.forward with torch.cuda.device(device): @@ -631,6 +642,42 @@ class BaseLayerMemoryManager: param._is_memory_managed = True +class EmbeddingLayerMemoryManager(BaseLayerMemoryManager): + """Offloads a (large, frozen) nn.Embedding by keeping the weight on cpu + and doing the row gather THERE: instead of staging a multi-GB vocab table + to the gpu, only the looked-up rows (tokens x dim, ~KBs) cross the bus. + Lookups happen once per prompt, so the cpu gather is free.""" + + def __init__(self, module: nn.Module, manager: "MemoryManager"): + super().__init__(module, manager) + + # cpu-resident weight; no pinning — the weight never crosses the bus + module.weight.data = module.weight.data.to("cpu") + + self._original_forward = module.forward + + def _mm_forward(input_ids, *args, **kwargs): + if args or kwargs: + return self._original_forward(input_ids, *args, **kwargs) + out_device = ( + input_ids.device + if input_ids.device.type == "cuda" + else self.manager.process_device + ) + out = F.embedding( + input_ids.to("cpu"), + self.module.weight, + self.module.padding_idx, + self.module.max_norm, + self.module.norm_type, + self.module.scale_grad_by_freq, + self.module.sparse, + ) + return out.to(out_device, non_blocking=True) + + module.forward = _mm_forward + + class LinearLayerMemoryManager(BaseLayerMemoryManager): def __init__( self, @@ -739,6 +786,11 @@ class OstrisLinearLayerMemoryManager(BaseLayerMemoryManager): # already resident on device return self._original_forward(x) + if x.device != device: + # cpu activations from a fully-offloaded pipeline: see + # _BouncingLinearFn.forward + x = x.to(device, non_blocking=True) + state = _get_device_state(device) d = state["depth"] idx = state["forward_clk"] diff --git a/toolkit/models/v2/PLANNING.md b/toolkit/models/v2/PLANNING.md index 773e26b..94560be 100644 --- a/toolkit/models/v2/PLANNING.md +++ b/toolkit/models/v2/PLANNING.md @@ -470,11 +470,75 @@ metrics_baseline.json). Profiling findings and fixes: - posix_fadvise(WILLNEED) readahead on single-file loads (mixin._readahead) — ~10-15% on cold reads, free when warm, no-op off POSIX. +Round 2 (same day): +- ARA path (`attach_ara_and_quantize`) now quantizes each hijacked linear on + the gpu (was: wherever it lived, historically the cpu) with the same + keep_on_device policy; the uint8 extras pass follows suit. Verified with + hidream + its uint3 ARA end to end. +- `quantize_model` delegates its non-ARA branch to `quantize_module` + (keep_on_device derived from model_config); dead `_quantize_model_blocks` + removed; `patch_dequantization_on_save` made idempotent (double-patching + would have recursed on save). +- Local-first hub loads: `_local_first` wraps from_pretrained / load_config / + AutoTokenizer / AutoProcessor — fully-cached repos load with + local_files_only (13x on VAE metadata, 0.53s -> 0.04s; a failed local + attempt is instant, so uncached repos fall through to the online path at + zero cost). Tokenizer-load slowness was DIAGNOSED as these per-process hub + etag HEADs, not the use_fast=False sites. + +Round 3 — vram spike elimination (2026-08-28): +- keep_on_device's extras pass had moved the whole unquantized remainder to + the gpu in bf16 before quantizing (+7-10GB transient over final residency + on big-extras models). Now `quantize(keep_on_quantize_device=True)` + quantizes each extra layer on the gpu and leaves it — transient is one + layer; non-quantized leftovers move only at final placement. Same for the + ARA uint8 pass. +- convrot8 `quantize_` chunked: takes the weight at stored dtype + (wants_fp32_weight=False; ConvRotIntN keeps True) and rotates+quantizes in + 256MB row-chunks — a 152k-vocab TE projection's transient fell 7.6GB -> + 1.4GB. Sub-threshold layers keep the one-shot path, verified bit-identical. +- qwen_image measured: standard load peak 35.3 -> 27.5GB (below the 28.0 + baseline; process peak now set by generation); low_vram load quantizes the + whole 20B stack never holding >2.4GB gpu. + Ideas not yet done: pinned-buffer staged h2d (~+60% warm-upload bandwidth, 9.5 vs 5.8GB/s measured), overlapping h2d with quantize kernels (small — the kernels are ~2% of the pass), wan22 low_vram generate choreography (165s of -unload/reload per sample), extending keep_on_device to the legacy -quantize_model/trainer path. +unload/reload per sample), holders' legacy "just to make sure" .to(device) +lines that move quantized TEs to gpu even under low_vram (qwen's Preparing +Model stage parks 7.6GB — policy question, not a quantize spike). + +## Offload hardening + final certification (2026-08-28) + +The harness gained a per-arch 100%-offload probe (attach MemoryManager at +offload_percent=1.0 to the DiT + TEs + connectors/vision towers, then a +2-step sample — 2 not 1: diffusers' shift_terminal stretch divides by zero +at a single step). It flushed out and fixed, in the manager: +- fully-offloaded modules now REPORT the compute device (class swap with + stringification-stable __module__/__qualname__ — transformers keys its + hidden-states capture registry by str(class)); bouncing ops also pull cpu + activations to the staged device +- weights tied to embeddings (lm_head <-> embed_tokens) are never pinned out + from under the unmanaged embedding +- ignore_modules are moved to the compute device at attach +- the nested attach walk no longer double-lists managed modules as unmanaged + (a managed Embedding was getting hauled back to gpu by pipeline.to) +- NEW: bouncing embeddings — vocab tables >64MB stay cpu-resident with the + row gather done ON cpu (only looked-up rows cross the bus); tied lm_heads + fall through to normal bouncing +- wan22's dual wrapper skips its whole-14B .to() swaps when memory-managed, + and offload attaches per sub-transformer (pipelines hold those directly) + +Measured offload-phase peaks: wan22_14b 14.7 -> 1.5GB, ltx2.5 14.0 -> +11.8GB (remainder: video activations + resident VAEs/vocoder by design), +most image archs 1-4GB vs 10-53GB resident. flux2's Mistral warnings +silenced behavior-preservingly (tie_word_embeddings=False, explicit +fix_mistral_regex=False, processor_kwargs — tokenization verified +bit-identical). ConvRot skip notices gated to layers >=1M weights. + +FINAL CERTIFICATION: full sweep, 33/34 PASS, 0 FAIL, 1 SKIP (mageflow +upstream 404) — load + convrot8 quantize + generate + 100%-offload forward +for every arch on the complete optimized stack. ## Testing diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index 35641e2..fb981b6 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -139,13 +139,29 @@ class OstrisModelMixin: @classmethod def aitk_from_pretrained(cls, path, subfolder=None, dtype=None, **kwargs): - return cls.from_pretrained( - path, subfolder=subfolder, torch_dtype=dtype, **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.load_config(path, subfolder=subfolder) + 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): @@ -247,6 +263,17 @@ class OstrisModelMixin: 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( @@ -254,20 +281,16 @@ class OstrisModelMixin: "network wiring" ) status_fn("Quantizing with accuracy recovery adapter") - attach_ara_and_quantize(base_model, self, ara_path, exclude=exclude) + 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})") - 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" - ) quantize_module( self, qtype, @@ -623,8 +646,8 @@ class OstrisModelMixin: os.path.join(source, subfolder) ): subfolder = None - return AutoTokenizer.from_pretrained( - source, subfolder=subfolder or "", **kwargs + return cls._local_first( + AutoTokenizer.from_pretrained, source, subfolder=subfolder or "", **kwargs ) @classmethod @@ -645,8 +668,8 @@ class OstrisModelMixin: os.path.join(source, subfolder) ): subfolder = None - return AutoProcessor.from_pretrained( - source, subfolder=subfolder or "", **kwargs + return cls._local_first( + AutoProcessor.from_pretrained, source, subfolder=subfolder or "", **kwargs ) # ------------------------------------------------------------------ @@ -778,15 +801,17 @@ class OstrisTransformersMixin(OstrisModelMixin): @classmethod def aitk_from_pretrained(cls, path, subfolder=None, dtype=None, **kwargs): - return cls.from_pretrained( - path, subfolder=subfolder or "", torch_dtype=dtype, **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 AutoConfig.from_pretrained(path, subfolder=subfolder or "") + return cls._local_first( + AutoConfig.from_pretrained, path, subfolder=subfolder or "" + ) @classmethod def aitk_from_config(cls, config): diff --git a/toolkit/models/wan21/wan21.py b/toolkit/models/wan21/wan21.py index 73fde4c..f22c4ab 100644 --- a/toolkit/models/wan21/wan21.py +++ b/toolkit/models/wan21/wan21.py @@ -433,7 +433,10 @@ class Wan21(BaseModel): pipe.transformer = pipe.transformer.to(self.device_torch) flush() - text_encoder.to(self.device_torch) + # low_vram: the text encoder stays on cpu; get_prompt_embeds moves it + # to the gpu on demand + if not self.model_config.low_vram: + text_encoder.to(self.device_torch) text_encoder.requires_grad_(False) text_encoder.eval() pipe.transformer = pipe.transformer.to(self.device_torch) diff --git a/toolkit/util/convrot_quant.py b/toolkit/util/convrot_quant.py index ec0645a..6fc5b76 100644 --- a/toolkit/util/convrot_quant.py +++ b/toolkit/util/convrot_quant.py @@ -662,7 +662,7 @@ class ConvRotQuantizer(OstrisQuantizer): d = module.in_features rot = self._rot_for(d) if d % BLOCK != 0 or module.out_features % BLOCK != 0 or rot < 16: - if d not in _skip_warned: + if d not in _skip_warned and d * module.out_features >= 1 << 20: _skip_warned.add(d) print_acc( f"ConvRot: skipping linears with in_features={d} " @@ -1312,7 +1312,7 @@ class ConvRotInt8Quantizer(OstrisQuantizer): def can_quantize(self, module: torch.nn.Linear) -> bool: d = module.in_features if d % BLOCK != 0 or module.out_features % 8 != 0 or self._rot_for(d) < 16: - if d not in _skip_warned: + if d not in _skip_warned and d * module.out_features >= 1 << 20: _skip_warned.add(d) print_acc( f"ConvRot: skipping linears with in_features={d} " @@ -1324,7 +1324,7 @@ class ConvRotInt8Quantizer(OstrisQuantizer): # (torch._int_mm raises CUBLAS_STATUS_NOT_SUPPORTED), e.g. # omnigen2's 64 -> 2520 x_embedder key = (d, module.out_features) - if key not in _skip_warned: + if key not in _skip_warned and d * module.out_features >= 1 << 20: _skip_warned.add(key) print_acc( f"ConvRot: skipping linears with in_features={d}, " @@ -1334,9 +1334,32 @@ class ConvRotInt8Quantizer(OstrisQuantizer): return False return True + # takes the weight in its stored dtype and casts per row-chunk below: a + # full-size fp32 copy of a huge layer (e.g. a TE's vocab projection) was + # a multi-GB vram transient on top of the full-size rotated fp32 tensor + wants_fp32_weight = False + + # fp32 working-set target per chunk; layers under this keep the one-shot + # path (identical math and bit-identical output to the historical code) + _QUANT_CHUNK_BYTES = 256 * 1024 * 1024 + def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None: rot = self._rot_for(module.in_features) - q, scales = quantize_int8_rows(rotate(weight_fp32, rot)) + rows, in_f = weight_fp32.shape + chunk_rows = max(1, self._QUANT_CHUNK_BYTES // (in_f * 4)) + if rows <= chunk_rows: + q, scales = quantize_int8_rows(rotate(weight_fp32.float(), rot)) + else: + # row-chunked: rotation and per-row scaling are row-independent, + # so the fp32 transient stays ~3 chunks instead of ~4x the layer + q = torch.empty(rows, in_f, dtype=torch.int8, device=weight_fp32.device) + scales = torch.empty(rows, dtype=torch.float32, device=weight_fp32.device) + for i in range(0, rows, chunk_rows): + qc, sc = quantize_int8_rows( + rotate(weight_fp32[i : i + chunk_rows].float(), rot) + ) + q[i : i + chunk_rows] = qc + scales[i : i + chunk_rows] = sc module.register_buffer("cr8_qdata", q, persistent=False) # fp32 scales stored as a uint8 byte view (see convrot4: module.to(dtype=...) # would otherwise cast them) @@ -2270,6 +2293,10 @@ class ConvRotIntNQuantizer(ConvRotInt8Quantizer): weight grid and bitpacked storage. One instance per bit-width, shareable across modules.""" + # this class's own quantize_ still expects the fp32 copy (the parent's + # chunked-cast path does not apply to it) + wants_fp32_weight = True + def __init__(self, bits: int, rot_size: int = 256): super().__init__(rot_size=rot_size) self.bits = bits diff --git a/toolkit/util/quantize.py b/toolkit/util/quantize.py index 0b55cbe..a4b2ad7 100644 --- a/toolkit/util/quantize.py +++ b/toolkit/util/quantize.py @@ -133,6 +133,7 @@ def quantize( include: Optional[Union[str, List[str]]] = None, exclude: Optional[Union[str, List[str]]] = None, quantize_device: Optional[torch.device] = None, + keep_on_quantize_device: bool = False, ): """Quantize the specified model submodules @@ -163,6 +164,11 @@ def quantize( If provided, each module is moved to this device to quantize, then moved back to the device its weights were on initially. Lets a CPU-resident model (low vram) quantize layer-by-layer on the GPU. + keep_on_quantize_device (`bool`): + With quantize_device set: leave each layer on the quantize device after + quantizing instead of moving it back. The gpu transient is then one + bf16 layer at a time — never the whole unquantized remainder at once — + for models whose final home is that device. """ if include is not None: include = [include] if isinstance(include, str) else include @@ -237,7 +243,7 @@ def quantize( optimizer=optimizer, ) finally: - if orig_device is not None: + if orig_device is not None and not keep_on_quantize_device: # quanto replaces the module in its parent, so re-fetch by name model.get_submodule(name).to(orig_device) except Exception as e: @@ -361,10 +367,21 @@ def quantize_module( status_fn(" - quantizing extras") if keep_on_device: - # blocks already live on the gpu; bring the remainder over and run the - # extras pass there (fast kernels instead of an all-cores cpu burst) + # blocks already live on the gpu; quantize each remaining layer there + # one at a time and leave it — the transient is one bf16 layer, never + # the whole unquantized remainder at once (that transient bump could + # exceed the model's final footprint by GBs on big-extras models) + quantize( + module, + weights=quantization_type, + exclude=exclude, + quantize_device=device, + keep_on_quantize_device=True, + **quantize_kwargs, + ) + # non-quantized leftovers (norms, embeddings, buffers) follow — final + # residency, not a transient module.to(device) - quantize(module, weights=quantization_type, exclude=exclude, **quantize_kwargs) else: # cpu-resident model: quantize each extra layer with a gpu round-trip quantize( @@ -395,15 +412,35 @@ def quantize_model( # sensitive modules to keep in full precision (fnmatch patterns) exclude_modules = base_model.get_quantization_exclude_modules() or [] - if base_model.model_config.accuracy_recovery_adapter is not None: + mc = base_model.model_config + device = base_model.device_torch + keep_on_device = ( + not mc.low_vram + and not (mc.layer_offloading and mc.layer_offloading_transformer_percent > 0) + and torch.device(device).type != "cpu" + ) + + if mc.accuracy_recovery_adapter is not None: attach_ara_and_quantize( base_model, model_to_quantize, - ara_path=base_model.model_config.accuracy_recovery_adapter, + ara_path=mc.accuracy_recovery_adapter, exclude=exclude_modules, + device=device, + keep_on_device=keep_on_device, ) else: - _quantize_model_blocks(base_model, model_to_quantize, exclude_modules) + quantize_module( + model_to_quantize, + mc.qtype, + device=device, + dtype=base_model.torch_dtype, + block_names=base_model.get_transformer_block_names(), + exclude=exclude_modules, + quantize_kwargs=mc.quantize_kwargs, + status_fn=base_model.print_and_status_update, + keep_on_device=keep_on_device, + ) @torch.no_grad() @@ -412,11 +449,18 @@ def attach_ara_and_quantize( model_to_quantize: torch.nn.Module, ara_path: str, exclude: Optional[List[str]] = None, + device=None, + keep_on_device: bool = False, ): """Load an accuracy recovery adapter as a live network on the module and quantize around it (adapter-hijacked linears at the configured qtype, everything else uint8). The network lands on - base_model.accuracy_recovery_adapter.""" + base_model.accuracy_recovery_adapter. + + ``device``: quantize each hijacked linear there (gpu kernels) instead of + wherever it happens to live (historically the cpu — slow). With + ``keep_on_device`` the quantized linears stay there (final home is that + gpu); otherwise each returns to the device it came from.""" from toolkit.config_modules import NetworkConfig from toolkit.lora_special import LoRASpecialNetwork @@ -537,12 +581,19 @@ def attach_ara_and_quantize( base_model.accuracy_recovery_adapter = network # quantize it + keep_on_device = keep_on_device and device is not None lora_exclude_modules = [] quantization_type = get_qtype(base_model.model_config.qtype) for lora_module in tqdm(network.unet_loras, desc="Attaching quantization"): # the lora has already hijacked the original module orig_module = lora_module.org_module[0] - orig_module.to(base_model.torch_dtype) + orig_device = None + if device is not None: + param = next(orig_module.parameters(), None) + orig_device = param.device if param is not None else None + orig_module.to(device, dtype=base_model.torch_dtype) + else: + orig_module.to(base_model.torch_dtype) # make the params not require gradients for param in orig_module.parameters(): param.requires_grad = False @@ -550,78 +601,23 @@ def attach_ara_and_quantize( freeze(orig_module) module_name = lora_module.lora_name.replace('$$', '.').replace('transformer.', '') lora_exclude_modules.append(module_name) - if base_model.model_config.low_vram: - # move it back to cpu + if not keep_on_device and orig_device is not None: + orig_module.to(orig_device) + elif base_model.model_config.low_vram and device is None: + # legacy behavior when no quantize device was given orig_module.to("cpu") - pass # quantize additional layers print_acc(" - quantizing additional layers") quantization_type = get_qtype('uint8') quantize( model_to_quantize, weights=quantization_type, - exclude=lora_exclude_modules + exclude_modules + exclude=lora_exclude_modules + exclude_modules, + quantize_device=device, + keep_on_quantize_device=keep_on_device, ) + if keep_on_device: + # non-quantized leftovers follow — final residency, not a transient + model_to_quantize.to(device) -def _quantize_model_blocks(base_model, model_to_quantize, exclude_modules): - if True: - # quantize model the original way without an accuracy recovery adapter - # move and quantize only certain pieces at a time. - quantization_type = get_qtype(base_model.model_config.qtype) - quantize_kwargs = base_model.model_config.quantize_kwargs or {} - # all_blocks = list(model_to_quantize.transformer_blocks) - all_blocks: List[torch.nn.Module] = [] - transformer_block_names = base_model.get_transformer_block_names() or [] - for name in transformer_block_names: - # name may be a dotted path for models that nest their blocks - # (e.g. hidream_o1's "model.language_model.layers"). - block_list = model_to_quantize - for part in name.split('.'): - block_list = getattr(block_list, part, None) - if block_list is None: - break - if block_list is not None: - all_blocks += list(block_list) - base_model.print_and_status_update( - f" - quantizing {len(all_blocks)} transformer blocks" - ) - already_quantized = 0 - for block in tqdm(all_blocks): - if not _has_quantizable_linear(block, quantization_type, exclude_modules): - # pre-quantized checkpoint with a matching qtype: nothing in - # this block would change — skip the device round-trip and the - # dtype cast entirely so the load stays byte-identical - already_quantized += 1 - continue - block.to(base_model.device_torch, dtype=base_model.torch_dtype, non_blocking=True) - # exclude patterns with a leading wildcard (e.g. "*adaln_proj*") - # also apply inside blocks, where names are block-relative - quantize( - block, - weights=quantization_type, - exclude=exclude_modules, - **quantize_kwargs, - ) - freeze(block) - # 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("cpu") - if already_quantized: - base_model.print_and_status_update( - f" - {already_quantized} blocks already quantized with a matching qtype; left untouched" - ) - - # todo, on extras find a universal way to quantize them on device and move them back to their original - # device without having to move the transformer blocks to the device first - base_model.print_and_status_update(" - quantizing extras") - # model_to_quantize.to(base_model.device_torch, dtype=base_model.torch_dtype) - quantize( - model_to_quantize, - weights=quantization_type, - exclude=exclude_modules, - **quantize_kwargs, - ) - freeze(model_to_quantize) From bb55d38958260b5a49d57ffca54c303b6820df16 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sat, 29 Aug 2026 08:58:15 -0600 Subject: [PATCH 14/14] Version bump --- version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/version.py b/version.py index aab9f70..364338e 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.12.27" +VERSION = "0.13.0"