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

Added legacy paths

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-28 07:55:19 -06:00
parent c3bc8b0b4e
commit 702254688d
45 changed files with 840 additions and 223 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
extensions_built_in/diffusion_models/anima/anima.py
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

10
extensions_built_in/diffusion_models/boogu_image/boogu_image.py
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)

9
extensions_built_in/diffusion_models/ernie_image/ernie_image.py
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()

6
extensions_built_in/diffusion_models/flux2/flux2_klein_model.py
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")

25
extensions_built_in/diffusion_models/flux2/flux2_model.py
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()

6
extensions_built_in/diffusion_models/hidream/hidream_model.py
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)

5
extensions_built_in/diffusion_models/hidream/hidream_o1_model.py
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,
)

8
extensions_built_in/diffusion_models/ideogram4/ideogram4.py
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()

12
extensions_built_in/diffusion_models/ltx2/convert_ltx2_to_diffusers.py
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

83
extensions_built_in/diffusion_models/ltx2/ltx2.py
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).
transformer = transformer.to(dtype)
if num_quantized_dit == 0:
transformer = transformer.to(dtype)
else:
# mixed-precision comfy file (fp32 scale_shift tables next to bf16
# weights): comfy-style per-op input casting, with the stored-fp32
# pieces pinned so parent .to(dtype) casts cannot downcast them
attach_per_op_casting(transformer)
pin_stored_fp32(transformer)
flush()
if self.model_config.quantize:
@@ -1424,7 +1417,11 @@ class LTX25Model(LTX2Model):
connectors, connectors_sd, "connectors"
)
del connectors_sd, dit_sd
connectors = connectors.to(dtype)
if num_quantized_connectors == 0:
connectors = connectors.to(dtype)
else:
attach_per_op_casting(connectors)
pin_stored_fp32(connectors)
flush()
# ---- text encoder (Gemma-4 12B, single comfy file) ----

4
extensions_built_in/diffusion_models/mageflow/mageflow.py
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)
),

30
extensions_built_in/diffusion_models/mageflow/src/vae.py
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,20 +602,32 @@ class MageVAE(nn.Module):
latent_channels = 128
downsample_factor = 16
def __init__(self, ckpt_path: str, sample_posterior: bool = True):
def __init__(self, ckpt_path: str = None, sample_posterior: bool = True):
super().__init__()
self.sample_posterior = sample_posterior
self.dconv_encoder = _DConvEncoder()
self.decoder_model = _DConvDenoiser()
sd = _load_state_dict(ckpt_path)
self._load_encoder(sd, ckpt_path)
self._load_decoder(sd, ckpt_path)
if ckpt_path is not None:
sd = _load_state_dict(ckpt_path)
self._load_encoder(sd, ckpt_path)
self._load_decoder(sd, ckpt_path)
# adaLN modulation depends only on t, and we always run at t=0.
# Precompute and drop the MLPs once at construction (~37M params saved).
self._freeze_adaln_cache()
# adaLN modulation depends only on t, and we always run at t=0.
# Precompute and drop the MLPs once at construction (~37M params saved).
self._freeze_adaln_cache()
@classmethod
def load_from_state_dict(cls, state_dict, dtype=None, **kwargs):
model = cls(
ckpt_path=None,
sample_posterior=kwargs.pop("sample_posterior", True),
)
model._load_encoder(state_dict, "<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."

40
extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py
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)

23
extensions_built_in/diffusion_models/minimax_h3/src/audio_vae.py
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)``

21
extensions_built_in/diffusion_models/minimax_h3/src/vae.py
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,

12
extensions_built_in/diffusion_models/omnigen2/__init__.py
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")

5
extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py
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)

2
extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py
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