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