github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Phase 2

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-27 18:18:31 -06:00
parent 45886f01b2
commit 520d96aac3
27 changed files with 333 additions and 253 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
extensions_built_in/diffusion_models/anima/anima.py
View File

@@ -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()

36
extensions_built_in/diffusion_models/chroma/chroma_model.py
View File

@@ -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')

63
extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py
View File

@@ -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')

17
extensions_built_in/diffusion_models/f_light/f_light.py
View File

@@ -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)

14
extensions_built_in/diffusion_models/flux2/flux2_model.py
View File

@@ -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

4
extensions_built_in/diffusion_models/flux2/src/model.py
View File

@@ -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()

17
extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py
View File

@@ -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)

7
extensions_built_in/diffusion_models/ideogram4/ideogram4.py
View File

@@ -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()

4
extensions_built_in/diffusion_models/ideogram4/src/transformer.py
View File

@@ -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:

12
extensions_built_in/diffusion_models/krea2/krea2.py
View File

@@ -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

4
extensions_built_in/diffusion_models/krea2/src/mmdit.py
View File

@@ -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

30
extensions_built_in/diffusion_models/ltx2/ltx2.py
View File

@@ -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):

13
extensions_built_in/diffusion_models/mageflow/mageflow.py
View File

@@ -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()

4
extensions_built_in/diffusion_models/mageflow/src/transformer.py
View File

@@ -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

42
extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py
View File

@@ -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

28
extensions_built_in/diffusion_models/minimax_h3/src/transformer.py
View File

@@ -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:

8
extensions_built_in/diffusion_models/omnigen2/__init__.py
View File

@@ -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

21
extensions_built_in/diffusion_models/wan22/wan22_14b_model.py
View File

@@ -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],

11
extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py
View File

@@ -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)

4
extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_transformer.py
View File

@@ -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()

80
toolkit/models/v2/PLANNING.md
View File

@@ -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.

41
toolkit/models/v2/_mixin.py
View File

@@ -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

10
toolkit/models/v2/diffusion_models/qwen_image.py
View File

@@ -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

51
toolkit/models/v2/diffusion_models/wan.py
View File

@@ -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

18
toolkit/models/wan21/wan21.py
View File

@@ -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')

28
toolkit/util/get_model.py
View File

@@ -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)}"
)

17
toolkit/util/quantize.py
View File

@@ -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)