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

Add support for Krea2 (#906)

Browse Source 
* Add support for krea2

* Update repo pointer to actual repo
This commit is contained in:
Jaret Burkett (Ostris)
2026-06-23 09:18:14 -06:00
committed by GitHub
parent af594061ab
commit 99be3d96a2
10 changed files with 1307 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

474
extensions_built_in/diffusion_models/krea2/krea2.py Normal file
View File

@@ -0,0 +1,474 @@
"""Krea 2 (K2) for ai-toolkit.
Krea 2 is a single-stream MMDiT text-to-image model:
- text encoder: Qwen3-VL-4B-Instruct (a stack of 12 hidden-state layers is fed
in; ``src/text_encoder.py``),
- autoencoder: the Qwen-Image VAE (f8, 16 latent channels, the same VAE the
``qwen_image`` arch uses),
- denoiser: ``SingleStreamDiT`` (``src/mmdit.py``), which fuses the text layers
with a small ``TextFusionTransformer`` and runs the packed [text | image]
sequence through ``SingleStreamBlock`` layers.
Flow-matching convention matches ai-toolkit exactly (t=1 noise -> t=0 clean,
target = noise - clean), so ``get_noise_prediction`` does no time flip / negation.
"""
import os
from typing import List, Optional
import torch
from safetensors.torch import load_file, save_file
import huggingface_hub
from huggingface_hub.errors import EntryNotFoundError
from diffusers import AutoencoderKLQwenImage
from transformers import (
AutoTokenizer,
Qwen2TokenizerFast,
Qwen3VLForConditionalGeneration,
)
from optimum.quanto import freeze, QTensor
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from toolkit.models.base_model import BaseModel
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 toolkit.metadata import get_meta_for_safetensors
from toolkit.util.quantize import quantize, get_qtype, quantize_model
from .src.mmdit import SingleStreamDiT, SingleMMDiTConfig
from .src.text_encoder import encode_krea_prompt, SELECT_LAYERS
from .src.pipeline import Krea2Pipeline, pad_text_features, predict_velocity
# The reference "single_mmdit_large_wide" architecture (oss_raw / oss_turbo share it).
KREA2_MMDIT_CONFIG = dict(
features=6144,
tdim=256,
txtdim=2560,
heads=48,
kvheads=12,
multiplier=4,
layers=28,
patch=2,
channels=16,
txtheads=20,
txtkvheads=20,
txtlayers=12,
)
# Krea 2's mu schedule is exponential time-shifting whose mu is linearly
# interpolated in image-token count between (256-res -> 0.5) and (1280-res ->
# 1.15) -- exactly what CustomFlowMatchEulerDiscreteScheduler's dynamic shifting
# does, so we mirror those endpoints here for the training timestep distribution.
# x1 = (256 // (8*2))**2 = 256
# x2 = (1280 // (8*2))**2 = 6400
scheduler_config = {
"base_image_seq_len": 256,
"max_image_seq_len": 6400,
"base_shift": 0.5,
"max_shift": 1.15,
"num_train_timesteps": 1000,
"shift": 1.0,
"use_dynamic_shifting": True,
"time_shift_type": "exponential",
}
# Defaults; both overridable via model.model_kwargs.
QWEN3_VL_PATH = "Qwen/Qwen3-VL-4B-Instruct"
QWEN_IMAGE_VAE_PATH = "Qwen/Qwen-Image"
HF_TOKEN = os.getenv("HF_TOKEN", None)
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.
``name_or_path`` may be: a ``.safetensors`` file, a directory containing one
(``filename`` or the lone ``.safetensors`` in it), or a hub repo id (the
file ``filename`` is downloaded, defaulting to ``model.safetensors``).
"""
if name_or_path.endswith(".safetensors") and os.path.isfile(name_or_path):
return load_file(name_or_path)
if os.path.isdir(name_or_path):
if filename is not None:
return load_file(os.path.join(name_or_path, filename))
candidates = [f for f in os.listdir(name_or_path) if f.endswith(".safetensors")]
if len(candidates) == 1:
return load_file(os.path.join(name_or_path, candidates[0]))
raise FileNotFoundError(
f"Could not pick an MMDiT checkpoint in {name_or_path}: found "
f"{candidates}. Set model.model_kwargs.checkpoint_filename."
)
# Treat as a hub repo id. When no filename is given, derive it from the repo
# name's trailing segment (e.g. "krea/Krea-2-Raw" -> "raw.safetensors",
# "krea/Krea-2-Turbo" -> "turbo.safetensors").
fname = filename or (name_or_path.split("/")[-1].split("-")[-1].lower() + ".safetensors")
try:
path = huggingface_hub.hf_hub_download(
repo_id=name_or_path, filename=fname, token=HF_TOKEN
)
except EntryNotFoundError as e:
raise FileNotFoundError(
f"Could not find {fname!r} in hub repo {name_or_path!r}. Set "
"model.model_kwargs.checkpoint_filename to the weight file name."
) from e
return load_file(path)
class Krea2Model(BaseModel):
arch = "krea2"
def __init__(
self,
device,
model_config: ModelConfig,
dtype="bf16",
custom_pipeline=None,
noise_scheduler=None,
**kwargs,
):
super().__init__(
device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
)
self.is_flow_matching = True
self.is_transformer = True
self.target_lora_modules = ["SingleStreamDiT"]
self.patch_size = KREA2_MMDIT_CONFIG["patch"]
self.vae_scale_factor = 8 # Qwen-Image VAE is f8
# Safety cap on prompt token length (truncation only); embeds are stored
# per-sample at natural length and padded to the batch max at the model call.
self.max_text_length = int(
self.model_config.model_kwargs.get("max_text_length", 512)
)
# Qwen2TokenizerFast used to tokenize the assistant suffix (matches the
# reference's separate processor pass).
self.processor = None
@staticmethod
def get_train_scheduler():
return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
def get_bucket_divisibility(self):
# 8 for the VAE downsample, 2 for the patch size.
return self.vae_scale_factor * self.patch_size
# ------------------------------------------------------------------
# Loading
# ------------------------------------------------------------------
def _load_transformer(self):
dtype = self.torch_dtype
self.print_and_status_update("Loading transformer (SingleStreamDiT)")
mmdit_kwargs = dict(KREA2_MMDIT_CONFIG)
mmdit_kwargs.update(self.model_config.model_kwargs.get("mmdit_config", {}))
config = SingleMMDiTConfig(**mmdit_kwargs)
# Build on meta, then materialize straight from the checkpoint.
with torch.device("meta"):
transformer = SingleStreamDiT(config)
self.print_and_status_update(" - fetching transformer weights")
state_dict = _load_mmdit_state_dict(
self.model_config.name_or_path,
self.model_config.model_kwargs.get("checkpoint_filename", None),
)
state_dict = {
k: (v.to(dtype) if v.is_floating_point() else v)
for k, v in state_dict.items()
}
self.print_and_status_update(" - loading transformer state dict")
transformer.load_state_dict(state_dict, strict=True, assign=True)
del state_dict
flush()
return transformer
def _load_text_encoder(self):
dtype = self.torch_dtype
te_path = self.model_config.model_kwargs.get("text_encoder_path", QWEN3_VL_PATH)
self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}")
tokenizer = AutoTokenizer.from_pretrained(
te_path, max_length=self.max_text_length, token=HF_TOKEN
)
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
)
# 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.eval()
text_encoder.requires_grad_(False)
flush()
return tokenizer, processor, text_encoder
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.eval()
vae.requires_grad_(False)
return vae
def load_model(self):
dtype = self.torch_dtype
self.print_and_status_update("Loading Krea 2 model")
transformer = self._load_transformer()
if self.model_config.quantize:
self.print_and_status_update("Quantizing transformer")
quantize_model(self, transformer)
flush()
if self.model_config.low_vram:
self.print_and_status_update("Moving transformer to CPU")
transformer.to("cpu")
else:
transformer.to(self.device_torch, dtype=dtype)
flush()
tokenizer, processor, text_encoder = self._load_text_encoder()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing text encoder")
text_encoder.to(self.device_torch)
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
if self.model_config.low_vram:
self.print_and_status_update("Moving text encoder to CPU")
text_encoder.to("cpu")
else:
text_encoder.to(self.device_torch)
flush()
vae = self._load_vae()
vae.to(self.vae_device_torch, dtype=self.vae_torch_dtype)
self.noise_scheduler = Krea2Model.get_train_scheduler()
self.vae = vae
self.text_encoder = text_encoder
self.tokenizer = tokenizer
self.processor = processor
self.model = transformer
self.pipeline = Krea2Pipeline(self)
self.print_and_status_update("Model Loaded")
# ------------------------------------------------------------------
# Generation (training previews)
# ------------------------------------------------------------------
def get_generation_pipeline(self):
return Krea2Pipeline(self)
def generate_single_image(
self,
pipeline: Krea2Pipeline,
gen_config: GenerateImageConfig,
conditional_embeds: AdvancedPromptEmbeds,
unconditional_embeds: AdvancedPromptEmbeds,
generator: torch.Generator,
extra: dict,
):
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
sc = self.get_bucket_divisibility()
gen_config.width = int(gen_config.width // sc * sc)
gen_config.height = int(gen_config.height // sc * sc)
img = pipeline(
conditional_embeds=conditional_embeds,
unconditional_embeds=unconditional_embeds,
height=gen_config.height,
width=gen_config.width,
num_inference_steps=gen_config.num_inference_steps,
guidance_scale=gen_config.guidance_scale,
latents=gen_config.latents,
generator=generator,
)[0]
return img
# ------------------------------------------------------------------
# Training hooks
# ------------------------------------------------------------------
def get_noise_prediction(
self,
latent_model_input: torch.Tensor, # (B, 16, h, w)
timestep: torch.Tensor, # 0..1000 scale
text_embeddings: AdvancedPromptEmbeds,
**kwargs,
):
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
# toolkit timestep (0..1000, 1000 = pure noise) -> Krea flow time t in
# [0, 1] with t=1 = pure noise. Same convention -> straight divide.
t = timestep.to(self.device_torch, dtype=torch.float32) / 1000.0
if t.dim() == 0:
t = t.unsqueeze(0)
if t.shape[0] != latent_model_input.shape[0]:
t = t.expand(latent_model_input.shape[0])
context, text_mask = pad_text_features(
text_embeddings.text_embeds, self.device_torch, self.torch_dtype
)
pred = predict_velocity(
self.transformer,
latent_model_input.to(self.device_torch, self.torch_dtype),
t,
context,
text_mask,
)
return pred
def get_prompt_embeds(self, prompt) -> AdvancedPromptEmbeds:
if isinstance(prompt, str):
prompt = [prompt]
if self.text_encoder.device == torch.device("cpu"):
self.text_encoder.to(self.device_torch)
# Encode each prompt at its natural length and store one (L, 12*2560)
# tensor per batch item. The (L, 12, 2560) stack is flattened to 2D so the
# toolkit's batching reads the list length (not the seq length) as the
# batch size; predict_velocity restores the layer axis. Padding to the
# batch max is deferred to the model call so caches stay small and any
# prompts can share a batch.
features_list = []
for p in prompt:
features = encode_krea_prompt(
self.text_encoder,
self.tokenizer,
self.processor,
p,
max_length=self.max_text_length,
select_layers=SELECT_LAYERS,
)
# (L, n, d) -> (L, n*d)
features = features.reshape(features.shape[0], -1)
features_list.append(features.to(self.torch_dtype))
return AdvancedPromptEmbeds(text_embeds=features_list)
def get_loss_target(self, *args, **kwargs):
# Flow-matching velocity target: noise - clean.
noise = kwargs.get("noise")
batch = kwargs.get("batch")
return (noise - batch.latents).detach()
def get_model_has_grad(self):
return False
def get_te_has_grad(self):
return False
# ------------------------------------------------------------------
# VAE (Qwen-Image AutoencoderKLQwenImage -- same handling as qwen_image arch)
# ------------------------------------------------------------------
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
images = self.vae.decode(latents).sample
images = images.squeeze(2) # drop frame dim
return images.to(device, dtype=dtype)
# ------------------------------------------------------------------
# Saving / bookkeeping
# ------------------------------------------------------------------
def save_model(self, output_path, meta, save_dtype):
if not output_path.endswith(".safetensors"):
output_path = output_path + ".safetensors"
transformer: SingleStreamDiT = unwrap_model(self.model)
state_dict = transformer.state_dict()
save_dict = {}
for k, v in state_dict.items():
if isinstance(v, QTensor):
v = v.dequantize()
save_dict[k] = v.clone().to("cpu", dtype=save_dtype)
meta = get_meta_for_safetensors(meta, name="krea2")
save_file(save_dict, output_path, metadata=meta)
def get_base_model_version(self):
return "krea2"
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["blocks"]
def convert_lora_weights_before_save(self, state_dict):
return {
k.replace("transformer.", "diffusion_model."): v
for k, v in state_dict.items()
}
def convert_lora_weights_before_load(self, state_dict):
return {
k.replace("diffusion_model.", "transformer."): v
for k, v in state_dict.items()
}