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

Phase 1

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-27 11:53:08 -06:00
parent e8d9cf6d35
commit 8db198ec0a
51 changed files with 985 additions and 1273 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
extensions_built_in/diffusion_models/boogu_image/boogu_image.py
View File

@@ -30,6 +30,7 @@ from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
from toolkit.basic import flush
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from toolkit.models.base_model import BaseModel
from toolkit.models.v2.text_encoders.qwen3_vl import patch_qwen_vl_patch_embed
from toolkit.samplers.custom_flowmatch_sampler import (
CustomFlowMatchEulerDiscreteScheduler,
)
@@ -71,33 +72,6 @@ SYSTEM_PROMPT_T2I = (
HF_TOKEN = os.getenv("HF_TOKEN", None)
def patch_qwen_vl_patch_embed(model) -> int:
"""Swap Qwen-VL's vision ``patch_embed`` Conv3d for the equivalent ``F.linear``.
Qwen-VL's patch_embed is a Conv3d whose kernel == stride, i.e. just a linear
projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and
falls back to a slow path that effectively locks up image caching for the edit
(TI2I) model. The weight is read lazily so this survives later ``.to()`` moves.
Returns the number of patch_embed modules patched. (Vendored from
extensions_built_in/captioner/Qwen3VLCaptioner.py.)
"""
patched = 0
for module in model.modules():
proj = getattr(module, "proj", None)
if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple(
proj.stride
):
def fast_forward(hidden_states, _proj=proj):
w = _proj.weight.reshape(_proj.weight.shape[0], -1)
x = hidden_states.view(-1, w.shape[1]).to(w.dtype)
return F.linear(x, w, _proj.bias)
module.forward = fast_forward
patched += 1
return patched
class BooguImageModel(BaseModel):
arch = "boogu_image"
# Default HF repo when model.name_or_path is unset (overridden by the edit model).
@@ -167,8 +141,8 @@ class BooguImageModel(BaseModel):
# deserialize -- use the bf16 repo and let ai-toolkit quantize if wanted.
self.print_and_status_update("Loading transformer")
try:
transformer = BooguImageTransformer2DModel.from_pretrained(
base, subfolder="transformer", torch_dtype=dtype, token=HF_TOKEN
transformer = BooguImageTransformer2DModel.load_model(
base, dtype=dtype, token=HF_TOKEN
)
except OSError as e:
raise OSError(

11
extensions_built_in/diffusion_models/boogu_image/src/transformer.py
View File

@@ -21,6 +21,8 @@ from typing import Any, Dict, List, Optional, Tuple, Union
import torch
import torch.nn as nn
from diffusers.configuration_utils import ConfigMixin, register_to_config
from toolkit.models.v2._mixin import OstrisModelMixin
from diffusers.loaders import PeftAdapterMixin
from diffusers.loaders.single_file_model import FromOriginalModelMixin
from diffusers.models.attention_processor import Attention
@@ -490,10 +492,17 @@ class BooguImageDoubleStreamTransformerBlock(nn.Module):
class BooguImageTransformer2DModel(
ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin
ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, OstrisModelMixin
):
"""Boogu-Image transformer with mixed double-stream -> single-stream topology."""
aitk_subfolder = "transformer"
@classmethod
def get_transformer_block_names(cls):
return ["double_stream_layers", "single_stream_layers"]
_supports_gradient_checkpointing = True
_no_split_modules = [
"BooguImageTransformerBlock",

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

@@ -5,8 +5,9 @@ import torch
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from PIL import Image
from toolkit.models.base_model import BaseModel
from toolkit.models.v2.text_encoders.t5 import T5TextEncoder
from toolkit.models.v2.vae.autoencoder_kl import KLVAE
from toolkit.basic import flush
from diffusers import AutoencoderKL
# from toolkit.pixel_shuffle_encoder import AutoencoderPixelMixer
from toolkit.prompt_utils import PromptEmbeds
from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
@@ -14,7 +15,6 @@ from toolkit.dequantize import patch_dequantization_on_save
from toolkit.accelerator import unwrap_model
from optimum.quanto import freeze, QTensor
from toolkit.util.quantize import quantize, get_qtype
from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer
from .pipeline import ChromaPipeline, prepare_latent_image_ids
from einops import rearrange, repeat
import random
@@ -148,37 +148,17 @@ class ChromaModel(BaseModel):
self.print_and_status_update("Loading transformer")
chroma_state_dict = load_file(model_path, 'cpu')
# determine number of double and single blocks
double_blocks = 0
single_blocks = 0
for key in chroma_state_dict.keys():
if "double_blocks" in key:
block_num = int(key.split(".")[1]) + 1
if block_num > double_blocks:
double_blocks = block_num
elif "single_blocks" in key:
block_num = int(key.split(".")[1]) + 1
if block_num > single_blocks:
single_blocks = block_num
print(f"Double Blocks: {double_blocks}")
print(f"Single Blocks: {single_blocks}")
chroma_params.depth = double_blocks
chroma_params.depth_single_blocks = single_blocks
transformer = Chroma(chroma_params)
if model_path.endswith(".safetensors"):
transformer = Chroma.load_model(model_path, dtype=dtype)
else:
transformer = Chroma.load_from_state_dict(load_file(model_path, "cpu"), dtype)
# add dtype, not sure why it doesnt have it
transformer.dtype = dtype
# load the state dict into the model
transformer.load_state_dict(chroma_state_dict)
transformer.to(self.quantize_device, dtype=dtype)
transformer.config = FakeConfig()
transformer.config.num_layers = double_blocks
transformer.config.num_single_layers = single_blocks
transformer.config.num_layers = transformer.params.depth
transformer.config.num_single_layers = transformer.params.depth_single_blocks
if self.model_config.quantize:
# patch the state dict method
@@ -195,21 +175,9 @@ class ChromaModel(BaseModel):
flush()
self.print_and_status_update("Loading T5")
tokenizer_2 = T5TokenizerFast.from_pretrained(
extras_path, subfolder="tokenizer_2", torch_dtype=dtype
)
text_encoder_2 = T5EncoderModel.from_pretrained(
extras_path, subfolder="text_encoder_2", torch_dtype=dtype
)
text_encoder_2.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing T5")
quantize(text_encoder_2, weights=get_qtype(
self.model_config.qtype))
freeze(text_encoder_2)
flush()
tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path)
text_encoder_2 = T5TextEncoder.load_model(extras_path, dtype=dtype)
self.prepare_text_encoder(text_encoder_2, dtype=dtype)
# self.print_and_status_update("Loading CLIP")
text_encoder = FakeCLIP()
@@ -219,12 +187,7 @@ class ChromaModel(BaseModel):
self.noise_scheduler = ChromaModel.get_train_scheduler()
self.print_and_status_update("Loading VAE")
vae = AutoencoderKL.from_pretrained(
extras_path,
subfolder="vae",
torch_dtype=dtype
)
vae = vae.to(self.device_torch, dtype=dtype)
vae = KLVAE.load_model(extras_path, dtype=dtype, device=self.device_torch)
self.print_and_status_update("Making pipe")

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

@@ -5,8 +5,8 @@ import torch
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from PIL import Image
from toolkit.models.base_model import BaseModel
from toolkit.models.v2.text_encoders.t5 import T5TextEncoder
from toolkit.basic import flush
from diffusers import AutoencoderKL
# from toolkit.pixel_shuffle_encoder import AutoencoderPixelMixer
from toolkit.prompt_utils import PromptEmbeds
from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
@@ -14,7 +14,6 @@ from toolkit.dequantize import patch_dequantization_on_save
from toolkit.accelerator import unwrap_model
from optimum.quanto import freeze, QTensor
from toolkit.util.quantize import quantize, get_qtype
from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer
from .pipeline import ChromaPipeline, prepare_latent_image_ids
from einops import rearrange, repeat
import random
@@ -152,38 +151,16 @@ class ChromaRadianceModel(BaseModel):
if model_path.endswith('.pth') or model_path.endswith('.pt'):
chroma_state_dict = torch.load(model_path, map_location='cpu', weights_only=True)
transformer = Chroma.load_from_state_dict(chroma_state_dict, dtype)
else:
chroma_state_dict = load_file(model_path, 'cpu')
# determine number of double and single blocks
double_blocks = 0
single_blocks = 0
for key in chroma_state_dict.keys():
if "double_blocks" in key:
block_num = int(key.split(".")[1]) + 1
if block_num > double_blocks:
double_blocks = block_num
elif "single_blocks" in key:
block_num = int(key.split(".")[1]) + 1
if block_num > single_blocks:
single_blocks = block_num
print(f"Double Blocks: {double_blocks}")
print(f"Single Blocks: {single_blocks}")
chroma_params.depth = double_blocks
chroma_params.depth_single_blocks = single_blocks
transformer = Chroma(chroma_params)
transformer = Chroma.load_model(model_path, dtype=dtype)
# add dtype, not sure why it doesnt have it
transformer.dtype = dtype
# load the state dict into the model
transformer.load_state_dict(chroma_state_dict)
transformer.to(self.quantize_device, dtype=dtype)
transformer.config = FakeConfig()
transformer.config.num_layers = double_blocks
transformer.config.num_single_layers = single_blocks
transformer.config.num_layers = transformer.params.depth
transformer.config.num_single_layers = transformer.params.depth_single_blocks
if self.model_config.quantize:
# patch the state dict method
@@ -200,21 +177,9 @@ class ChromaRadianceModel(BaseModel):
flush()
self.print_and_status_update("Loading T5")
tokenizer_2 = T5TokenizerFast.from_pretrained(
extras_path, subfolder="tokenizer_2", torch_dtype=dtype
)
text_encoder_2 = T5EncoderModel.from_pretrained(
extras_path, subfolder="text_encoder_2", torch_dtype=dtype
)
text_encoder_2.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing T5")
quantize(text_encoder_2, weights=get_qtype(
self.model_config.qtype))
freeze(text_encoder_2)
flush()
tokenizer_2 = T5TextEncoder.load_tokenizer(extras_path)
text_encoder_2 = T5TextEncoder.load_model(extras_path, dtype=dtype)
self.prepare_text_encoder(text_encoder_2, dtype=dtype)
# self.print_and_status_update("Loading CLIP")
text_encoder = FakeCLIP()

35
extensions_built_in/diffusion_models/chroma/src/model.py
View File

@@ -1,4 +1,6 @@
from dataclasses import dataclass
from dataclasses import dataclass, replace
from toolkit.models.v2._mixin import OstrisModelMixin
import torch
from torch import Tensor, nn
@@ -86,11 +88,40 @@ def modify_mask_to_attend_padding(mask, max_seq_length, num_extra_padding=8):
return modified_mask
class Chroma(nn.Module):
class Chroma(nn.Module, OstrisModelMixin):
"""
Transformer model for flow matching on sequences.
"""
@classmethod
def aitk_config_from_state_dict(cls, state_dict):
# block counts come from the checkpoint's key indices
double_blocks = 0
single_blocks = 0
for key in state_dict.keys():
if "double_blocks" in key:
block_num = int(key.split(".")[1]) + 1
if block_num > double_blocks:
double_blocks = block_num
elif "single_blocks" in key:
block_num = int(key.split(".")[1]) + 1
if block_num > single_blocks:
single_blocks = block_num
print(f"Double Blocks: {double_blocks}")
print(f"Single Blocks: {single_blocks}")
return replace(
chroma_params, depth=double_blocks, depth_single_blocks=single_blocks
)
@classmethod
def aitk_from_config(cls, config):
with torch.device("meta"):
return cls(config)
@classmethod
def get_transformer_block_names(cls):
return ["double_blocks", "single_blocks"]
def __init__(self, params: ChromaParams):
super().__init__()
self.params = params

35
extensions_built_in/diffusion_models/chroma/src/radiance.py
View File

@@ -1,4 +1,6 @@
from dataclasses import dataclass
from dataclasses import dataclass, replace
from toolkit.models.v2._mixin import OstrisModelMixin
import torch
from torch import Tensor, nn
@@ -100,11 +102,40 @@ def modify_mask_to_attend_padding(mask, max_seq_length, num_extra_padding=8):
return modified_mask
class Chroma(nn.Module):
class Chroma(nn.Module, OstrisModelMixin):
"""
Transformer model for flow matching on sequences.
"""
@classmethod
def aitk_config_from_state_dict(cls, state_dict):
# block counts come from the checkpoint's key indices
double_blocks = 0
single_blocks = 0
for key in state_dict.keys():
if "double_blocks" in key:
block_num = int(key.split(".")[1]) + 1
if block_num > double_blocks:
double_blocks = block_num
elif "single_blocks" in key:
block_num = int(key.split(".")[1]) + 1
if block_num > single_blocks:
single_blocks = block_num
print(f"Double Blocks: {double_blocks}")
print(f"Single Blocks: {single_blocks}")
return replace(
chroma_params, depth=double_blocks, depth_single_blocks=single_blocks
)
@classmethod
def aitk_from_config(cls, config):
with torch.device("meta"):
return cls(config)
@classmethod
def get_transformer_block_names(cls):
return ["double_blocks", "single_blocks"]
def __init__(self, params: ChromaParams):
super().__init__()
self.params = params

10
extensions_built_in/diffusion_models/ernie_image/ernie_image.py
View File

@@ -79,20 +79,14 @@ class ErnieImageModel(BaseModel):
self.print_and_status_update("Loading transformer")
transformer_path = model_path
transformer_subfolder = "transformer"
if os.path.exists(transformer_path):
transformer_subfolder = None
transformer_path = os.path.join(transformer_path, "transformer")
if os.path.exists(model_path):
# check if the path is a full checkpoint.
te_folder_path = os.path.join(model_path, "text_encoder")
# if we have the te, this folder is a full checkpoint, use it as the base
if os.path.exists(te_folder_path):
base_model_path = model_path
transformer = ErnieImageTransformer2DModel.from_pretrained(
transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype
)
transformer = ErnieImageTransformer2DModel.load_model(model_path, dtype=dtype)
if self.model_config.quantize:
self.print_and_status_update("Quantizing Transformer")

10
extensions_built_in/diffusion_models/ernie_image/transformer.py
View File

@@ -33,6 +33,8 @@ from diffusers.models.attention_dispatch import dispatch_attention_fn
from diffusers.models.attention_processor import Attention
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
from diffusers.models.modeling_utils import ModelMixin
from toolkit.models.v2._mixin import OstrisModelMixin
from diffusers.models.normalization import RMSNorm
@@ -288,8 +290,14 @@ class ErnieImageAdaLNContinuous(nn.Module):
return x
class ErnieImageTransformer2DModel(ModelMixin, ConfigMixin):
class ErnieImageTransformer2DModel(ModelMixin, ConfigMixin, OstrisModelMixin):
_supports_gradient_checkpointing = True
aitk_subfolder = "transformer"
@classmethod
def get_transformer_block_names(cls):
return ["layers"]
@register_to_config
def __init__(

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

@@ -6,15 +6,15 @@ import yaml
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from PIL import Image
from toolkit.models.base_model import BaseModel
from toolkit.models.v2.text_encoders.t5 import T5TextEncoder
from toolkit.models.v2.vae.autoencoder_kl import KLVAE
from toolkit.basic import flush
from diffusers import AutoencoderKL
from toolkit.prompt_utils import PromptEmbeds
from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
from toolkit.dequantize import patch_dequantization_on_save
from toolkit.accelerator import unwrap_model
from optimum.quanto import freeze, QTensor
from toolkit.util.quantize import quantize, get_qtype
from transformers import T5TokenizerFast, T5EncoderModel
from .src import FLitePipeline, DiT
if TYPE_CHECKING:
@@ -74,11 +74,7 @@ class FLiteModel(BaseModel):
self.print_and_status_update("Loading transformer")
transformer = DiT.from_pretrained(
model_path,
subfolder="dit_model",
torch_dtype=dtype,
)
transformer = DiT.load_model(model_path, dtype=dtype)
transformer.to(self.quantize_device, dtype=dtype)
@@ -97,31 +93,16 @@ class FLiteModel(BaseModel):
flush()
self.print_and_status_update("Loading T5")
tokenizer = T5TokenizerFast.from_pretrained(
extras_path, subfolder="tokenizer", torch_dtype=dtype
tokenizer = T5TextEncoder.load_tokenizer(extras_path, subfolder="tokenizer")
text_encoder = T5TextEncoder.load_model(
extras_path, dtype=dtype, subfolder="text_encoder"
)
text_encoder = T5EncoderModel.from_pretrained(
extras_path, subfolder="text_encoder", torch_dtype=dtype
)
text_encoder.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing T5")
quantize(text_encoder, weights=get_qtype(
self.model_config.qtype))
freeze(text_encoder)
flush()
self.prepare_text_encoder(text_encoder, dtype=dtype)
self.noise_scheduler = FLiteModel.get_train_scheduler()
self.print_and_status_update("Loading VAE")
vae = AutoencoderKL.from_pretrained(
extras_path,
subfolder="vae",
torch_dtype=dtype
)
vae = vae.to(self.device_torch, dtype=dtype)
vae = KLVAE.load_model(extras_path, dtype=dtype, device=self.device_torch)
self.print_and_status_update("Making pipe")

10
extensions_built_in/diffusion_models/f_light/src/model.py
View File

@@ -7,6 +7,8 @@ import torch.nn.functional as F
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
from diffusers.models.modeling_utils import ModelMixin
from toolkit.models.v2._mixin import OstrisModelMixin
from diffusers.utils.accelerate_utils import apply_forward_hook
from einops import rearrange
from peft import get_peft_model_state_dict, set_peft_model_state_dict
@@ -302,7 +304,13 @@ def apply_rotary_emb(x, cos, sin):
return torch.cat([y1, y2], 3).to(dtype=orig_dtype)
class DiT(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin): # type: ignore[misc]
class DiT(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin, OstrisModelMixin): # type: ignore[misc]
aitk_subfolder = "dit_model"
@classmethod
def get_transformer_block_names(cls):
return ["blocks"]
_supports_gradient_checkpointing = True
@register_to_config

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

@@ -21,7 +21,11 @@ from toolkit.util.quantize import quantize, get_qtype, quantize_model
from transformers import AutoProcessor, Mistral3ForConditionalGeneration
from .src.model import Flux2, Flux2Params
from .src.pipeline import Flux2Pipeline
from .src.autoencoder import AutoEncoder, AutoEncoderParams, AutoEncoderSmallDecoderParams
from toolkit.models.v2.vae.flux2_kl import (
AutoEncoder,
AutoEncoderParams,
AutoEncoderSmallDecoderParams,
)
from safetensors.torch import load_file, save_file
from PIL import Image
import torch.nn.functional as F

435
extensions_built_in/diffusion_models/flux2/src/autoencoder.py
View File

@@ -1,435 +0,0 @@
from dataclasses import dataclass, field
import torch
from einops import rearrange
from torch import Tensor, nn
import math
import torch.utils.checkpoint as ckpt
@dataclass
class AutoEncoderParams:
resolution: int = 256
in_channels: int = 3
ch: int = 128
out_ch: int = 3
ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4])
num_res_blocks: int = 2
z_channels: int = 32
@dataclass
class AutoEncoderSmallDecoderParams:
resolution: int = 256
in_channels: int = 3
ch: int = 128
ch_encoder: int = 96
out_ch: int = 3
ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4])
num_res_blocks: int = 2
z_channels: int = 32
def swish(x: Tensor) -> Tensor:
return x * torch.sigmoid(x)
class AttnBlock(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
self.in_channels = in_channels
self.norm = nn.GroupNorm(
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
)
self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
def attention(self, h_: Tensor) -> Tensor:
h_ = self.norm(h_)
q = self.q(h_)
k = self.k(h_)
v = self.v(h_)
b, c, h, w = q.shape
q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
h_ = nn.functional.scaled_dot_product_attention(q, k, v)
return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
def forward(self, x: Tensor) -> Tensor:
return x + self.proj_out(self.attention(x))
class ResnetBlock(nn.Module):
def __init__(self, in_channels: int, out_channels: int):
super().__init__()
self.in_channels = in_channels
out_channels = in_channels if out_channels is None else out_channels
self.out_channels = out_channels
self.norm1 = nn.GroupNorm(
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
)
self.conv1 = nn.Conv2d(
in_channels, out_channels, kernel_size=3, stride=1, padding=1
)
self.norm2 = nn.GroupNorm(
num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
)
self.conv2 = nn.Conv2d(
out_channels, out_channels, kernel_size=3, stride=1, padding=1
)
if self.in_channels != self.out_channels:
self.nin_shortcut = nn.Conv2d(
in_channels, out_channels, kernel_size=1, stride=1, padding=0
)
def forward(self, x):
h = x
h = self.norm1(h)
h = swish(h)
h = self.conv1(h)
h = self.norm2(h)
h = swish(h)
h = self.conv2(h)
if self.in_channels != self.out_channels:
x = self.nin_shortcut(x)
return x + h
class Downsample(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
# no asymmetric padding in torch conv, must do it ourselves
self.conv = nn.Conv2d(
in_channels, in_channels, kernel_size=3, stride=2, padding=0
)
def forward(self, x: Tensor):
pad = (0, 1, 0, 1)
x = nn.functional.pad(x, pad, mode="constant", value=0)
x = self.conv(x)
return x
class Upsample(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
self.conv = nn.Conv2d(
in_channels, in_channels, kernel_size=3, stride=1, padding=1
)
def forward(self, x: Tensor):
x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
x = self.conv(x)
return x
class Encoder(nn.Module):
def __init__(
self,
resolution: int,
in_channels: int,
ch: int,
ch_mult: list[int],
num_res_blocks: int,
z_channels: int,
):
super().__init__()
self.quant_conv = torch.nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
self.ch = ch
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
self.resolution = resolution
self.in_channels = in_channels
# downsampling
self.conv_in = nn.Conv2d(
in_channels, self.ch, kernel_size=3, stride=1, padding=1
)
curr_res = resolution
in_ch_mult = (1,) + tuple(ch_mult)
self.in_ch_mult = in_ch_mult
self.down = nn.ModuleList()
block_in = self.ch
for i_level in range(self.num_resolutions):
block = nn.ModuleList()
attn = nn.ModuleList()
block_in = ch * in_ch_mult[i_level]
block_out = ch * ch_mult[i_level]
for _ in range(self.num_res_blocks):
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
block_in = block_out
down = nn.Module()
down.block = block
down.attn = attn
if i_level != self.num_resolutions - 1:
down.downsample = Downsample(block_in)
curr_res = curr_res // 2
self.down.append(down)
# middle
self.mid = nn.Module()
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
self.mid.attn_1 = AttnBlock(block_in)
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
# end
self.norm_out = nn.GroupNorm(
num_groups=32, num_channels=block_in, eps=1e-6, affine=True
)
self.conv_out = nn.Conv2d(
block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1
)
self.gradient_checkpointing = False
def enable_gradient_checkpointing(self):
self.gradient_checkpointing = True
def forward(self, x: Tensor) -> Tensor:
# downsampling
hs = [self.conv_in(x)]
for i_level in range(self.num_resolutions):
for i_block in range(self.num_res_blocks):
if torch.is_grad_enabled() and self.gradient_checkpointing:
h = ckpt.checkpoint(self.down[i_level].block[i_block], hs[-1])
if len(self.down[i_level].attn) > 0:
h = ckpt.checkpoint(self.down[i_level].attn[i_block], h)
else:
h = self.down[i_level].block[i_block](hs[-1])
if len(self.down[i_level].attn) > 0:
h = self.down[i_level].attn[i_block](h)
hs.append(h)
if i_level != self.num_resolutions - 1:
if torch.is_grad_enabled() and self.gradient_checkpointing:
hs.append(ckpt.checkpoint(self.down[i_level].downsample, hs[-1]))
else:
hs.append(self.down[i_level].downsample(hs[-1]))
# middle
h = hs[-1]
if torch.is_grad_enabled() and self.gradient_checkpointing:
h = ckpt.checkpoint(self.mid.block_1, h)
h = ckpt.checkpoint(self.mid.attn_1, h)
h = ckpt.checkpoint(self.mid.block_2, h)
else:
h = self.mid.block_1(h)
h = self.mid.attn_1(h)
h = self.mid.block_2(h)
# end
h = self.norm_out(h)
h = swish(h)
h = self.conv_out(h)
h = self.quant_conv(h)
return h
class Decoder(nn.Module):
def __init__(
self,
ch: int,
out_ch: int,
ch_mult: list[int],
num_res_blocks: int,
in_channels: int,
resolution: int,
z_channels: int,
):
super().__init__()
self.post_quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1)
self.ch = ch
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
self.resolution = resolution
self.in_channels = in_channels
self.ffactor = 2 ** (self.num_resolutions - 1)
# compute in_ch_mult, block_in and curr_res at lowest res
block_in = ch * ch_mult[self.num_resolutions - 1]
curr_res = resolution // 2 ** (self.num_resolutions - 1)
self.z_shape = (1, z_channels, curr_res, curr_res)
# z to block_in
self.conv_in = nn.Conv2d(
z_channels, block_in, kernel_size=3, stride=1, padding=1
)
# middle
self.mid = nn.Module()
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
self.mid.attn_1 = AttnBlock(block_in)
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
# upsampling
self.up = nn.ModuleList()
for i_level in reversed(range(self.num_resolutions)):
block = nn.ModuleList()
attn = nn.ModuleList()
block_out = ch * ch_mult[i_level]
for _ in range(self.num_res_blocks + 1):
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
block_in = block_out
up = nn.Module()
up.block = block
up.attn = attn
if i_level != 0:
up.upsample = Upsample(block_in)
curr_res = curr_res * 2
self.up.insert(0, up) # prepend to get consistent order
# end
self.norm_out = nn.GroupNorm(
num_groups=32, num_channels=block_in, eps=1e-6, affine=True
)
self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
self.gradient_checkpointing = False
def enable_gradient_checkpointing(self):
self.gradient_checkpointing = True
def forward(self, z: Tensor) -> Tensor:
z = self.post_quant_conv(z)
# get dtype for proper tracing
upscale_dtype = next(self.up.parameters()).dtype
# z to block_in
h = self.conv_in(z)
# middle
if torch.is_grad_enabled() and self.gradient_checkpointing:
h = ckpt.checkpoint(self.mid.block_1, h)
h = ckpt.checkpoint(self.mid.attn_1, h)
h = ckpt.checkpoint(self.mid.block_2, h)
else:
h = self.mid.block_1(h)
h = self.mid.attn_1(h)
h = self.mid.block_2(h)
# cast to proper dtype
h = h.to(upscale_dtype)
# upsampling
for i_level in reversed(range(self.num_resolutions)):
for i_block in range(self.num_res_blocks + 1):
if torch.is_grad_enabled() and self.gradient_checkpointing:
h = ckpt.checkpoint(self.up[i_level].block[i_block], h)
if len(self.up[i_level].attn) > 0:
h = ckpt.checkpoint(self.up[i_level].attn[i_block], h)
else:
h = self.up[i_level].block[i_block](h)
if len(self.up[i_level].attn) > 0:
h = self.up[i_level].attn[i_block](h)
if i_level != 0:
if torch.is_grad_enabled() and self.gradient_checkpointing:
h = ckpt.checkpoint(self.up[i_level].upsample, h)
else:
h = self.up[i_level].upsample(h)
# end
h = self.norm_out(h)
h = swish(h)
h = self.conv_out(h)
return h
class AutoEncoder(nn.Module):
def __init__(self, params: AutoEncoderParams):
super().__init__()
self.params = params
self.encoder = Encoder(
resolution=params.resolution,
in_channels=params.in_channels,
ch=params.ch,
ch_mult=params.ch_mult,
num_res_blocks=params.num_res_blocks,
z_channels=params.z_channels,
)
decoder_ch = params.ch
if hasattr(params, "ch_encoder"):
decoder_ch = params.ch_encoder
self.decoder = Decoder(
resolution=params.resolution,
in_channels=params.in_channels,
ch=decoder_ch,
out_ch=params.out_ch,
ch_mult=params.ch_mult,
num_res_blocks=params.num_res_blocks,
z_channels=params.z_channels,
)
self.bn_eps = 1e-4
self.bn_momentum = 0.1
self.ps = [2, 2]
self.bn = torch.nn.BatchNorm2d(
math.prod(self.ps) * params.z_channels,
eps=self.bn_eps,
momentum=self.bn_momentum,
affine=False,
track_running_stats=True,
)
self._gradient_checkpointing = False
@property
def gradient_checkpointing(self):
return self._gradient_checkpointing
@gradient_checkpointing.setter
def gradient_checkpointing(self, value: bool):
self._gradient_checkpointing = value
self.encoder.gradient_checkpointing = value
self.decoder.gradient_checkpointing = value
@property
def device(self):
return next(self.parameters()).device
@property
def dtype(self):
return next(self.parameters()).dtype
def enable_gradient_checkpointing(self):
self.gradient_checkpointing = True
self.encoder.enable_gradient_checkpointing()
self.decoder.enable_gradient_checkpointing()
def normalize(self, z):
self.bn.eval()
return self.bn(z)
def inv_normalize(self, z):
self.bn.eval()
s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.bn_eps)
m = self.bn.running_mean.view(1, -1, 1, 1)
return z * s + m
def encode(self, x: Tensor) -> Tensor:
moments = self.encoder(x)
mean = torch.chunk(moments, 2, dim=1)[0]
z = rearrange(
mean,
"... c (i pi) (j pj) -> ... (c pi pj) i j",
pi=self.ps[0],
pj=self.ps[1],
)
z = self.normalize(z)
return z
def decode(self, z: Tensor) -> Tensor:
z = self.inv_normalize(z)
z = rearrange(
z,
"... (c pi pj) i j -> ... c (i pi) (j pj)",
pi=self.ps[0],
pj=self.ps[1],
)
dec = self.decoder(z)
return dec

2
extensions_built_in/diffusion_models/flux2/src/pipeline.py
View File

@@ -12,7 +12,7 @@ from diffusers.utils import (
from diffusers.utils.torch_utils import randn_tensor
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.utils import BaseOutput
from .autoencoder import AutoEncoder
from toolkit.models.v2.vae.flux2_kl import AutoEncoder
from .model import Flux2
from einops import rearrange
from transformers import AutoProcessor, Mistral3ForConditionalGeneration

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

@@ -8,7 +8,11 @@ from toolkit import train_tools
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from PIL import Image
from toolkit.models.base_model import BaseModel
from diffusers import FluxTransformer2DModel, AutoencoderKL, FluxKontextPipeline
from toolkit.models.v2.text_encoders.t5 import T5TextEncoder
from toolkit.models.v2.text_encoders.clip import CLIPTextEncoder
from toolkit.models.v2.vae.autoencoder_kl import KLVAE
from diffusers import FluxKontextPipeline
from toolkit.models.v2.diffusion_models.flux import FluxTransformer2DModel
from toolkit.basic import flush
from toolkit.prompt_utils import PromptEmbeds
from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
@@ -18,7 +22,6 @@ from toolkit.accelerator import get_accelerator, unwrap_model
from optimum.quanto import freeze, QTensor
from toolkit.util.mask import generate_random_mask, random_dialate_mask
from toolkit.util.quantize import quantize, get_qtype
from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer
from einops import rearrange, repeat
import random
import torch.nn.functional as F
@@ -80,11 +83,7 @@ class FluxKontextModel(BaseModel):
# so we need this for the VAE, te, etc
base_model_path = self.model_config.extras_name_or_path
transformer_path = model_path
transformer_subfolder = 'transformer'
if os.path.exists(transformer_path):
transformer_subfolder = None
transformer_path = os.path.join(transformer_path, 'transformer')
if os.path.exists(model_path):
# check if the path is a full checkpoint.
te_folder_path = os.path.join(model_path, 'text_encoder')
# if we have the te, this folder is a full checkpoint, use it as the base
@@ -92,11 +91,7 @@ class FluxKontextModel(BaseModel):
base_model_path = model_path
self.print_and_status_update("Loading transformer")
transformer = FluxTransformer2DModel.from_pretrained(
transformer_path,
subfolder=transformer_subfolder,
torch_dtype=dtype
)
transformer = FluxTransformer2DModel.load_model(model_path, dtype=dtype)
transformer.to(self.quantize_device, dtype=dtype)
if self.model_config.quantize:
@@ -114,32 +109,18 @@ class FluxKontextModel(BaseModel):
flush()
self.print_and_status_update("Loading T5")
tokenizer_2 = T5TokenizerFast.from_pretrained(
base_model_path, subfolder="tokenizer_2", torch_dtype=dtype
)
text_encoder_2 = T5EncoderModel.from_pretrained(
base_model_path, subfolder="text_encoder_2", torch_dtype=dtype
)
text_encoder_2.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing T5")
quantize(text_encoder_2, weights=get_qtype(
self.model_config.qtype))
freeze(text_encoder_2)
flush()
tokenizer_2 = T5TextEncoder.load_tokenizer(base_model_path)
text_encoder_2 = T5TextEncoder.load_model(base_model_path, dtype=dtype)
self.prepare_text_encoder(text_encoder_2, dtype=dtype)
self.print_and_status_update("Loading CLIP")
text_encoder = CLIPTextModel.from_pretrained(
base_model_path, subfolder="text_encoder", torch_dtype=dtype)
tokenizer = CLIPTokenizer.from_pretrained(
base_model_path, subfolder="tokenizer", torch_dtype=dtype)
text_encoder.to(self.device_torch, dtype=dtype)
text_encoder = CLIPTextEncoder.load_model(
base_model_path, dtype=dtype, device=self.device_torch
)
tokenizer = CLIPTextEncoder.load_tokenizer(base_model_path, use_fast=False)
self.print_and_status_update("Loading VAE")
vae = AutoencoderKL.from_pretrained(
base_model_path, subfolder="vae", torch_dtype=dtype)
vae = KLVAE.load_model(base_model_path, dtype=dtype)
self.noise_scheduler = FluxKontextModel.get_train_scheduler()

2
extensions_built_in/diffusion_models/hidream/hidream_e1_model.py
View File

@@ -7,7 +7,7 @@ from toolkit.accelerator import unwrap_model
import torch
from toolkit.prompt_utils import PromptEmbeds
from toolkit.config_modules import GenerateImageConfig
from diffusers.models import HiDreamImageTransformer2DModel
from toolkit.models.v2.diffusion_models.hidream import HiDreamImageTransformer2DModel
import torch.nn.functional as F
from PIL import Image

68
extensions_built_in/diffusion_models/hidream/hidream_model.py
View File

@@ -9,7 +9,10 @@ from toolkit import train_tools
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from PIL import Image
from toolkit.models.base_model import BaseModel
from diffusers import AutoencoderKL, TorchAoConfig
from toolkit.models.v2.text_encoders.t5 import T5TextEncoder
from toolkit.models.v2.text_encoders.clip import CLIPTextEncoderWithProjection
from toolkit.models.v2.vae.autoencoder_kl import KLVAE
from diffusers import TorchAoConfig
from toolkit.basic import flush
from toolkit.prompt_utils import PromptEmbeds
from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
@@ -19,7 +22,6 @@ from toolkit.accelerator import get_accelerator, unwrap_model
from optimum.quanto import freeze, QTensor
from toolkit.util.mask import generate_random_mask, random_dialate_mask
from toolkit.util.quantize import quantize, get_qtype
from transformers import T5TokenizerFast, T5EncoderModel, CLIPTextModel, CLIPTokenizer, TorchAoConfig as TorchAoConfigTransformers
from .src.pipelines.hidream_image.pipeline_hidream_image import HiDreamImagePipeline
from .src.models.transformers.transformer_hidream_image import HiDreamImageTransformer2DModel
from .src.schedulers.fm_solvers_unipc import FlowUniPCMultistepScheduler
@@ -28,15 +30,6 @@ from einops import rearrange, repeat
import random
import torch.nn.functional as F
from tqdm import tqdm
from transformers import (
CLIPTextModelWithProjection,
CLIPTokenizer,
T5EncoderModel,
T5Tokenizer,
LlamaForCausalLM,
PreTrainedTokenizerFast
)
if TYPE_CHECKING:
from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
@@ -125,10 +118,8 @@ class HidreamModel(BaseModel):
self.print_and_status_update("Loading transformer")
transformer = self.hidream_transformer_class.from_pretrained(
model_path,
subfolder="transformer",
torch_dtype=torch.bfloat16
transformer = self.hidream_transformer_class.load_model(
model_path, dtype=torch.bfloat16
)
if not self.low_vram:
@@ -164,44 +155,34 @@ class HidreamModel(BaseModel):
self.print_and_status_update("Loading vae")
vae = AutoencoderKL.from_pretrained(
extras_path,
subfolder="vae",
torch_dtype=torch.bfloat16
).to(self.device_torch, dtype=dtype)
vae = KLVAE.load_model(extras_path, dtype=torch.bfloat16).to(
self.device_torch, dtype=dtype
)
self.print_and_status_update("Loading clip encoders")
text_encoder = CLIPTextModelWithProjection.from_pretrained(
extras_path,
subfolder="text_encoder",
torch_dtype=torch.bfloat16
text_encoder = CLIPTextEncoderWithProjection.load_model(
extras_path, dtype=torch.bfloat16
).to(self.device_torch, dtype=dtype)
tokenizer = CLIPTokenizer.from_pretrained(
extras_path,
subfolder="tokenizer"
tokenizer = CLIPTextEncoderWithProjection.load_tokenizer(
extras_path, use_fast=False
)
text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
extras_path,
subfolder="text_encoder_2",
torch_dtype=torch.bfloat16
text_encoder_2 = CLIPTextEncoderWithProjection.load_model(
extras_path, dtype=torch.bfloat16, subfolder="text_encoder_2"
).to(self.device_torch, dtype=dtype)
tokenizer_2 = CLIPTokenizer.from_pretrained(
extras_path,
subfolder="tokenizer_2"
tokenizer_2 = CLIPTextEncoderWithProjection.load_tokenizer(
extras_path, subfolder="tokenizer_2", use_fast=False
)
flush()
self.print_and_status_update("Loading T5 encoders")
text_encoder_3 = T5EncoderModel.from_pretrained(
extras_path,
subfolder="text_encoder_3",
torch_dtype=torch.bfloat16
text_encoder_3 = T5TextEncoder.load_model(
extras_path, dtype=torch.bfloat16, subfolder="text_encoder_3"
).to(self.device_torch, dtype=dtype)
if self.model_config.quantize_te:
@@ -211,9 +192,8 @@ class HidreamModel(BaseModel):
freeze(text_encoder_3)
flush()
tokenizer_3 = T5Tokenizer.from_pretrained(
extras_path,
subfolder="tokenizer_3"
tokenizer_3 = T5TextEncoder.load_tokenizer(
extras_path, subfolder="tokenizer_3", use_fast=False
)
flush()

10
extensions_built_in/diffusion_models/hidream/src/models/transformers/transformer_hidream_image.py
View File

@@ -8,6 +8,8 @@ from einops import repeat
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
from diffusers.models.modeling_utils import ModelMixin
from toolkit.models.v2._mixin import OstrisModelMixin
from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
from diffusers.utils.torch_utils import maybe_allow_in_graph
from diffusers.models.modeling_outputs import Transformer2DModelOutput
@@ -228,9 +230,15 @@ class HiDreamImageBlock(nn.Module):
)
class HiDreamImageTransformer2DModel(
ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin
ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, OstrisModelMixin
):
_supports_gradient_checkpointing = True
aitk_subfolder = "transformer"
@classmethod
def get_transformer_block_names(cls):
return ["double_stream_blocks", "single_stream_blocks"]
_no_split_modules = ["HiDreamImageBlock"]
@register_to_config

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

@@ -26,7 +26,11 @@ from huggingface_hub.errors import EntryNotFoundError
from transformers import AutoModel, AutoTokenizer
from .src.transformer import Ideogram4Config, Ideogram4Transformer2DModel
from .src.vae import AutoEncoder, AutoEncoderParams, convert_diffusers_state_dict
from toolkit.models.v2.vae.flux2_kl import (
AutoEncoder,
AutoEncoderParams,
convert_diffusers_state_dict,
)
from .src.latent_norm import get_latent_norm
from .src.pipeline import (
Ideogram4Pipeline,

482
extensions_built_in/diffusion_models/ideogram4/src/vae.py
View File

@@ -1,482 +0,0 @@
"""Flux2 KL autoencoder."""
from __future__ import annotations
import math
import re
from dataclasses import dataclass, field
import torch
import torch.utils.checkpoint as ckpt
from einops import rearrange
from torch import Tensor, nn
@dataclass
class AutoEncoderParams:
resolution: int = 256
in_channels: int = 3
ch: int = 128
out_ch: int = 3
ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4])
num_res_blocks: int = 2
z_channels: int = 32
def swish(x: Tensor) -> Tensor:
return x * torch.sigmoid(x)
class AttnBlock(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
self.in_channels = in_channels
self.norm = nn.GroupNorm(
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
)
self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
def attention(self, h_: Tensor) -> Tensor:
h_ = self.norm(h_)
q = self.q(h_)
k = self.k(h_)
v = self.v(h_)
b, c, h, w = q.shape
q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
h_ = nn.functional.scaled_dot_product_attention(q, k, v)
return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
def forward(self, x: Tensor) -> Tensor:
return x + self.proj_out(self.attention(x))
class ResnetBlock(nn.Module):
def __init__(self, in_channels: int, out_channels: int):
super().__init__()
self.in_channels = in_channels
out_channels = in_channels if out_channels is None else out_channels
self.out_channels = out_channels
self.norm1 = nn.GroupNorm(
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
)
self.conv1 = nn.Conv2d(
in_channels, out_channels, kernel_size=3, stride=1, padding=1
)
self.norm2 = nn.GroupNorm(
num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
)
self.conv2 = nn.Conv2d(
out_channels, out_channels, kernel_size=3, stride=1, padding=1
)
if self.in_channels != self.out_channels:
self.nin_shortcut = nn.Conv2d(
in_channels, out_channels, kernel_size=1, stride=1, padding=0
)
def forward(self, x):
h = x
h = self.norm1(h)
h = swish(h)
h = self.conv1(h)
h = self.norm2(h)
h = swish(h)
h = self.conv2(h)
if self.in_channels != self.out_channels:
x = self.nin_shortcut(x)
return x + h
class Downsample(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
# no asymmetric padding in torch conv, must do it ourselves
self.conv = nn.Conv2d(
in_channels, in_channels, kernel_size=3, stride=2, padding=0
)
def forward(self, x: Tensor):
pad = (0, 1, 0, 1)
x = nn.functional.pad(x, pad, mode="constant", value=0)
x = self.conv(x)
return x
class Upsample(nn.Module):
def __init__(self, in_channels: int):
super().__init__()
self.conv = nn.Conv2d(
in_channels, in_channels, kernel_size=3, stride=1, padding=1
)
def forward(self, x: Tensor):
x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
x = self.conv(x)
return x
class Encoder(nn.Module):
def __init__(
self,
resolution: int,
in_channels: int,
ch: int,
ch_mult: list[int],
num_res_blocks: int,
z_channels: int,
):
super().__init__()
self.quant_conv = torch.nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
self.ch = ch
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
self.resolution = resolution
self.in_channels = in_channels
# downsampling
self.conv_in = nn.Conv2d(
in_channels, self.ch, kernel_size=3, stride=1, padding=1
)
curr_res = resolution
in_ch_mult = (1,) + tuple(ch_mult)
self.in_ch_mult = in_ch_mult
self.down = nn.ModuleList()
block_in = self.ch
for i_level in range(self.num_resolutions):
block = nn.ModuleList()
attn = nn.ModuleList()
block_in = ch * in_ch_mult[i_level]
block_out = ch * ch_mult[i_level]
for _ in range(self.num_res_blocks):
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
block_in = block_out
down = nn.Module()
down.block = block
down.attn = attn
if i_level != self.num_resolutions - 1:
down.downsample = Downsample(block_in)
curr_res = curr_res // 2
self.down.append(down)
# middle
self.mid = nn.Module()
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
self.mid.attn_1 = AttnBlock(block_in)
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
# end
self.norm_out = nn.GroupNorm(
num_groups=32, num_channels=block_in, eps=1e-6, affine=True
)
self.conv_out = nn.Conv2d(
block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1
)
self.gradient_checkpointing = False
def enable_gradient_checkpointing(self):
self.gradient_checkpointing = True
def forward(self, x: Tensor) -> Tensor:
# downsampling
hs = [self.conv_in(x)]
for i_level in range(self.num_resolutions):
for i_block in range(self.num_res_blocks):
if torch.is_grad_enabled() and self.gradient_checkpointing:
h = ckpt.checkpoint(self.down[i_level].block[i_block], hs[-1]) # type: ignore[index, operator]
if len(self.down[i_level].attn) > 0: # type: ignore[arg-type]
h = ckpt.checkpoint(self.down[i_level].attn[i_block], h) # type: ignore[index, operator]
else:
h = self.down[i_level].block[i_block](hs[-1]) # type: ignore[index, operator]
if len(self.down[i_level].attn) > 0: # type: ignore[arg-type]
h = self.down[i_level].attn[i_block](h) # type: ignore[index, operator]
hs.append(h)
if i_level != self.num_resolutions - 1:
if torch.is_grad_enabled() and self.gradient_checkpointing:
hs.append(ckpt.checkpoint(self.down[i_level].downsample, hs[-1])) # type: ignore[operator]
else:
hs.append(self.down[i_level].downsample(hs[-1])) # type: ignore[operator]
# middle
h = hs[-1]
if torch.is_grad_enabled() and self.gradient_checkpointing:
h = ckpt.checkpoint(self.mid.block_1, h) # type: ignore[operator]
h = ckpt.checkpoint(self.mid.attn_1, h) # type: ignore[operator]
h = ckpt.checkpoint(self.mid.block_2, h) # type: ignore[operator]
else:
h = self.mid.block_1(h) # type: ignore[operator]
h = self.mid.attn_1(h) # type: ignore[operator]
h = self.mid.block_2(h) # type: ignore[operator]
# end
h = self.norm_out(h)
h = swish(h)
h = self.conv_out(h)
h = self.quant_conv(h)
return h
class Decoder(nn.Module):
def __init__(
self,
ch: int,
out_ch: int,
ch_mult: list[int],
num_res_blocks: int,
in_channels: int,
resolution: int,
z_channels: int,
):
super().__init__()
self.post_quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1)
self.ch = ch
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
self.resolution = resolution
self.in_channels = in_channels
self.ffactor = 2 ** (self.num_resolutions - 1)
# compute in_ch_mult, block_in and curr_res at lowest res
block_in = ch * ch_mult[self.num_resolutions - 1]
curr_res = resolution // 2 ** (self.num_resolutions - 1)
self.z_shape = (1, z_channels, curr_res, curr_res)
# z to block_in
self.conv_in = nn.Conv2d(
z_channels, block_in, kernel_size=3, stride=1, padding=1
)
# middle
self.mid = nn.Module()
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
self.mid.attn_1 = AttnBlock(block_in)
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
# upsampling
self.up = nn.ModuleList()
for i_level in reversed(range(self.num_resolutions)):
block = nn.ModuleList()
attn = nn.ModuleList()
block_out = ch * ch_mult[i_level]
for _ in range(self.num_res_blocks + 1):
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
block_in = block_out
up = nn.Module()
up.block = block
up.attn = attn
if i_level != 0:
up.upsample = Upsample(block_in)
curr_res = curr_res * 2
self.up.insert(0, up) # prepend to get consistent order
# end
self.norm_out = nn.GroupNorm(
num_groups=32, num_channels=block_in, eps=1e-6, affine=True
)
self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
self.gradient_checkpointing = False
def enable_gradient_checkpointing(self):
self.gradient_checkpointing = True
def forward(self, z: Tensor) -> Tensor:
z = self.post_quant_conv(z)
# get dtype for proper tracing
upscale_dtype = next(self.up.parameters()).dtype
# z to block_in
h = self.conv_in(z)
# middle
if torch.is_grad_enabled() and self.gradient_checkpointing:
h = ckpt.checkpoint(self.mid.block_1, h) # type: ignore[operator]
h = ckpt.checkpoint(self.mid.attn_1, h) # type: ignore[operator]
h = ckpt.checkpoint(self.mid.block_2, h) # type: ignore[operator]
else:
h = self.mid.block_1(h) # type: ignore[operator]
h = self.mid.attn_1(h) # type: ignore[operator]
h = self.mid.block_2(h) # type: ignore[operator]
# cast to proper dtype
h = h.to(upscale_dtype)
# upsampling
for i_level in reversed(range(self.num_resolutions)):
for i_block in range(self.num_res_blocks + 1):
if torch.is_grad_enabled() and self.gradient_checkpointing:
h = ckpt.checkpoint(self.up[i_level].block[i_block], h) # type: ignore[index, operator]
if len(self.up[i_level].attn) > 0: # type: ignore[arg-type]
h = ckpt.checkpoint(self.up[i_level].attn[i_block], h) # type: ignore[index, operator]
else:
h = self.up[i_level].block[i_block](h) # type: ignore[index, operator]
if len(self.up[i_level].attn) > 0: # type: ignore[arg-type]
h = self.up[i_level].attn[i_block](h) # type: ignore[index, operator]
if i_level != 0:
if torch.is_grad_enabled() and self.gradient_checkpointing:
h = ckpt.checkpoint(self.up[i_level].upsample, h) # type: ignore[operator]
else:
h = self.up[i_level].upsample(h) # type: ignore[operator]
# end
h = self.norm_out(h)
h = swish(h)
h = self.conv_out(h)
return h
class AutoEncoder(nn.Module):
def __init__(self, params: AutoEncoderParams):
super().__init__()
self.params = params
self.encoder = Encoder(
resolution=params.resolution,
in_channels=params.in_channels,
ch=params.ch,
ch_mult=params.ch_mult,
num_res_blocks=params.num_res_blocks,
z_channels=params.z_channels,
)
self.decoder = Decoder(
resolution=params.resolution,
in_channels=params.in_channels,
ch=params.ch,
out_ch=params.out_ch,
ch_mult=params.ch_mult,
num_res_blocks=params.num_res_blocks,
z_channels=params.z_channels,
)
self.bn_eps = 1e-4
self.bn_momentum = 0.1
self.ps = [2, 2]
self.bn = torch.nn.BatchNorm2d(
math.prod(self.ps) * params.z_channels,
eps=self.bn_eps,
momentum=self.bn_momentum,
affine=False,
track_running_stats=True,
)
self._gradient_checkpointing = False
@property
def gradient_checkpointing(self) -> bool:
return self._gradient_checkpointing
@gradient_checkpointing.setter
def gradient_checkpointing(self, value: bool):
self._gradient_checkpointing = value
self.encoder.gradient_checkpointing = value
self.decoder.gradient_checkpointing = value
def enable_gradient_checkpointing(self):
self.gradient_checkpointing = True
self.encoder.enable_gradient_checkpointing()
self.decoder.enable_gradient_checkpointing()
@property
def device(self) -> torch.device:
return next(self.parameters()).device
@property
def dtype(self) -> torch.dtype:
return next(self.parameters()).dtype
_NUM_RESOLUTIONS = 4
def convert_diffusers_state_dict(src: dict[str, Tensor]) -> dict[str, Tensor]:
out: dict[str, Tensor] = {}
attn_substrings = (".mid.attn_1.",)
for src_key, tensor in src.items():
dst_key = _rewrite_diffusers_key(src_key)
if dst_key is None:
raise KeyError(f"Unrecognized diffusers VAE state-dict key: {src_key}")
if (
any(s in dst_key for s in attn_substrings)
and dst_key.endswith(".weight")
and tensor.ndim == 2
):
tensor = tensor.unsqueeze(-1).unsqueeze(-1)
out[dst_key] = tensor
return out
def _rewrite_diffusers_key(key: str) -> str | None:
if key.startswith("bn."):
return key
if key.startswith("quant_conv."):
return key.replace("quant_conv.", "encoder.quant_conv.", 1)
if key.startswith("post_quant_conv."):
return key.replace("post_quant_conv.", "decoder.post_quant_conv.", 1)
if key == "encoder.conv_norm_out.weight":
return "encoder.norm_out.weight"
if key == "encoder.conv_norm_out.bias":
return "encoder.norm_out.bias"
if key == "decoder.conv_norm_out.weight":
return "decoder.norm_out.weight"
if key == "decoder.conv_norm_out.bias":
return "decoder.norm_out.bias"
m = re.match(r"^(encoder|decoder)\.mid_block\.resnets\.(\d+)\.(.+)$", key)
if m:
side, idx, rest = m.group(1), int(m.group(2)), m.group(3)
rest = rest.replace("conv_shortcut", "nin_shortcut")
return f"{side}.mid.block_{idx + 1}.{rest}"
m = re.match(r"^(encoder|decoder)\.mid_block\.attentions\.0\.(.+)$", key)
if m:
side, rest = m.group(1), m.group(2)
rest = (
rest.replace("group_norm.", "norm.")
.replace("to_q.", "q.")
.replace("to_k.", "k.")
.replace("to_v.", "v.")
.replace("to_out.0.", "proj_out.")
)
return f"{side}.mid.attn_1.{rest}"
m = re.match(r"^encoder\.down_blocks\.(\d+)\.resnets\.(\d+)\.(.+)$", key)
if m:
level, res_idx, rest = m.group(1), m.group(2), m.group(3)
rest = rest.replace("conv_shortcut", "nin_shortcut")
return f"encoder.down.{level}.block.{res_idx}.{rest}"
m = re.match(r"^encoder\.down_blocks\.(\d+)\.downsamplers\.0\.conv\.(.+)$", key)
if m:
return f"encoder.down.{m.group(1)}.downsample.conv.{m.group(2)}"
m = re.match(r"^decoder\.up_blocks\.(\d+)\.resnets\.(\d+)\.(.+)$", key)
if m:
diffusers_idx = int(m.group(1))
res_idx = m.group(2)
rest = m.group(3).replace("conv_shortcut", "nin_shortcut")
return (
f"decoder.up.{_NUM_RESOLUTIONS - 1 - diffusers_idx}.block.{res_idx}.{rest}"
)
m = re.match(r"^decoder\.up_blocks\.(\d+)\.upsamplers\.0\.conv\.(.+)$", key)
if m:
diffusers_idx = int(m.group(1))
return f"decoder.up.{_NUM_RESOLUTIONS - 1 - diffusers_idx}.upsample.conv.{m.group(2)}"
if key.startswith(
(
"encoder.conv_in.",
"encoder.conv_out.",
"decoder.conv_in.",
"decoder.conv_out.",
)
):
return key
return None

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

@@ -25,18 +25,21 @@ from safetensors.torch import load_file, save_file
import huggingface_hub
from huggingface_hub.errors import EntryNotFoundError
from diffusers import AutoencoderKLQwenImage
from transformers import (
AutoProcessor,
AutoTokenizer,
Qwen2TokenizerFast,
Qwen3VLForConditionalGeneration,
)
from optimum.quanto import freeze
from toolkit.config_modules import GenerateImageConfig, ModelConfig, NetworkConfig
from toolkit.lora_special import LoRASpecialNetwork
from toolkit.models.base_model import BaseModel
from toolkit.models.v2.vae.qwen_image import QwenImageVAE, QwenImageVAEHolderMixin
from toolkit.models.v2.text_encoders.qwen3_vl import (
Qwen3VLTextEncoder,
patch_qwen_vl_patch_embed,
)
from toolkit.basic import flush
from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
from toolkit.samplers.custom_flowmatch_sampler import (
@@ -100,30 +103,6 @@ QWEN_IMAGE_VAE_PATH = "Qwen/Qwen-Image"
HF_TOKEN = os.getenv("HF_TOKEN", None)
def patch_qwen_vl_patch_embed(model):
"""Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain
linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and
falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear
(a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves.
Returns the number of patch_embed modules patched. (Same patch as the
Qwen3VLCaptioner extension.)"""
patched = 0
for module in model.modules():
proj = getattr(module, "proj", None)
if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple(
proj.stride
):
def fast_forward(hidden_states, _proj=proj):
w = _proj.weight.reshape(_proj.weight.shape[0], -1)
x = hidden_states.view(-1, w.shape[1]).to(w.dtype)
return F.linear(x, w, _proj.bias)
module.forward = fast_forward
patched += 1
return patched
def _load_mmdit_state_dict(name_or_path: str, filename: Optional[str]) -> dict:
"""Load the MMDiT weights from a local safetensors file/dir or the HF hub.
@@ -163,7 +142,7 @@ def _load_mmdit_state_dict(name_or_path: str, filename: Optional[str]) -> dict:
return load_file(path)
class Krea2Model(BaseModel):
class Krea2Model(QwenImageVAEHolderMixin, BaseModel):
arch = "krea2"
def __init__(
@@ -268,8 +247,8 @@ class Krea2Model(BaseModel):
processor = Qwen2TokenizerFast.from_pretrained(
te_path, max_length=self.max_text_length, token=HF_TOKEN
)
text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
te_path, torch_dtype=dtype, token=HF_TOKEN
text_encoder = Qwen3VLTextEncoder.load_model(
te_path, dtype=dtype, subfolder="", token=HF_TOKEN
)
vl_processor = None
if self.is_edit:
@@ -281,8 +260,7 @@ class Krea2Model(BaseModel):
else:
# We only ever encode text, so the vision tower is dead weight -- drop it to
# free VRAM and skip loading its (bf16-slow) Conv3d patch_embed onto the GPU.
if getattr(text_encoder.model, "visual", None) is not None:
text_encoder.model.visual = None
text_encoder.drop_vision_tower()
text_encoder.eval()
text_encoder.requires_grad_(False)
flush()
@@ -291,8 +269,8 @@ class Krea2Model(BaseModel):
def _load_vae(self):
vae_path = self.model_config.model_kwargs.get("vae_path", QWEN_IMAGE_VAE_PATH)
self.print_and_status_update(f"Loading Qwen-Image VAE from {vae_path}")
vae = AutoencoderKLQwenImage.from_pretrained(
vae_path, subfolder="vae", torch_dtype=self.vae_torch_dtype, token=HF_TOKEN
vae = QwenImageVAE.load_model(
vae_path, dtype=self.vae_torch_dtype, token=HF_TOKEN
)
vae.eval()
vae.requires_grad_(False)
@@ -772,76 +750,9 @@ class Krea2Model(BaseModel):
return False
# ------------------------------------------------------------------
# VAE (Qwen-Image AutoencoderKLQwenImage -- same handling as qwen_image arch)
# VAE (Qwen-Image AutoencoderKLQwenImage -- shared QwenImageVAEHolderMixin)
# ------------------------------------------------------------------
def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None):
if device is None:
device = self.vae_device_torch
if dtype is None:
dtype = self.vae_torch_dtype
if self.vae.device == torch.device("cpu"):
self.vae.to(device)
self.vae.eval()
self.vae.requires_grad_(False)
image_list = [image.to(device, dtype=dtype) for image in image_list]
images = torch.stack(image_list).to(device, dtype=dtype)
# AutoencoderKLQwenImage is a video VAE: add a frame dim.
images = images.unsqueeze(2)
latents = self.vae.encode(images).latent_dist.sample()
latents_mean = (
torch.tensor(self.vae.config.latents_mean)
.view(1, self.vae.config.z_dim, 1, 1, 1)
.to(latents.device, latents.dtype)
)
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(
1, self.vae.config.z_dim, 1, 1, 1
).to(latents.device, latents.dtype)
latents = (latents - latents_mean) * latents_std
latents = latents.squeeze(2) # drop frame dim
return latents.to(device, dtype=dtype)
def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
if device is None:
device = self.vae_device_torch
if dtype is None:
dtype = self.vae_torch_dtype
if self.vae.device == torch.device("cpu"):
self.vae.to(device)
latents = latents.to(device, dtype=dtype)
latents = latents.unsqueeze(2) # add frame dim
latents_mean = (
torch.tensor(self.vae.config.latents_mean)
.view(1, self.vae.config.z_dim, 1, 1, 1)
.to(latents.device, latents.dtype)
)
latents_std = (
torch.tensor(self.vae.config.latents_std)
.view(1, self.vae.config.z_dim, 1, 1, 1)
.to(latents.device, latents.dtype)
)
latents = latents * latents_std + latents_mean
# Full-resolution decode spikes VRAM; tile it when low on VRAM (decode
# only -- encode stays untiled).
tiled = self.model_config.low_vram
if tiled:
self.vae.enable_tiling()
try:
images = self.vae.decode(latents).sample
finally:
if tiled:
self.vae.disable_tiling()
images = images.squeeze(2) # drop frame dim
return images.to(device, dtype=dtype)
vae_decode_tiled_on_low_vram = True
# ------------------------------------------------------------------
# Saving / bookkeeping
# ------------------------------------------------------------------

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

@@ -34,7 +34,8 @@ from torchvision.transforms.functional import to_tensor
from safetensors.torch import load_file, save_file
import huggingface_hub
from transformers import AutoProcessor, AutoTokenizer, Qwen3VLForConditionalGeneration
from transformers import AutoProcessor, AutoTokenizer
from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder
from optimum.quanto import freeze
from toolkit.config_modules import GenerateImageConfig, ModelConfig
@@ -209,8 +210,11 @@ class MageFlowModel(BaseModel):
self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}")
tokenizer = AutoTokenizer.from_pretrained(te_path, token=HF_TOKEN, **te_kwargs)
text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
te_path, torch_dtype=dtype, token=HF_TOKEN, **te_kwargs
text_encoder = Qwen3VLTextEncoder.load_model(
te_path,
dtype=dtype,
subfolder=te_kwargs.get("subfolder", ""),
token=HF_TOKEN,
)
vl_processor = None
if self.is_edit:
@@ -224,8 +228,7 @@ class MageFlowModel(BaseModel):
else:
# We only ever encode text, so the vision tower is dead weight --
# drop it to free VRAM.
if getattr(text_encoder.model, "visual", None) is not None:
text_encoder.model.visual = None
text_encoder.drop_vision_tower()
text_encoder.eval()
text_encoder.requires_grad_(False)
flush()

27
extensions_built_in/diffusion_models/mageflow/src/text_encoder.py
View File

@@ -18,7 +18,8 @@ import math
from typing import List, Optional
import torch
import torch.nn.functional as F
from toolkit.models.v2.text_encoders.qwen3_vl import patch_qwen_vl_patch_embed
# Prompt templates from the reference (mage_flow/models/utils.py). ``start_idx``
@@ -57,30 +58,6 @@ def edit_prompt_body(instruction: str, num_refs: int) -> str:
return prefix + instruction
def patch_qwen_vl_patch_embed(model):
"""Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain
linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and
falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear
(a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves.
Returns the number of patch_embed modules patched. (Same patch as the krea2
extension / Qwen3VLCaptioner.)"""
patched = 0
for module in model.modules():
proj = getattr(module, "proj", None)
if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple(
proj.stride
):
def fast_forward(hidden_states, _proj=proj):
w = _proj.weight.reshape(_proj.weight.shape[0], -1)
x = hidden_states.view(-1, w.shape[1]).to(w.dtype)
return F.linear(x, w, _proj.bias)
module.forward = fast_forward
patched += 1
return patched
def resize_vl_images(
images: List[torch.Tensor], max_long_edge: int = 384
) -> List["PIL.Image.Image"]:

89
extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py
View File

@@ -6,21 +6,25 @@ import torch
import yaml
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from toolkit.models.base_model import BaseModel
from toolkit.models.v2.vae.qwen_image import QwenImageVAE, QwenImageVAEHolderMixin
from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder
from toolkit.basic import flush
from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
from toolkit.samplers.custom_flowmatch_sampler import (
CustomFlowMatchEulerDiscreteScheduler,
)
from toolkit.accelerator import unwrap_model
from optimum.quanto import freeze
from toolkit.util.quantize import quantize, get_qtype, quantize_model
from toolkit.util.quantize import quantize_model
from toolkit.memory_management import MemoryManager
from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
from transformers import Qwen3VLProcessor
import torch.nn.functional as F
try:
from diffusers import NucleusMoEImagePipeline, NucleusMoEImageTransformer2DModel, AutoencoderKLQwenImage
from diffusers import NucleusMoEImagePipeline, AutoencoderKLQwenImage
from toolkit.models.v2.diffusion_models.nucleus_image import (
NucleusMoEImageTransformer2DModel,
)
from diffusers.models.transformers.transformer_nucleusmoe_image import SwiGLUExperts
except ImportError:
raise ImportError(
@@ -46,7 +50,7 @@ scheduler_config = {
}
class NucleusImageModel(BaseModel):
class NucleusImageModel(QwenImageVAEHolderMixin, BaseModel):
arch = "nucleus_image"
def __init__(
@@ -81,20 +85,14 @@ class NucleusImageModel(BaseModel):
self.print_and_status_update("Loading transformer")
transformer_path = model_path
transformer_subfolder = "transformer"
if os.path.exists(transformer_path):
transformer_subfolder = None
transformer_path = os.path.join(transformer_path, "transformer")
if os.path.exists(model_path):
# check if the path is a full checkpoint.
te_folder_path = os.path.join(model_path, "text_encoder")
# if we have the te, this folder is a full checkpoint, use it as the base
if os.path.exists(te_folder_path):
base_model_path = model_path
transformer = NucleusMoEImageTransformer2DModel.from_pretrained(
transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype
)
transformer = NucleusMoEImageTransformer2DModel.load_model(model_path, dtype=dtype)
# handle versions of pytorch that don't have grouped mm, by disabling it in the SwiGLUExperts
if not hasattr(torch.nn.functional, "grouped_mm"):
@@ -129,33 +127,13 @@ class NucleusImageModel(BaseModel):
tokenizer = Qwen3VLProcessor.from_pretrained(
base_model_path, subfolder="processor", torch_dtype=dtype
)
text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
base_model_path, subfolder="text_encoder", torch_dtype=dtype
)
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
text_encoder.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder")
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
text_encoder = Qwen3VLTextEncoder.load_model(base_model_path, dtype=dtype)
self.prepare_text_encoder(text_encoder, dtype=dtype)
self.print_and_status_update("Loading VAE")
vae = AutoencoderKLQwenImage.from_pretrained(
base_model_path, subfolder="vae", torch_dtype=dtype
).to(self.device_torch, dtype=dtype)
vae = QwenImageVAE.load_model(
base_model_path, dtype=dtype, device=self.device_torch
)
self.noise_scheduler = NucleusImageModel.get_train_scheduler()
@@ -214,41 +192,6 @@ class NucleusImageModel(BaseModel):
return pipeline
def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None):
if device is None:
device = self.vae_device_torch
if dtype is None:
dtype = self.vae_torch_dtype
# Move to vae to device if on cpu
if self.vae.device == torch.device("cpu"):
self.vae.to(device)
self.vae.eval()
self.vae.requires_grad_(False)
# move to device and dtype
image_list = [image.to(device, dtype=dtype) for image in image_list]
images = torch.stack(image_list).to(device, dtype=dtype)
# it uses wan vae, so add dim for frame count
images = images.unsqueeze(2)
latents = self.vae.encode(images).latent_dist.sample()
latents_mean = (
torch.tensor(self.vae.config.latents_mean)
.view(1, self.vae.config.z_dim, 1, 1, 1)
.to(latents.device, latents.dtype)
)
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(
1, self.vae.config.z_dim, 1, 1, 1
).to(latents.device, latents.dtype)
latents = (latents - latents_mean) * latents_std
latents = latents.to(device, dtype=dtype)
latents = latents.squeeze(2) # remove the frame count dimension
return latents
def generate_single_image(
self,
pipeline: NucleusMoEImagePipeline,

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

@@ -96,8 +96,8 @@ class OmniGen2Model(BaseModel):
self.print_and_status_update("Loading transformer")
transformer = OmniGen2Transformer2DModel.from_pretrained(
model_path, subfolder="transformer", torch_dtype=torch.bfloat16
transformer = OmniGen2Transformer2DModel.load_model(
model_path, dtype=torch.bfloat16
)
if not self.low_vram:

12
extensions_built_in/diffusion_models/omnigen2/src/models/transformers/transformer_omnigen2.py
View File

@@ -15,6 +15,8 @@ from diffusers.models.attention_processor import Attention
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.modeling_utils import ModelMixin
from toolkit.models.v2._mixin import OstrisModelMixin
from ..attention_processor import OmniGen2AttnProcessorFlash2Varlen, OmniGen2AttnProcessor
from .repo import OmniGen2RotaryPosEmbed
from .block_lumina2 import LuminaLayerNormContinuous, LuminaRMSNormZero, LuminaFeedForward, Lumina2CombinedTimestepCaptionEmbedding
@@ -177,7 +179,15 @@ class OmniGen2TransformerBlock(nn.Module):
return hidden_states
class OmniGen2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
class OmniGen2Transformer2DModel(
ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, OstrisModelMixin
):
aitk_subfolder = "transformer"
@classmethod
def get_transformer_block_names(cls):
return ["noise_refiner", "ref_image_refiner", "context_refiner", "layers"]
"""
OmniGen2 Transformer 2D Model.

4
extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py
View File

@@ -126,9 +126,7 @@ class PRXPixelT2IModel(BaseModel):
self.print_and_status_update("Loading transformer")
# from_pretrained reads config.json (bottleneck_size, resolution_embeds,
# in_channels=3, ...) and the safetensors in one shot.
transformer = PRXTransformer2DModel.from_pretrained(
model_path, subfolder="transformer", torch_dtype=dtype
)
transformer = PRXTransformer2DModel.load_model(model_path, dtype=dtype)
transformer.to(dtype=dtype)
flush()

10
extensions_built_in/diffusion_models/prx_pixel_t2i/src/transformer_prx.py
View File

@@ -40,6 +40,8 @@ from diffusers.models.attention_dispatch import dispatch_attention_fn
from diffusers.models.embeddings import get_timestep_embedding
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.modeling_utils import ModelMixin
from toolkit.models.v2._mixin import OstrisModelMixin
from diffusers.models.normalization import RMSNorm
@@ -666,7 +668,7 @@ def seq2img(seq: torch.Tensor, patch_size: int, shape: torch.Tensor) -> torch.Te
return seq
class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin, OstrisModelMixin):
r"""
Transformer-based 2D model for text to image generation.
@@ -702,6 +704,12 @@ class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
`PRXResolutionEmbedder`. Used by the PRX-7B variant.
"""
aitk_subfolder = "transformer"
@classmethod
def get_transformer_block_names(cls):
return ["blocks"]
config_name = "config.json"
_supports_gradient_checkpointing = True

145
extensions_built_in/diffusion_models/qwen_image/qwen_image.py
View File

@@ -7,28 +7,25 @@ from toolkit import train_tools
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from PIL import Image
from toolkit.models.base_model import BaseModel
from toolkit.models.v2.vae.qwen_image import QwenImageVAE, QwenImageVAEHolderMixin
from toolkit.basic import flush
from toolkit.prompt_utils import PromptEmbeds
from toolkit.samplers.custom_flowmatch_sampler import (
CustomFlowMatchEulerDiscreteScheduler,
)
from toolkit.accelerator import get_accelerator, unwrap_model
from optimum.quanto import freeze, QTensor
from toolkit.util.quantize import quantize, get_qtype, quantize_model
from toolkit.util.quantize import quantize_model
import torch.nn.functional as F
from toolkit.memory_management import MemoryManager
from safetensors.torch import load_file
from diffusers import (
QwenImagePipeline,
QwenImageTransformer2DModel,
AutoencoderKLQwenImage,
)
from transformers import (
Qwen2_5_VLForConditionalGeneration,
Qwen2Tokenizer,
Qwen2VLProcessor,
)
from transformers import Qwen2VLProcessor
from toolkit.models.v2.diffusion_models.qwen_image import QwenImageTransformer2DModel
from toolkit.models.v2.text_encoders.qwen25_vl import Qwen25VLTextEncoder
from tqdm import tqdm
from toolkit.util.qwen_vae_gradient_checkpointing import patch_qwen_vae_gradient_checkpointing
@@ -55,7 +52,7 @@ scheduler_config = {
}
class QwenImageModel(BaseModel):
class QwenImageModel(QwenImageVAEHolderMixin, BaseModel):
arch = "qwen_image"
_qwen_image_keep_visual = False
_qwen_pipeline = QwenImagePipeline
@@ -97,31 +94,14 @@ class QwenImageModel(BaseModel):
self.print_and_status_update("Loading transformer")
if model_path.endswith(".safetensors"):
# load the safetensors file
transformer = QwenImageTransformer2DModel.from_single_file(
model_path,
config="Qwen/Qwen-Image",
subfolder="transformer",
torch_dtype=model_dtype,
)
transformer.to(model_dtype)
if not model_path.endswith(".safetensors") and os.path.exists(model_path):
# check if the path is a full checkpoint.
te_folder_path = os.path.join(model_path, "text_encoder")
# if we have the te, this folder is a full checkpoint, use it as the base
if os.path.exists(te_folder_path):
base_model_path = model_path
else:
transformer_path = model_path
transformer_subfolder = "transformer"
if os.path.exists(transformer_path):
transformer_subfolder = None
transformer_path = os.path.join(transformer_path, "transformer")
# check if the path is a full checkpoint.
te_folder_path = os.path.join(model_path, "text_encoder")
# if we have the te, this folder is a full checkpoint, use it as the base
if os.path.exists(te_folder_path):
base_model_path = model_path
transformer = QwenImageTransformer2DModel.from_pretrained(
transformer_path, subfolder=transformer_subfolder, torch_dtype=dtype
)
transformer = QwenImageTransformer2DModel.load_model(model_path, dtype=model_dtype)
if self.model_config.quantize:
self.print_and_status_update("Quantizing Transformer")
@@ -145,41 +125,18 @@ class QwenImageModel(BaseModel):
flush()
self.print_and_status_update("Text Encoder")
tokenizer = Qwen2Tokenizer.from_pretrained(
base_model_path, subfolder="tokenizer", torch_dtype=dtype
)
text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
base_model_path, subfolder="text_encoder", torch_dtype=dtype
)
tokenizer = Qwen25VLTextEncoder.load_tokenizer(base_model_path, use_fast=False)
text_encoder = Qwen25VLTextEncoder.load_model(base_model_path, dtype=dtype)
# remove the visual model as it is not needed for image generation
self.processor = None
if not self._qwen_image_keep_visual:
text_encoder.model.visual = None
text_encoder.drop_vision_tower()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
text_encoder.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder")
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
self.prepare_text_encoder(text_encoder, dtype=dtype)
self.print_and_status_update("Loading VAE")
vae = AutoencoderKLQwenImage.from_pretrained(
base_model_path, subfolder="vae", torch_dtype=dtype
)
vae = QwenImageVAE.load_model(base_model_path, dtype=dtype)
self.noise_scheduler = QwenImageModel.get_train_scheduler()
@@ -417,69 +374,3 @@ class QwenImageModel(BaseModel):
lora_keys_use_comfy_prefix = True
def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None):
if device is None:
device = self.vae_device_torch
if dtype is None:
dtype = self.vae_torch_dtype
# Move to vae to device if on cpu
if self.vae.device == torch.device("cpu"):
self.vae.to(device)
self.vae.eval()
self.vae.requires_grad_(False)
# move to device and dtype
image_list = [image.to(device, dtype=dtype) for image in image_list]
images = torch.stack(image_list).to(device, dtype=dtype)
# it uses wan vae, so add dim for frame count
images = images.unsqueeze(2)
latents = self.vae.encode(images).latent_dist.sample()
latents_mean = (
torch.tensor(self.vae.config.latents_mean)
.view(1, self.vae.config.z_dim, 1, 1, 1)
.to(latents.device, latents.dtype)
)
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(
1, self.vae.config.z_dim, 1, 1, 1
).to(latents.device, latents.dtype)
latents = (latents - latents_mean) * latents_std
latents = latents.to(device, dtype=dtype)
latents = latents.squeeze(2) # remove the frame count dimension
return latents
def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
if device is None:
device = self.vae_device_torch
if dtype is None:
dtype = self.vae_torch_dtype
if self.vae.device == torch.device("cpu"):
self.vae.to(device)
latents = latents.to(device, dtype=dtype)
# add frame count dim for wan vae
latents = latents.unsqueeze(2)
latents_mean = (
torch.tensor(self.vae.config.latents_mean)
.view(1, self.vae.config.z_dim, 1, 1, 1)
.to(latents.device, latents.dtype)
)
latents_std = (
torch.tensor(self.vae.config.latents_std)
.view(1, self.vae.config.z_dim, 1, 1, 1)
.to(latents.device, latents.dtype)
)
latents = latents * latents_std + latents_mean
images = self.vae.decode(latents).sample
images = images.squeeze(2) # remove the frame count dimension
return images.to(device, dtype=dtype)

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

@@ -17,7 +17,7 @@ from toolkit.samplers.custom_flowmatch_sampler import (
)
from toolkit.util.quantize import quantize_model
from .wan22_pipeline import Wan22Pipeline
from diffusers import WanTransformer3DModel
from toolkit.models.v2.diffusion_models.wan import WanTransformer3DModel
from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
from torchvision.transforms import functional as TF
@@ -287,11 +287,9 @@ class Wan2214bModel(Wan21):
self.print_and_status_update("Loading transformer 1")
dtype = self.torch_dtype
transformer_1 = WanTransformer3DModel.from_pretrained(
transformer_path_1,
subfolder=subfolder_1,
torch_dtype=dtype,
).to(dtype=dtype)
transformer_1 = WanTransformer3DModel.load_model(
transformer_path_1, dtype=dtype, subfolder=subfolder_1
)
flush()
@@ -317,11 +315,9 @@ class Wan2214bModel(Wan21):
self.print_and_status_update("Loading transformer 2")
dtype = self.torch_dtype
transformer_2 = WanTransformer3DModel.from_pretrained(
transformer_path_2,
subfolder=subfolder_2,
torch_dtype=dtype,
).to(dtype=dtype)
transformer_2 = WanTransformer3DModel.load_model(
transformer_path_2, dtype=dtype, subfolder=subfolder_2
)
flush()

39
extensions_built_in/diffusion_models/z_image/z_image.py
View File

@@ -13,19 +13,16 @@ from toolkit.samplers.custom_flowmatch_sampler import (
CustomFlowMatchEulerDiscreteScheduler,
)
from toolkit.accelerator import unwrap_model
from optimum.quanto import freeze
from toolkit.util.quantize import (
quantize,
get_qtype,
quantize_model,
dequantize_if_quantized,
)
from toolkit.memory_management import MemoryManager
from toolkit.metadata import get_meta_for_safetensors
from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder
from toolkit.models.v2.vae.autoencoder_kl import KLVAE
from safetensors.torch import load_file, save_file
from transformers import AutoTokenizer, Qwen3ForCausalLM
from diffusers import AutoencoderKL
try:
from diffusers import ZImagePipeline
@@ -255,36 +252,12 @@ class ZImageModel(BaseModel):
flush()
self.print_and_status_update("Text Encoder")
tokenizer = AutoTokenizer.from_pretrained(
base_model_path, subfolder="tokenizer", torch_dtype=dtype
)
text_encoder = Qwen3ForCausalLM.from_pretrained(
base_model_path, subfolder="text_encoder", torch_dtype=dtype
)
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
text_encoder.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder")
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path)
text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype)
self.prepare_text_encoder(text_encoder, dtype=dtype)
self.print_and_status_update("Loading VAE")
vae = AutoencoderKL.from_pretrained(
base_model_path, subfolder="vae", torch_dtype=dtype
)
vae = KLVAE.load_model(base_model_path, dtype=dtype)
self.noise_scheduler = ZImageModel.get_train_scheduler()

33
extensions_built_in/diffusion_models/z_image/z_image_l2p_model.py
View File

@@ -9,11 +9,10 @@ import torch.nn.functional as F
import yaml
from toolkit.basic import flush
from toolkit.accelerator import unwrap_model
from optimum.quanto import freeze
from toolkit.util.quantize import quantize, get_qtype, quantize_model
from toolkit.util.quantize import quantize_model
from toolkit.memory_management import MemoryManager
from transformers import AutoTokenizer, Qwen3ForCausalLM
from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder
from toolkit.models.FakeVAE import FakeVAE
from toolkit.paths import MODELS_PATH
from safetensors.torch import load_file, save_file
@@ -468,31 +467,9 @@ class ZImageL2PModel(ZImageModel):
flush()
self.print_and_status_update("Text Encoder")
tokenizer = AutoTokenizer.from_pretrained(
base_model_path, subfolder="tokenizer", torch_dtype=dtype
)
text_encoder = Qwen3ForCausalLM.from_pretrained(
base_model_path, subfolder="text_encoder", torch_dtype=dtype
)
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
text_encoder.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder")
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path)
text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype)
self.prepare_text_encoder(text_encoder, dtype=dtype)
self.print_and_status_update("Loading VAE")
# vae = AutoencoderKL.from_pretrained(

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

@@ -13,14 +13,13 @@ from toolkit.samplers.custom_flowmatch_sampler import (
CustomFlowMatchEulerDiscreteScheduler,
)
from toolkit.accelerator import unwrap_model
from optimum.quanto import freeze
from toolkit.util.quantize import quantize, get_qtype, quantize_model
from toolkit.util.quantize import quantize_model
from toolkit.memory_management import MemoryManager
from safetensors.torch import load_file
from optimum.quanto import QTensor
from toolkit.metadata import get_meta_for_safetensors
from safetensors.torch import load_file, save_file
from transformers import AutoTokenizer, Qwen3ForCausalLM
from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder
from diffusers import AutoencoderKL
from toolkit.models.FakeVAE import FakeVAE
from .zeta_chroma_transformer import ZImageDCT, ZImageDCTParams, vae_flatten, vae_unflatten, prepare_latent_image_ids, make_text_position_ids
@@ -142,31 +141,9 @@ class ZetaChromaModel(BaseModel):
flush()
self.print_and_status_update("Text Encoder")
tokenizer = AutoTokenizer.from_pretrained(
base_model_path, subfolder="tokenizer", torch_dtype=dtype
)
text_encoder = Qwen3ForCausalLM.from_pretrained(
base_model_path, subfolder="text_encoder", torch_dtype=dtype
)
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
text_encoder.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder")
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path)
text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype)
self.prepare_text_encoder(text_encoder, dtype=dtype)
self.print_and_status_update("Loading VAE")
vae = FakeVAE(scaling_factor=1.0)