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

Added legacy paths

Browse Source
This commit is contained in:
Jaret Burkett committed 2026-08-28 07:55:19 -06:00
1 parent c3bc8b0b4e
commit 702254688d
45 files changed
+832 -215
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

extensions_built_in/audio_models/ace_step/src/model.py
+67 -52
View File
@@ -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...")
extensions_built_in/diffusion_models/anima/anima.py
+30 -8
View File
@@ -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
extensions_built_in/diffusion_models/boogu_image/boogu_image.py
+5 -5
View File
@@ -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)
extensions_built_in/diffusion_models/ernie_image/ernie_image.py
+5 -4
View File
@@ -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()
extensions_built_in/diffusion_models/flux2/flux2_klein_model.py
+3 -3
View File
@@ -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")
extensions_built_in/diffusion_models/flux2/flux2_model.py
+5 -20
View File
@@ -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()
extensions_built_in/diffusion_models/hidream/hidream_model.py
+4 -2
View File
@@ -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)
extensions_built_in/diffusion_models/hidream/hidream_o1_model.py
+3 -2
View File
@@ -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,
)
extensions_built_in/diffusion_models/ideogram4/ideogram4.py
+4 -4
View File
@@ -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()
extensions_built_in/diffusion_models/ltx2/convert_ltx2_to_diffusers.py
+8 -4
View File
@@ -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
extensions_built_in/diffusion_models/ltx2/ltx2.py
+38 -41
View File
@@ -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).
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
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) ----
extensions_built_in/diffusion_models/mageflow/mageflow.py
+2 -2
View File
@@ -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)
),
extensions_built_in/diffusion_models/mageflow/src/vae.py
+16 -2
View File
@@ -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,13 +602,14 @@ 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()
if ckpt_path is not None:
sd = _load_state_dict(ckpt_path)
self._load_encoder(sd, ckpt_path)
self._load_decoder(sd, ckpt_path)
@@ -615,6 +618,17 @@ class MageVAE(nn.Module):
# 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, "<state dict>")
model._load_decoder(state_dict, "<state dict>")
model._freeze_adaln_cache()
return model
def _load_encoder(self, sd, ckpt_path):
prefix = "student.dconv_encoder."
enc_sd = {k[len(prefix) :]: v for k, v in sd.items() if k.startswith(prefix)}
extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py
+6 -34
View File
@@ -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)
extensions_built_in/diffusion_models/minimax_h3/src/audio_vae.py
+22 -1
View File
@@ -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)``
extensions_built_in/diffusion_models/minimax_h3/src/vae.py
+20 -1
View File
@@ -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,
extensions_built_in/diffusion_models/omnigen2/__init__.py
+7 -5
View File
@@ -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")
extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py
+2 -3
View File
@@ -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)
extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py
+1 -1
View File
@@ -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
testing/test_model_loading.py
+44 -4
View File
@@ -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
toolkit/models/v2/PLANNING.md
+43 -2
View File
@@ -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.
toolkit/models/v2/_mixin.py
+22
View File
@@ -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
toolkit/models/v2/diffusion_models/auraflow.py
+13
View File
@@ -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"]
toolkit/models/v2/diffusion_models/cosmos.py
+15
View File
@@ -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"]
toolkit/models/v2/diffusion_models/ltx2.py
+34
View File
@@ -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"
toolkit/models/v2/diffusion_models/lumina2.py
+13
View File
@@ -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"]
toolkit/models/v2/diffusion_models/pixart.py
+24
View File
@@ -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"]
toolkit/models/v2/diffusion_models/sd3.py
+11
View File
@@ -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"]
toolkit/models/v2/diffusion_models/unet.py
+13
View File
@@ -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"]
toolkit/models/v2/text_encoders/anima.py
+9
View File
@@ -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"
toolkit/models/v2/text_encoders/gemma2.py
+14
View File
@@ -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"]
toolkit/models/v2/text_encoders/gemma3.py
+32
View File
@@ -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
toolkit/models/v2/text_encoders/llama.py
+14
View File
@@ -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"]
toolkit/models/v2/text_encoders/mistral3.py
+25
View File
@@ -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"]
toolkit/models/v2/text_encoders/qwen3.py
+12 -1
View File
@@ -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"]
toolkit/models/v2/text_encoders/qwen3_vl.py
+37 -1
View File
@@ -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"]
toolkit/models/v2/vae/autoencoder_kl_flux2.py
+10
View File
@@ -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"
toolkit/models/v2/vae/flux2_kl.py
+10 -1
View File
@@ -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
toolkit/models/v2/vae/ltx2.py
+16
View File
@@ -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"
toolkit/models/v2/vae/wan.py
+9
View File
@@ -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"
toolkit/models/v2/vision_encoders/clip_vision.py
+13
View File
@@ -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"]
toolkit/models/wan21/wan21.py
+3 -4
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.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")
toolkit/models/wan21/wan21_i2v.py
+5 -8
View File
@@ -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()
toolkit/stable_diffusion_model.py
+50
View File
@@ -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")
toolkit/util/mixed_precision.py
+93
View File
@@ -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