407 lines
16 KiB
Python
407 lines
16 KiB
Python
"""Boogu-Image base (text-to-image) integration for ai-toolkit.
|
|
|
|
Boogu-Image is a Lumina2-style mixed double-/single-stream flow-matching DiT
|
|
conditioned on Qwen3-VL instruction features. This wires up the base T2I model
|
|
for LoRA / fine-tune training and preview sampling.
|
|
|
|
Only the base text-to-image path is implemented here (no reference-image / edit
|
|
conditioning). The architecture lives under ``./src`` (vendored & trimmed from the
|
|
upstream Boogu repo); nothing is imported from the original repo.
|
|
|
|
Weights are pulled from the bf16 release ``Boogu/Boogu-Image-0.1-Base`` (clean
|
|
safetensors). The ``-fp8`` sibling ships torchao float8 ``.bin`` weights that
|
|
need a matching torchao/cache_dit to deserialize and is not supported here --
|
|
use the bf16 repo and set ``quantize: true`` to run the transformer in fp8 via
|
|
ai-toolkit's own quantization.
|
|
"""
|
|
|
|
import os
|
|
from typing import List, Optional
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import yaml
|
|
from safetensors.torch import save_file
|
|
|
|
from transformers import AutoModel, AutoProcessor
|
|
|
|
from toolkit.accelerator import unwrap_model
|
|
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.models.v2.text_encoders.qwen3_vl import Qwen3VLModelEncoder
|
|
from toolkit.models.v2.vae.autoencoder_kl import KLVAE
|
|
from toolkit.samplers.custom_flowmatch_sampler import (
|
|
CustomFlowMatchEulerDiscreteScheduler,
|
|
)
|
|
|
|
|
|
from optimum.quanto import QTensor
|
|
from diffusers import AutoencoderKL
|
|
|
|
from .src.transformer import BooguImageTransformer2DModel
|
|
from .src.rope import get_freqs_cis
|
|
from .src.pipeline import (
|
|
BooguImagePipeline,
|
|
pad_instruction_features,
|
|
run_boogu_transformer,
|
|
)
|
|
|
|
|
|
# ai-toolkit uses CustomFlowMatchEulerDiscreteScheduler for training and (via our
|
|
# pipeline) sampling. ``shift`` warps timesteps toward the high-noise end; 3.0 is a
|
|
# reasonable high-resolution default and Boogu's own time-shift is applied in the
|
|
# preview sampler (see src/pipeline.boogu_time_schedule).
|
|
scheduler_config = {
|
|
"num_train_timesteps": 1000,
|
|
"use_dynamic_shifting": False,
|
|
"shift": 3.0,
|
|
}
|
|
|
|
# Released weights. The "-fp8" sibling ships torchao float8 weights that need
|
|
# cache_dit/torchao to deserialize; the plain repo ships clean bf16 safetensors,
|
|
# which load directly and let ai-toolkit do its own (optional) quantization.
|
|
BOOGU_BASE_PATH = "Boogu/Boogu-Image-0.1-Base"
|
|
|
|
# System prompt the base T2I model was trained with (SYSTEM_PROMPT_4_T2I upstream).
|
|
SYSTEM_PROMPT_T2I = (
|
|
"You are a helpful assistant that generates high-quality images based on user "
|
|
"instructions. The instructions are as follows."
|
|
)
|
|
|
|
HF_TOKEN = os.getenv("HF_TOKEN", None)
|
|
|
|
|
|
class BooguImageModel(BaseModel):
|
|
arch = "boogu_image"
|
|
# Default HF repo when model.name_or_path is unset (overridden by the edit model).
|
|
default_repo = BOOGU_BASE_PATH
|
|
use_old_lokr_format = False
|
|
|
|
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 = ["BooguImageTransformer2DModel"]
|
|
|
|
self.patch_size = 2
|
|
self.vae_scale_factor = 8
|
|
# Safety cap on instruction token length (truncation only). Each caption is
|
|
# encoded at its natural length and padded to the batch max at the model
|
|
# call, so this is just an upper bound.
|
|
self.max_text_length = int(
|
|
self.model_config.model_kwargs.get("max_text_length", 1024)
|
|
)
|
|
|
|
# Lazily-built, resolution-independent rotary frequency tables.
|
|
self._freqs_cis = None
|
|
|
|
@property
|
|
def text_embedding_space_version(self):
|
|
return self.arch + "_v1"
|
|
|
|
@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
|
|
|
|
def get_freqs_cis(self):
|
|
"""Precompute (once) the per-axis rotary frequency tables for the model."""
|
|
if self._freqs_cis is None:
|
|
cfg = unwrap_model(self.model).config
|
|
self._freqs_cis = get_freqs_cis(
|
|
cfg.axes_dim_rope, cfg.axes_lens, theta=10000
|
|
)
|
|
return self._freqs_cis
|
|
|
|
# ------------------------------------------------------------------
|
|
# Loading
|
|
# ------------------------------------------------------------------
|
|
def load_model(self):
|
|
dtype = self.torch_dtype
|
|
self.print_and_status_update("Loading Boogu-Image model")
|
|
base = self.model_config.name_or_path or self.default_repo
|
|
|
|
# --- transformer ---
|
|
# Loads the bf16 release (clean safetensors). The "-fp8" sibling ships
|
|
# torchao float8 .bin weights that need a matching torchao/cache_dit to
|
|
# deserialize -- use the bf16 repo and let ai-toolkit quantize if wanted.
|
|
self.print_and_status_update("Loading transformer")
|
|
try:
|
|
transformer = BooguImageTransformer2DModel.load_model(
|
|
base, dtype=dtype, token=HF_TOKEN
|
|
)
|
|
except OSError as e:
|
|
raise OSError(
|
|
f"Could not load Boogu transformer safetensors from '{base}'. The "
|
|
f"'-fp8' release ships torchao float8 .bin weights, which are not "
|
|
f"supported here -- point model.name_or_path at the bf16 repo "
|
|
f"'{BOOGU_BASE_PATH}' instead."
|
|
) from e
|
|
transformer.eval()
|
|
flush()
|
|
|
|
# Attention defaults to torch SDPA ("native"); opt into Flash Attention 2
|
|
# with model_kwargs.attention_backend: "flash" (needs the flash_attn pkg).
|
|
attention_backend = self.model_config.model_kwargs.get(
|
|
"attention_backend", "native"
|
|
)
|
|
if attention_backend != "native":
|
|
transformer.set_attention_backend(attention_backend)
|
|
|
|
# quantize + offload + placement, all driven by model_config
|
|
transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
|
|
flush()
|
|
|
|
# --- instruction encoder (Qwen3-VL) + processor ---
|
|
te_path = self.model_config.model_kwargs.get("text_encoder_path", base)
|
|
te_subfolder = self.model_config.model_kwargs.get(
|
|
"text_encoder_subfolder", "mllm"
|
|
)
|
|
self.print_and_status_update("Loading Qwen3-VL instruction encoder")
|
|
processor = AutoProcessor.from_pretrained(
|
|
te_path, subfolder="processor", token=HF_TOKEN
|
|
)
|
|
# AutoModel yields the inner Qwen3VLModel (the ``.model`` of the
|
|
# *ForConditionalGeneration), whose last_hidden_state is exactly the
|
|
# instruction feature the Boogu pipeline consumes.
|
|
text_encoder = Qwen3VLModelEncoder.load_model(
|
|
te_path, dtype=dtype, subfolder=te_subfolder, token=HF_TOKEN
|
|
)
|
|
text_encoder.eval()
|
|
text_encoder.requires_grad_(False)
|
|
# The vision tower's bf16 Conv3d patch_embed has no fast kernel and stalls
|
|
# image caching for the edit model -- swap it for an equivalent F.linear.
|
|
# No-op for the base T2I model (it never runs the vision tower).
|
|
n_patched = patch_qwen_vl_patch_embed(text_encoder)
|
|
if n_patched:
|
|
self.print_and_status_update(
|
|
f" - patched {n_patched} Qwen-VL Conv3d patch_embed -> linear"
|
|
)
|
|
flush()
|
|
|
|
# quantize + offload + placement, all driven by model_config
|
|
text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
|
|
flush()
|
|
|
|
# --- VAE (FLUX AutoencoderKL) ---
|
|
self.print_and_status_update("Loading VAE")
|
|
vae = KLVAE.load_model(base, dtype=self.vae_torch_dtype, token=HF_TOKEN)
|
|
vae.to(self.vae_device_torch, dtype=self.vae_torch_dtype)
|
|
vae.eval()
|
|
vae.requires_grad_(False)
|
|
flush()
|
|
|
|
self.noise_scheduler = BooguImageModel.get_train_scheduler()
|
|
self.vae = vae
|
|
self.text_encoder = text_encoder
|
|
self.tokenizer = processor
|
|
self.model = transformer
|
|
self.pipeline = BooguImagePipeline(self)
|
|
self.print_and_status_update("Model Loaded")
|
|
|
|
# ------------------------------------------------------------------
|
|
# Generation
|
|
# ------------------------------------------------------------------
|
|
def get_generation_pipeline(self):
|
|
return BooguImagePipeline(self)
|
|
|
|
def generate_single_image(
|
|
self,
|
|
pipeline: BooguImagePipeline,
|
|
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 (1000 = pure noise)
|
|
text_embeddings: AdvancedPromptEmbeds,
|
|
**kwargs,
|
|
):
|
|
if self.model.device == torch.device("cpu"):
|
|
self.model.to(self.device_torch)
|
|
|
|
# toolkit timestep (0..1000, 1000=noise) -> Boogu native time (0=noise, 1=clean)
|
|
t01 = timestep.to(self.device_torch, dtype=torch.float32) / 1000.0
|
|
if t01.dim() == 0:
|
|
t01 = t01.unsqueeze(0)
|
|
if t01.shape[0] != latent_model_input.shape[0]:
|
|
t01 = t01.expand(latent_model_input.shape[0])
|
|
boogu_t = 1.0 - t01
|
|
|
|
instr_feats, instr_mask = pad_instruction_features(
|
|
text_embeddings.text_embeds, self.device_torch, self.torch_dtype
|
|
)
|
|
|
|
# Model predicts clean - noise; negate to return the toolkit velocity
|
|
# (noise - clean), matching get_loss_target / the scheduler.
|
|
raw_velocity = run_boogu_transformer(
|
|
self.transformer,
|
|
latent_model_input.to(self.device_torch, self.torch_dtype),
|
|
boogu_t,
|
|
instr_feats,
|
|
instr_mask,
|
|
self.get_freqs_cis(),
|
|
)
|
|
return -raw_velocity
|
|
|
|
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)
|
|
device = self.text_encoder.device
|
|
|
|
# Encode each instruction at its natural length (no cross-sample padding);
|
|
# padding to a common length is deferred to the model call. The system
|
|
# prompt + chat template match the base T2I training setup.
|
|
features_list = []
|
|
for p in prompt:
|
|
messages = [
|
|
{
|
|
"role": "system",
|
|
"content": [{"type": "text", "text": SYSTEM_PROMPT_T2I}],
|
|
},
|
|
{"role": "user", "content": [{"type": "text", "text": p}]},
|
|
]
|
|
inputs = self.tokenizer.apply_chat_template(
|
|
[messages],
|
|
tokenize=True,
|
|
return_dict=True,
|
|
return_tensors="pt",
|
|
add_generation_prompt=False,
|
|
truncation=True,
|
|
max_length=self.max_text_length,
|
|
)
|
|
input_ids = inputs["input_ids"].to(device)
|
|
attention_mask = inputs["attention_mask"].to(device)
|
|
|
|
with torch.no_grad():
|
|
output = self.text_encoder(
|
|
input_ids=input_ids, attention_mask=attention_mask
|
|
)
|
|
# (L, D) -- drop the batch dim, one tensor per prompt
|
|
features_list.append(output.last_hidden_state[0].to(self.torch_dtype))
|
|
|
|
return AdvancedPromptEmbeds(text_embeds=features_list)
|
|
|
|
def get_loss_target(self, *args, **kwargs):
|
|
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
|
|
# ------------------------------------------------------------------
|
|
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(self.vae_device_torch)
|
|
|
|
if isinstance(image_list, list):
|
|
images = torch.stack(image_list, dim=0)
|
|
else:
|
|
images = image_list
|
|
images = images.to(device, dtype=dtype)
|
|
|
|
latents = self.vae.encode(images).latent_dist.sample()
|
|
shift = self.vae.config["shift_factor"] or 0
|
|
latents = (latents - shift) * self.vae.config["scaling_factor"]
|
|
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(self.vae_device_torch)
|
|
|
|
latents = latents.to(device, dtype=dtype)
|
|
shift = self.vae.config["shift_factor"] or 0
|
|
latents = latents / self.vae.config["scaling_factor"] + shift
|
|
return self.vae.decode(latents).sample
|
|
|
|
# ------------------------------------------------------------------
|
|
# Saving / misc
|
|
# ------------------------------------------------------------------
|
|
def save_model(self, output_path, meta, save_dtype):
|
|
transformer: BooguImageTransformer2DModel = unwrap_model(self.model)
|
|
transformer_dir = os.path.join(output_path, "transformer")
|
|
os.makedirs(transformer_dir, exist_ok=True)
|
|
|
|
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)
|
|
save_file(
|
|
save_dict,
|
|
os.path.join(transformer_dir, "diffusion_pytorch_model.safetensors"),
|
|
)
|
|
# config.json so the saved transformer can be reloaded with from_pretrained.
|
|
transformer.save_config(transformer_dir)
|
|
with open(os.path.join(output_path, "aitk_meta.yaml"), "w") as f:
|
|
yaml.dump(meta, f)
|
|
|
|
def get_base_model_version(self):
|
|
return "boogu_image.0.1"
|
|
|
|
def get_transformer_block_names(self) -> Optional[List[str]]:
|
|
return ["double_stream_layers", "single_stream_layers"]
|
|
|
|
lora_keys_use_comfy_prefix = True
|
|
|