diff --git a/extensions_built_in/diffusion_models/__init__.py b/extensions_built_in/diffusion_models/__init__.py index a410b1e..24bcde7 100644 --- a/extensions_built_in/diffusion_models/__init__.py +++ b/extensions_built_in/diffusion_models/__init__.py @@ -15,6 +15,7 @@ from .hidream.hidream_o1_model import HidreamO1Model from .z_image.z_image_l2p_model import ZImageL2PModel from .ideogram4 import Ideogram4Model from .prx_pixel_t2i import PRXPixelT2IModel +from .boogu_image import BooguImageModel, BooguImageEditModel AI_TOOLKIT_MODELS = [ # put a list of models here @@ -44,4 +45,6 @@ AI_TOOLKIT_MODELS = [ ZImageL2PModel, Ideogram4Model, PRXPixelT2IModel, + BooguImageModel, + BooguImageEditModel, ] diff --git a/extensions_built_in/diffusion_models/boogu_image/__init__.py b/extensions_built_in/diffusion_models/boogu_image/__init__.py new file mode 100644 index 0000000..e966e37 --- /dev/null +++ b/extensions_built_in/diffusion_models/boogu_image/__init__.py @@ -0,0 +1,4 @@ +from .boogu_image import BooguImageModel +from .boogu_image_edit import BooguImageEditModel + +__all__ = ["BooguImageModel", "BooguImageEditModel"] diff --git a/extensions_built_in/diffusion_models/boogu_image/boogu_image.py b/extensions_built_in/diffusion_models/boogu_image/boogu_image.py new file mode 100644 index 0000000..fe5b75b --- /dev/null +++ b/extensions_built_in/diffusion_models/boogu_image/boogu_image.py @@ -0,0 +1,459 @@ +"""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.samplers.custom_flowmatch_sampler import ( + CustomFlowMatchEulerDiscreteScheduler, +) +from toolkit.util.quantize import quantize, get_qtype, quantize_model + +from optimum.quanto import freeze, 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) + + +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). + 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.from_pretrained( + base, subfolder="transformer", torch_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) + + 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() + + # --- 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 = AutoModel.from_pretrained( + te_path, subfolder=te_subfolder, torch_dtype=dtype, 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() + + if self.model_config.quantize_te: + self.print_and_status_update("Quantizing instruction 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 instruction encoder to CPU") + text_encoder.to("cpu") + else: + text_encoder.to(self.device_torch) + flush() + + # --- VAE (FLUX AutoencoderKL) --- + self.print_and_status_update("Loading VAE") + vae = AutoencoderKL.from_pretrained( + base, subfolder="vae", torch_dtype=self.vae_torch_dtype, token=HF_TOKEN + ) + vae.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"] + + 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() + } diff --git a/extensions_built_in/diffusion_models/boogu_image/boogu_image_edit.py b/extensions_built_in/diffusion_models/boogu_image/boogu_image_edit.py new file mode 100644 index 0000000..cd91035 --- /dev/null +++ b/extensions_built_in/diffusion_models/boogu_image/boogu_image_edit.py @@ -0,0 +1,386 @@ +"""Boogu-Image edit (TI2I) integration for ai-toolkit. + +The edit model is the same Lumina2-style transformer + Qwen3-VL encoder as the +base T2I model, with reference-image conditioning. A reference image feeds the +model in TWO places: + + 1. Into the Qwen3-VL instruction encoder as image content alongside the edit + instruction (so the *text embeddings* already encode the reference image). + This is why ``encode_control_in_text_embeddings = True``. + 2. Into the transformer as reference-image VAE latents + (``ref_image_hidden_states``), which the ref-image refiner + double-stream + blocks attend to. + +Everything else (transformer, VAE, scheduler, time/velocity convention, saving) +is inherited from ``BooguImageModel`` -- this file only overrides the pieces +that change for TI2I. +""" + +import math +from typing import TYPE_CHECKING, List, Optional + +import torch +import torch.nn.functional as F +from PIL import Image +from torchvision.transforms.functional import to_tensor + +from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds +from toolkit.config_modules import GenerateImageConfig, ModelConfig + +from .boogu_image import BooguImageModel +from .src.pipeline import ( + BooguImagePipeline, + pad_instruction_features, + run_boogu_transformer, +) + +if TYPE_CHECKING: + from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO + + +# Edit release (clean bf16 safetensors); same layout as the base repo. +BOOGU_EDIT_PATH = "Boogu/Boogu-Image-0.1-Edit" + +# System prompt the edit model was trained with (SYSTEM_PROMPT_4_TI2I upstream). +SYSTEM_PROMPT_TI2I = ( + "Describe the key features of the input image (color, shape, size, texture, " + "objects, background), then explain how the user's text instruction should " + "alter or modify the image. Generate a new image that meets the user's " + "requirements while maintaining consistency with the original input where " + "appropriate." +) + + +class BooguImageEditModel(BooguImageModel): + arch = "boogu_image_edit" + default_repo = BOOGU_EDIT_PATH + + 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 + ) + # The reference image is encoded into the Qwen3-VL instruction features, + # so get_prompt_embeds receives the control image(s). + self.encode_control_in_text_embeddings = True + # Boogu supports up to 5 reference images -> they arrive as a list. + self.has_multiple_control_images = True + # Reference images keep their own aspect/size (not resized to the target). + self.use_raw_control_images = True + + @property + def text_embedding_space_version(self): + # Distinct from the base T2I cache: the edit features fold in the ref image. + return self.arch + "_v1" + + # ------------------------------------------------------------------ + # Reference-image helpers + # ------------------------------------------------------------------ + def _vlm_resize_hw(self, h, w, max_pixels, max_side, factor=16): + """Boogu's VLM image downscale (BooguImageProcessor.get_new_height_width). + + Scale down (never up) to fit BOTH ``max_pixels`` (area) and + ``max_side_length``, then round each dim down to a multiple of ``factor`` + (the image processor's ``vae_scale_factor`` = 16 for this model). The Qwen + processor's own smart_resize runs afterwards, exactly as upstream. + """ + longest = h if h > w else w + ratio_side = max_side / longest + ratio_pixels = (max_pixels / (h * w)) ** 0.5 + ratio = min(ratio_pixels, ratio_side, 1.0) + nh = max(factor, int(h * ratio) // factor * factor) + nw = max(factor, int(w * ratio) // factor * factor) + return nh, nw + + def _ref_target_pixels(self, target_pixels: Optional[int]) -> int: + """Decide the pixel budget each reference image is resized to fit within. + + - default: ``control_image_max_pixels`` model_kwarg (1 MP) -- a hard cap so + raw, full-size control images don't blow up the token count / VRAM. + - ``match_target_res`` model_kwarg: use the target generation area instead, + matching Boogu's recommendation of ``max_input_image_pixels ~= H*W``. + """ + max_pixels = int( + self.model_config.model_kwargs.get("control_image_max_pixels", 1024 * 1024) + ) + if ( + self.model_config.model_kwargs.get("match_target_res", False) + and target_pixels + ): + return int(target_pixels) + return max_pixels + + def _encode_ref_latents( + self, control_tensors, target_pixels: Optional[int] = None + ) -> List[torch.Tensor]: + """Encode ``[0, 1]`` reference image tensors to VAE latents. + + Returns a list of ``(16, h, w)`` latents (one per reference image). Each + control image is resized so its area fits within the pixel budget (see + ``_ref_target_pixels``) -- preserving aspect ratio -- then snapped so the + latent grid is divisible by the patch size. ``control_tensors`` is a list + of ``(C, H, W)`` or ``(1, C, H, W)`` tensors in ``[0, 1]``. + """ + sc = self.get_bucket_divisibility() # 16: VAE(8) * patch(2) + budget = self._ref_target_pixels(target_pixels) + match = self.model_config.model_kwargs.get("match_target_res", False) + + latents = [] + for img in control_tensors: + if img.dim() == 3: + img = img.unsqueeze(0) + img = img.to(self.device_torch, dtype=self.torch_dtype) + + h, w = img.shape[2], img.shape[3] + # match_target_res: scale area *to* the budget; otherwise only scale + # *down* when the image is larger than the budget. + area = h * w + if match or area > budget: + ratio = h / w + new_h = math.sqrt(budget * ratio) + new_w = new_h / ratio + else: + new_h, new_w = float(h), float(w) + + # snap to a multiple of the bucket divisibility so the VAE latent grid + # is patchifiable (the transformer rearranges 2x2 latent patches). + new_h = max(sc, int(round(new_h / sc)) * sc) + new_w = max(sc, int(round(new_w / sc)) * sc) + if (new_h, new_w) != (h, w): + img = F.interpolate(img, size=(new_h, new_w), mode="bilinear") + + # encode_images expects [-1, 1]; control tensors arrive in [0, 1]. + latent = self.encode_images( + img * 2 - 1, device=self.device_torch, dtype=self.torch_dtype + ) + latents.append(latent[0]) # drop batch dim -> (16, h, w) + return latents + + def _batch_ref_latents_from_batch( + self, + batch: "DataLoaderBatchDTO", + batch_size: int, + target_pixels: Optional[int] = None, + ) -> Optional[List[List[torch.Tensor]]]: + """Build the transformer's ``ref_image_hidden_states`` from a train batch.""" + control_list = batch.control_tensor_list + if control_list is None and batch.control_tensor is not None: + control_list = [batch.control_tensor[b : b + 1] for b in range(batch_size)] + if control_list is None: + return None + if len(control_list) != batch_size: + raise ValueError("Control tensor list length does not match batch size") + return [ + self._encode_ref_latents(controls, target_pixels=target_pixels) + for controls in control_list + ] + + # ------------------------------------------------------------------ + # Conditioning + # ------------------------------------------------------------------ + def get_prompt_embeds(self, prompt, control_images=None) -> AdvancedPromptEmbeds: + if isinstance(prompt, str): + prompt = [prompt] + + if control_images is None: + raise ValueError("BooguImageEditModel requires control (reference) images") + + # Normalize to List[List[Tensor]] (per-prompt list of reference images), the + # same convention qwen_image_edit_plus uses. + if not isinstance(control_images, list): + control_images = [control_images] + if not isinstance(control_images[0], list): + control_images = [control_images] + if len(prompt) != len(control_images): + raise ValueError( + "Number of prompts must match number of control image sets" + ) + + if self.text_encoder.device == torch.device("cpu"): + self.text_encoder.to(self.device_torch) + device = self.text_encoder.device + + features_list = [] + for p, ctrl in zip(prompt, control_images): + # Keep reference images as tensors the whole way (no GPU->CPU->PIL + # round-trip). Match Boogu's VLM preprocessing: downscale each control + # image to fit max_pixels (384^2) AND max_side_length (768) -- the MLLM + # only needs a coarse understanding of the reference (high-res detail + # flows through the VAE ref latents), and this keeps the instruction + # sequence well under the transformer rope axes_lens (~144 tokens/ref). + max_pixels = int( + self.model_config.model_kwargs.get("vlm_max_pixels", 384 * 384) + ) + max_side = int( + self.model_config.model_kwargs.get("vlm_max_side_length", 768) + ) + images = [] + for img in ctrl: + if img.dim() == 4: + img = img[0] + img = img.to(device) + nh, nw = self._vlm_resize_hw( + img.shape[1], img.shape[2], max_pixels, max_side + ) + if (nh, nw) != (img.shape[1], img.shape[2]): + img = ( + F.interpolate( + img.unsqueeze(0), + size=(nh, nw), + mode="bicubic", + antialias=True, + ) + .squeeze(0) + .clamp(0, 1) + ) + images.append(img) + + # Build just the text template with image placeholders (tokenize=False), + # then let the processor expand the image tokens from the real grid size. + user_content = [{"type": "image"} for _ in images] + user_content.append({"type": "text", "text": p}) + messages = [ + { + "role": "system", + "content": [{"type": "text", "text": SYSTEM_PROMPT_TI2I}], + }, + {"role": "user", "content": user_content}, + ] + text = self.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=False + ) + # do_rescale=False: control tensors are already [0, 1] (the image + # normalizer maps them to [-1, 1]). No size override -- the images are + # already at Boogu's target size, the processor just snaps to its grid. + inputs = self.tokenizer( + text=[text], + images=images, + return_tensors="pt", + do_rescale=False, + ) + model_inputs = {} + for k, v in inputs.items(): + if isinstance(v, torch.Tensor): + v = v.to(device) + # cast image pixels to the encoder dtype; leave ids/masks as ints + if v.is_floating_point(): + v = v.to(self.torch_dtype) + model_inputs[k] = v + + with torch.no_grad(): + output = self.text_encoder(**model_inputs) + features_list.append(output.last_hidden_state[0].to(self.torch_dtype)) + + return AdvancedPromptEmbeds(text_embeds=features_list) + + 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, + batch: "DataLoaderBatchDTO" = None, + **kwargs, + ): + if self.model.device == torch.device("cpu"): + self.model.to(self.device_torch) + + with torch.no_grad(): + # target pixel area from the noise latents (h, w are VAE-downsampled) + _, _, lh, lw = latent_model_input.shape + target_pixels = (lh * self.vae_scale_factor) * (lw * self.vae_scale_factor) + ref_latents = ( + self._batch_ref_latents_from_batch( + batch, latent_model_input.shape[0], target_pixels=target_pixels + ) + if batch is not None + else None + ) + + # 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. + 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(), + ref_image_hidden_states=ref_latents, + ) + return -raw_velocity + + # ------------------------------------------------------------------ + # Sampling previews + # ------------------------------------------------------------------ + 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) + + # Load the reference image(s) for the transformer ref latents. The MLLM + # side already saw them (baked into conditional/unconditional embeds). + ctrl_paths = [ + p + for p in ( + gen_config.ctrl_img, + gen_config.ctrl_img_1, + gen_config.ctrl_img_2, + gen_config.ctrl_img_3, + ) + if p is not None + ] + ref_latents = None + if ctrl_paths: + ctrl_tensors = [ + to_tensor(Image.open(path).convert("RGB")) for path in ctrl_paths + ] + target_pixels = gen_config.width * gen_config.height + # one batch item (preview batch size is 1) -> List[List[(16, h, w)]] + ref_latents = [ + self._encode_ref_latents(ctrl_tensors, target_pixels=target_pixels) + ] + + 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, + ref_latents=ref_latents, + )[0] + return img + + def get_base_model_version(self): + return "boogu_image_edit.0.1" diff --git a/extensions_built_in/diffusion_models/boogu_image/src/__init__.py b/extensions_built_in/diffusion_models/boogu_image/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/extensions_built_in/diffusion_models/boogu_image/src/attention_processor.py b/extensions_built_in/diffusion_models/boogu_image/src/attention_processor.py new file mode 100644 index 0000000..6c52446 --- /dev/null +++ b/extensions_built_in/diffusion_models/boogu_image/src/attention_processor.py @@ -0,0 +1,491 @@ +# Vendored from the Boogu-Image repository (boogu/models/attention_processor.py). +# Original work: Copyright 2025 BAAI / OmniGen2 / HuggingFace. Apache-2.0. +# +# Attention here defaults to torch's ``scaled_dot_product_attention`` (the +# "native" backend) so the model has NO hard dependency on flash-attn. Flash +# Attention 2 is an OPTIONAL backend: each processor carries an +# ``attention_backend`` flag (set in bulk via +# ``BooguImageTransformer2DModel.set_attention_backend``) and only the "flash" +# branch touches the ``flash_attn`` package, so importing it stays lazy/guarded. +import math +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.attention_processor import Attention +from einops import repeat + +from .embeddings import apply_rotary_emb + +try: + from flash_attn import flash_attn_varlen_func + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input + + _FLASH_ATTN_AVAILABLE = True +except ImportError: # flash-attn is optional; "native" SDPA needs none of this. + flash_attn_varlen_func = None + index_first_axis = pad_input = unpad_input = None + _FLASH_ATTN_AVAILABLE = False + +# Supported attention backends. "native" -> SDPA, "flash" -> Flash Attention 2. +ATTENTION_BACKENDS = ("native", "flash") + + +def _get_unpad_data(mask_2d: torch.Tensor): + """Indices / cu_seqlens / max_seqlen from a 2D padding mask [B, L].""" + seqlens_in_batch = mask_2d.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(mask_2d.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + return indices, cu_seqlens, max_seqlen_in_batch + + +def _upad_input(query, key, value, attention_mask, query_length, num_heads): + """Unpad q/k/v for ``flash_attn_varlen_func`` given a [B, L] padding mask.""" + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + batch_size, kv_seq_len, num_key_value_heads, head_dim = key.shape + + key = index_first_axis( + key.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + value = index_first_axis( + value.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + + if query_length == kv_seq_len: + query = index_first_axis( + query.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k + ) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query.device + ) + indices_q = cu_seqlens_q[:-1] + query = query.squeeze(1) + else: + q_mask = attention_mask[:, -query_length:] + query, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input( + query, q_mask + ) + + return ( + query, + key, + value, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +def _flash_varlen_attention(query, key, value, attention_mask, attn, softmax_scale): + """Run flash-attn varlen over a [B, L, heads, head_dim] q/k/v with a 2D mask. + + Returns the attention output flattened back to [B, L, heads * head_dim]. + """ + batch_size, sequence_length = query.shape[0], query.shape[1] + kv_heads = key.shape[2] + + mask_2d = attention_mask.bool() if attention_mask is not None else None + ( + query_states, + key_states, + value_states, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_q, max_seqlen_k), + ) = _upad_input(query, key, value, mask_2d, sequence_length, attn.heads) + + if kv_heads < attn.heads: + key_states = repeat(key_states, "l h c -> l (h k) c", k=attn.heads // kv_heads) + value_states = repeat( + value_states, "l h c -> l (h k) c", k=attn.heads // kv_heads + ) + + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + dropout_p=0.0, + causal=False, + softmax_scale=softmax_scale, + ) + hidden_states = pad_input(attn_output_unpad, indices_q, batch_size, sequence_length) + return hidden_states.flatten(-2) + + +class BooguImageDoubleStreamSelfAttnProcessor(nn.Module): + """ + Double-stream self-attention processor. + + Instruction and image features each get their own q/k/v projections; the two + streams are concatenated (instruction first), attended jointly, then split + back and projected with separate output heads. Uses torch SDPA by default; + set ``attention_backend = "flash"`` for Flash Attention 2. + """ + + def __init__( + self, + head_dim: int, + num_attention_heads: int, + num_kv_heads: int, + qkv_bias: bool = False, + ) -> None: + super().__init__() + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError( + "BooguImageDoubleStreamSelfAttnProcessor requires PyTorch 2.0+." + ) + + self.head_dim = head_dim + self.num_attention_heads = num_attention_heads + self.num_kv_heads = num_kv_heads + self.attention_backend = "native" + + query_dim = head_dim * num_attention_heads + kv_dim = head_dim * num_kv_heads + + self.img_to_q = nn.Linear(query_dim, query_dim, bias=qkv_bias) + self.img_to_k = nn.Linear(query_dim, kv_dim, bias=qkv_bias) + self.img_to_v = nn.Linear(query_dim, kv_dim, bias=qkv_bias) + + self.instruct_to_q = nn.Linear(query_dim, query_dim, bias=qkv_bias) + self.instruct_to_k = nn.Linear(query_dim, kv_dim, bias=qkv_bias) + self.instruct_to_v = nn.Linear(query_dim, kv_dim, bias=qkv_bias) + + self.instruct_out = nn.Linear(query_dim, query_dim, bias=qkv_bias) + self.img_out = nn.Linear(query_dim, query_dim, bias=qkv_bias) + + self.initialize_weights() + + def initialize_weights(self) -> None: + nn.init.xavier_uniform_(self.img_to_q.weight) + nn.init.xavier_uniform_(self.img_to_k.weight) + nn.init.xavier_uniform_(self.img_to_v.weight) + nn.init.xavier_uniform_(self.instruct_to_q.weight) + nn.init.xavier_uniform_(self.instruct_to_k.weight) + nn.init.xavier_uniform_(self.instruct_to_v.weight) + nn.init.xavier_uniform_(self.instruct_out.weight) + nn.init.xavier_uniform_(self.img_out.weight) + + if self.img_to_q.bias is not None: + nn.init.zeros_(self.img_to_q.bias) + nn.init.zeros_(self.img_to_k.bias) + nn.init.zeros_(self.img_to_v.bias) + nn.init.zeros_(self.instruct_to_q.bias) + nn.init.zeros_(self.instruct_to_k.bias) + nn.init.zeros_(self.instruct_to_v.bias) + nn.init.zeros_(self.instruct_out.bias) + nn.init.zeros_(self.img_out.bias) + + def _concat_instruction_image_features( + self, + img_hidden_states_list: List[torch.Tensor], + instruct_hidden_states_list: List[torch.Tensor], + encoder_seq_lengths: List[int], + seq_lengths: List[int], + ) -> List[torch.Tensor]: + """Concatenate instruction then image features into one joint sequence.""" + batch_size = img_hidden_states_list[0].shape[0] + max_seq_len = max(seq_lengths) + + concatenated_list = [] + for img_tensor, instruct_tensor in zip( + img_hidden_states_list, instruct_hidden_states_list + ): + device = img_tensor.device + if instruct_tensor.device != device: + instruct_tensor = instruct_tensor.to(device) + + feature_dim = img_tensor.shape[-1] + concatenated = img_tensor.new_zeros(batch_size, max_seq_len, feature_dim) + + for i, (encoder_seq_len, seq_len) in enumerate( + zip(encoder_seq_lengths, seq_lengths) + ): + concatenated[i, :encoder_seq_len] = instruct_tensor[i, :encoder_seq_len] + concatenated[i, encoder_seq_len:seq_len] = img_tensor[ + i, : seq_len - encoder_seq_len + ] + + concatenated_list.append(concatenated) + + return concatenated_list + + def _split_instruction_image_features( + self, + hidden_states_list: List[torch.Tensor], + encoder_seq_lengths: List[int], + seq_lengths: List[int], + ) -> List[Tuple[torch.Tensor, torch.Tensor]]: + """Inverse of ``_concat_instruction_image_features``.""" + result_list = [] + for hidden_states in hidden_states_list: + batch_size = hidden_states.shape[0] + feature_dim = hidden_states.shape[-1] + + max_instruct_len = max(encoder_seq_lengths) + max_img_len = max( + seq_len - encoder_seq_len + for seq_len, encoder_seq_len in zip(seq_lengths, encoder_seq_lengths) + ) + + instruct_hidden_states = hidden_states.new_zeros( + batch_size, max_instruct_len, feature_dim + ) + img_hidden_states = hidden_states.new_zeros( + batch_size, max_img_len, feature_dim + ) + + for i, (encoder_seq_len, seq_len) in enumerate( + zip(encoder_seq_lengths, seq_lengths) + ): + img_len = seq_len - encoder_seq_len + instruct_hidden_states[i, :encoder_seq_len] = hidden_states[ + i, :encoder_seq_len + ] + img_hidden_states[i, :img_len] = hidden_states[ + i, encoder_seq_len:seq_len + ] + + result_list.append((instruct_hidden_states, img_hidden_states)) + + return result_list + + def __call__( + self, + attn: Attention, + img_hidden_states: torch.Tensor, + instruct_hidden_states: torch.Tensor, + joint_attention_mask: Optional[torch.Tensor] = None, + rotary_emb: Optional[torch.Tensor] = None, + encoder_seq_lengths: List[int] = None, + seq_lengths: List[int] = None, + base_sequence_length: Optional[int] = None, + ) -> torch.Tensor: + batch_size = img_hidden_states.shape[0] + + img_query = self.img_to_q(img_hidden_states) + img_key = self.img_to_k(img_hidden_states) + img_value = self.img_to_v(img_hidden_states) + + instruct_query = self.instruct_to_q(instruct_hidden_states) + instruct_key = self.instruct_to_k(instruct_hidden_states) + instruct_value = self.instruct_to_v(instruct_hidden_states) + + img_list = [img_query, img_key, img_value] + instruct_list = [instruct_query, instruct_key, instruct_value] + concatenated_list = self._concat_instruction_image_features( + img_list, instruct_list, encoder_seq_lengths, seq_lengths + ) + query, key, value = concatenated_list + + sequence_length = max(seq_lengths) + + query_dim = query.shape[-1] + inner_dim = key.shape[-1] + head_dim = query_dim // attn.heads + dtype = query.dtype + + kv_heads = inner_dim // head_dim + + query = query.view(batch_size, -1, attn.heads, head_dim) + key = key.view(batch_size, -1, kv_heads, head_dim) + value = value.view(batch_size, -1, kv_heads, head_dim) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + if rotary_emb is not None: + query = apply_rotary_emb(query, rotary_emb, use_real=False) + key = apply_rotary_emb(key, rotary_emb, use_real=False) + + query, key = query.to(dtype), key.to(dtype) + + if base_sequence_length is not None: + softmax_scale = ( + math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale + ) + else: + softmax_scale = attn.scale + + if self.attention_backend == "flash": + # q/k/v are [B, L, heads, head_dim]; the joint padding mask is 2D. + hidden_states = _flash_varlen_attention( + query, key, value, joint_attention_mask, attn, softmax_scale + ) + hidden_states = hidden_states.type_as(query) + else: + if joint_attention_mask is not None: + joint_attention_mask = joint_attention_mask.bool() + if joint_attention_mask.dim() == 2: + joint_attention_mask = joint_attention_mask.view( + batch_size, 1, 1, -1 + ) + elif joint_attention_mask.dim() == 3: + joint_attention_mask = joint_attention_mask.unsqueeze(1) + else: + raise ValueError( + f"Unsupported joint_attention_mask shape: {joint_attention_mask.shape}" + ) + + q = query.transpose(1, 2) + k = key.transpose(1, 2) + v = value.transpose(1, 2) + + # explicitly repeat key/value to avoid the slow MATH SDPA backend that + # enable_gqa triggers on some torch builds + k = k.repeat_interleave(q.size(-3) // k.size(-3), -3) + v = v.repeat_interleave(q.size(-3) // v.size(-3), -3) + + hidden_states = F.scaled_dot_product_attention( + q, k, v, attn_mask=joint_attention_mask, scale=softmax_scale + ) + hidden_states = hidden_states.transpose(1, 2).reshape( + batch_size, -1, attn.heads * head_dim + ) + hidden_states = hidden_states.type_as(query) + + split_results = self._split_instruction_image_features( + [hidden_states], encoder_seq_lengths, seq_lengths + ) + instruct_hidden_states, img_hidden_states = split_results[0] + + instruct_projected = self.instruct_out(instruct_hidden_states) + img_projected = self.img_out(img_hidden_states) + + merged_list = self._concat_instruction_image_features( + [img_projected], [instruct_projected], encoder_seq_lengths, seq_lengths + ) + hidden_states = merged_list[0] + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + + return hidden_states + + +class BooguImageAttnProcessor: + """ + Single-stream self-attention processor with RoPE + QK norm. + + Uses torch SDPA by default; set ``attention_backend = "flash"`` for Flash + Attention 2 (requires the ``flash_attn`` package). + """ + + def __init__(self) -> None: + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("BooguImageAttnProcessor requires PyTorch 2.0+.") + self.attention_backend = "native" + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + base_sequence_length: Optional[int] = None, + ) -> torch.Tensor: + batch_size, sequence_length, _ = hidden_states.shape + + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + query_dim = query.shape[-1] + inner_dim = key.shape[-1] + head_dim = query_dim // attn.heads + dtype = query.dtype + + kv_heads = inner_dim // head_dim + + query = query.view(batch_size, -1, attn.heads, head_dim) + key = key.view(batch_size, -1, kv_heads, head_dim) + value = value.view(batch_size, -1, kv_heads, head_dim) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + if image_rotary_emb is not None: + query = apply_rotary_emb(query, image_rotary_emb, use_real=False) + key = apply_rotary_emb(key, image_rotary_emb, use_real=False) + + query, key = query.to(dtype), key.to(dtype) + + if base_sequence_length is not None: + softmax_scale = ( + math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale + ) + else: + softmax_scale = attn.scale + + if self.attention_backend == "flash" and ( + attention_mask is None or attention_mask.dim() == 2 + ): + mask = ( + attention_mask + if attention_mask is not None + else query.new_ones(batch_size, sequence_length, dtype=torch.bool) + ) + hidden_states = _flash_varlen_attention( + query, key, value, mask, attn, softmax_scale + ) + hidden_states = hidden_states.type_as(query) + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + if attention_mask is not None: + attention_mask = attention_mask.bool() + if attention_mask.dim() == 2: + attention_mask = attention_mask.view(batch_size, 1, 1, -1) + elif attention_mask.dim() == 3: + B, L, _ = attention_mask.shape + diag_valid = torch.diagonal(attention_mask, dim1=-2, dim2=-1) + lengths = diag_valid.sum(dim=-1) + arange_L = torch.arange(L, device=attention_mask.device) + q_valid = arange_L.unsqueeze(0) < lengths.unsqueeze(1) + k_valid = q_valid + causal = torch.tril( + torch.ones(L, L, dtype=torch.bool, device=attention_mask.device) + ) + combined = causal & q_valid.unsqueeze(-1) & k_valid.unsqueeze(-2) + attention_mask = combined.unsqueeze(1) + else: + raise ValueError( + f"Unsupported attention_mask shape: {attention_mask.shape}" + ) + + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + key = key.repeat_interleave(query.size(-3) // key.size(-3), -3) + value = value.repeat_interleave(query.size(-3) // value.size(-3), -3) + + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, scale=softmax_scale + ) + hidden_states = hidden_states.transpose(1, 2).reshape( + batch_size, -1, attn.heads * head_dim + ) + hidden_states = hidden_states.type_as(query) + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + + return hidden_states diff --git a/extensions_built_in/diffusion_models/boogu_image/src/block_lumina2.py b/extensions_built_in/diffusion_models/boogu_image/src/block_lumina2.py new file mode 100644 index 0000000..d243d58 --- /dev/null +++ b/extensions_built_in/diffusion_models/boogu_image/src/block_lumina2.py @@ -0,0 +1,164 @@ +# Vendored from the Boogu-Image repository (boogu/models/transformers/block_lumina2.py). +# Original work: Copyright 2025 BAAI / OmniGen2 / HuggingFace. Apache-2.0. +# +# The optional triton RMSNorm and flash-attn SwiGLU fast paths are dropped here; +# we always use torch.nn.RMSNorm and a plain SwiGLU so the model runs anywhere. +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.embeddings import Timesteps +from torch.nn import RMSNorm + +from .embeddings import TimestepEmbedding + + +def swiglu(x, y): + return F.silu(x.float(), inplace=False).to(x.dtype) * y + + +class LuminaRMSNormZero(nn.Module): + """Adaptive RMS normalization with a zero-initialized modulation projection.""" + + def __init__( + self, + embedding_dim: int, + norm_eps: float, + norm_elementwise_affine: bool, + ): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear( + min(embedding_dim, 1024), + 4 * embedding_dim, + bias=True, + ) + + self.norm = RMSNorm(embedding_dim, eps=norm_eps) + + def forward( + self, + x: torch.Tensor, + emb: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb = self.linear(self.silu(emb)) + scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1) + x = self.norm(x) * (1 + scale_msa[:, None]) + return x, gate_msa, scale_mlp, gate_mlp + + +class LuminaLayerNormContinuous(nn.Module): + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + elementwise_affine=True, + eps=1e-5, + bias=True, + norm_type="layer_norm", + out_dim: Optional[int] = None, + ): + super().__init__() + + # AdaLN + self.silu = nn.SiLU() + self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias) + + if norm_type == "layer_norm": + self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) + elif norm_type == "rms_norm": + self.norm = RMSNorm( + embedding_dim, eps=eps, elementwise_affine=elementwise_affine + ) + else: + raise ValueError(f"unknown norm_type {norm_type}") + + self.linear_2 = None + if out_dim is not None: + self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias) + + def forward( + self, + x: torch.Tensor, + conditioning_embedding: torch.Tensor, + ) -> torch.Tensor: + emb = self.linear_1(self.silu(conditioning_embedding).to(x.dtype)) + scale = emb + x = self.norm(x) * (1 + scale)[:, None, :] + + if self.linear_2 is not None: + x = self.linear_2(x) + + return x + + +class LuminaFeedForward(nn.Module): + """A SwiGLU feed-forward layer with a multiple-of-256 inner dim.""" + + def __init__( + self, + dim: int, + inner_dim: int, + multiple_of: Optional[int] = 256, + ffn_dim_multiplier: Optional[float] = None, + ): + super().__init__() + self.swiglu = swiglu + + if ffn_dim_multiplier is not None: + inner_dim = int(ffn_dim_multiplier * inner_dim) + inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of) + + self.linear_1 = nn.Linear(dim, inner_dim, bias=False) + self.linear_2 = nn.Linear(inner_dim, dim, bias=False) + self.linear_3 = nn.Linear(dim, inner_dim, bias=False) + + def forward(self, x): + h1, h2 = self.linear_1(x), self.linear_3(x) + return self.linear_2(self.swiglu(h1, h2)) + + +class Lumina2CombinedTimestepCaptionEmbedding(nn.Module): + def __init__( + self, + hidden_size: int = 4096, + instruction_feat_dim: int = 2048, + frequency_embedding_size: int = 256, + norm_eps: float = 1e-5, + timestep_scale: float = 1.0, + ) -> None: + super().__init__() + + self.time_proj = Timesteps( + num_channels=frequency_embedding_size, + flip_sin_to_cos=True, + downscale_freq_shift=0.0, + scale=timestep_scale, + ) + + self.timestep_embedder = TimestepEmbedding( + in_channels=frequency_embedding_size, time_embed_dim=min(hidden_size, 1024) + ) + + self.caption_embedder = nn.Sequential( + RMSNorm(instruction_feat_dim, eps=norm_eps), + nn.Linear(instruction_feat_dim, hidden_size, bias=True), + ) + + self._initialize_weights() + + def _initialize_weights(self): + nn.init.trunc_normal_(self.caption_embedder[1].weight, std=0.02) + nn.init.zeros_(self.caption_embedder[1].bias) + + def forward( + self, + timestep: torch.Tensor, + instruction_hidden_states: torch.Tensor, + dtype: torch.dtype, + ) -> Tuple[torch.Tensor, torch.Tensor]: + timestep_proj = self.time_proj(timestep).to(dtype=dtype) + time_embed = self.timestep_embedder(timestep_proj) + caption_embed = self.caption_embedder(instruction_hidden_states) + return time_embed, caption_embed diff --git a/extensions_built_in/diffusion_models/boogu_image/src/embeddings.py b/extensions_built_in/diffusion_models/boogu_image/src/embeddings.py new file mode 100644 index 0000000..f89429e --- /dev/null +++ b/extensions_built_in/diffusion_models/boogu_image/src/embeddings.py @@ -0,0 +1,112 @@ +# Vendored from the Boogu-Image repository (boogu/models/embeddings.py). +# Original work: Copyright 2024 The HuggingFace Team. Apache-2.0. +# +# Only the pieces the Boogu transformer actually needs are kept here: +# ``TimestepEmbedding`` and ``apply_rotary_emb``. +from typing import Optional, Tuple, Union + +import torch +from diffusers.models.activations import get_activation +from torch import nn + + +class TimestepEmbedding(nn.Module): + def __init__( + self, + in_channels: int, + time_embed_dim: int, + act_fn: str = "silu", + out_dim: int = None, + post_act_fn: Optional[str] = None, + cond_proj_dim=None, + sample_proj_bias=True, + ): + super().__init__() + + self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias) + + if cond_proj_dim is not None: + self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) + else: + self.cond_proj = None + + self.act = get_activation(act_fn) + + if out_dim is not None: + time_embed_dim_out = out_dim + else: + time_embed_dim_out = time_embed_dim + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) + + if post_act_fn is None: + self.post_act = None + else: + self.post_act = get_activation(post_act_fn) + + self.initialize_weights() + + def initialize_weights(self): + nn.init.normal_(self.linear_1.weight, std=0.02) + nn.init.zeros_(self.linear_1.bias) + nn.init.normal_(self.linear_2.weight, std=0.02) + nn.init.zeros_(self.linear_2.bias) + + def forward(self, sample, condition=None): + if condition is not None: + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + + if self.act is not None: + sample = self.act(sample) + + sample = self.linear_2(sample) + + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + +def apply_rotary_emb( + x: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], + use_real: bool = True, + use_real_unbind_dim: int = -1, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Apply rotary embeddings to input tensors using the given frequency tensor. + + Boogu always calls this with ``use_real=False`` (the Lumina-style complex + path): ``freqs_cis`` is a complex tensor and ``x`` is reinterpreted as + complex, multiplied, and returned as real. + """ + if use_real: + cos, sin = freqs_cis # [S, D] + cos = cos[None, None] + sin = sin[None, None] + cos, sin = cos.to(x.device), sin.to(x.device) + + if use_real_unbind_dim == -1: + # Used for flux, cogvideox, hunyuan-dit + x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) + elif use_real_unbind_dim == -2: + # Used for Stable Audio, Boogu and CogView4 + x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) + x_rotated = torch.cat([-x_imag, x_real], dim=-1) + else: + raise ValueError( + f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2." + ) + + out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) + + return out + else: + # used for lumina / boogu + x_rotated = torch.view_as_complex( + x.float().reshape(*x.shape[:-1], x.shape[-1] // 2, 2) + ) + freqs_cis = freqs_cis.unsqueeze(2) + x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) + + return x_out.type_as(x) diff --git a/extensions_built_in/diffusion_models/boogu_image/src/pipeline.py b/extensions_built_in/diffusion_models/boogu_image/src/pipeline.py new file mode 100644 index 0000000..780af4f --- /dev/null +++ b/extensions_built_in/diffusion_models/boogu_image/src/pipeline.py @@ -0,0 +1,231 @@ +"""Packing / sampling helpers for Boogu-Image (base T2I). + +This module glues the Qwen3-VL instruction features and the image latents into +the call the Boogu transformer expects, and provides a minimal flow-matching +sampler used to render preview images during training. + +Time convention +--------------- +Boogu's native flow time is ``t in [0, 1]`` with ``t=0`` pure noise and ``t=1`` +clean; the transformer predicts ``clean - noise``. ai-toolkit's scheduler uses +the opposite convention (``t=1`` noise, velocity ``noise - clean``). The +conversion lives in ``BooguImageModel.get_noise_prediction``; this sampler runs +entirely in Boogu's native domain via :func:`run_boogu_transformer`. +""" + +from __future__ import annotations + +import math +from typing import List, Optional + +import numpy as np +import torch +from PIL import Image +from diffusers.utils.torch_utils import randn_tensor + +from .transformer import BooguImageTransformer2DModel + + +# --------------------------------------------------------------------------- +# Instruction feature padding. +# --------------------------------------------------------------------------- + + +def pad_instruction_features( + features_list: List[torch.Tensor], + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Right-pad per-sample ``(L_i, D)`` instruction features into a batch. + + Captions are stored per-sample at their natural length and only padded to the + batch max here, right before the model call. Returns ``(features (B, L, D), + attention_mask (B, L))`` with the mask 1 for real tokens, 0 for padding. + """ + lengths = [f.shape[0] for f in features_list] + max_len = max(lengths) + dim = features_list[0].shape[-1] + batch_size = len(features_list) + + features = torch.zeros(batch_size, max_len, dim, device=device, dtype=dtype) + mask = torch.zeros(batch_size, max_len, dtype=torch.long, device=device) + for i, f in enumerate(features_list): + n = f.shape[0] + features[i, :n] = f.to(device, dtype) + mask[i, :n] = 1 + return features, mask + + +# --------------------------------------------------------------------------- +# Time-shift schedule (mirrors the released Boogu base scheduler: v1 shift). +# --------------------------------------------------------------------------- + + +def _lin_shift( + num_tokens: float, + x1: float = 256.0, + y1: float = 0.5, + x2: float = 4096.0, + y2: float = 1.15, +) -> float: + """Linear token-count -> mu mapping (Boogu base_shift/max_shift defaults).""" + m = (y2 - y1) / (x2 - x1) + b = y1 - m * x1 + return m * num_tokens + b + + +def boogu_time_schedule( + num_steps: int, + num_patch_tokens: int, + device: Optional[torch.device] = None, +) -> torch.Tensor: + """Boogu native-domain timesteps (0=noise .. 1=clean) with v1 time shift. + + Returns a length ``num_steps + 1`` tensor; the trailing ``1.0`` is the clean + endpoint, matching the ``_timesteps`` tail in the reference scheduler. + """ + t_arr = np.linspace(0.0, 1.0, num_steps + 1, dtype=np.float32)[:-1] + + mu = _lin_shift(max(1, int(num_patch_tokens))) + eps = 1e-8 + t1 = np.clip(1.0 - t_arr, eps, 1.0 - eps) + num = math.exp(mu) + denom = num + (1.0 / t1 - 1.0) + t_arr = (1.0 - num / denom).astype(np.float32) + + times = np.concatenate([t_arr, np.ones(1, dtype=np.float32)]) + return torch.from_numpy(times).to(device=device, dtype=torch.float32) + + +# --------------------------------------------------------------------------- +# Transformer call (Boogu native time domain). +# --------------------------------------------------------------------------- + + +def run_boogu_transformer( + transformer: BooguImageTransformer2DModel, + latents: torch.Tensor, # (B, 16, H, W) + boogu_t: torch.Tensor, # (B,) in [0, 1], 0=noise, 1=clean + instruction_features: torch.Tensor, # (B, L, instruction_feat_dim) + instruction_mask: torch.Tensor, # (B, L) 1 for real tokens + freqs_cis, # precomputed per-axis rotary tables + ref_image_hidden_states=None, # edit/TI2I: List[List[(16, H, W)]] per batch item +) -> torch.Tensor: + """Run the transformer and return the raw model velocity (``clean - noise``). + + Shapes pass straight through: the prediction comes back as ``(B, 16, H, W)`` + in the same latent layout as ``latents``. ``ref_image_hidden_states`` stays + ``None`` for the base T2I model and carries reference-image VAE latents for + the edit (TI2I) model. + """ + out = transformer( + hidden_states=latents, + timestep=boogu_t, + instruction_hidden_states=instruction_features, + freqs_cis=freqs_cis, + instruction_attention_mask=instruction_mask, + ref_image_hidden_states=ref_image_hidden_states, + return_dict=False, + ) + return out + + +# --------------------------------------------------------------------------- +# Minimal sampling pipeline (for training previews). +# --------------------------------------------------------------------------- + + +class BooguImagePipeline: + """Lightweight flow-matching sampler used by ai-toolkit's preview generation.""" + + def __init__(self, model): + # ``model`` is the BooguImageModel so we can reuse its encode/decode and + # latent helpers without duplicating state. + self.model = model + + @property + def device(self): + return self.model.device_torch + + def to(self, *args, **kwargs): + return self + + @torch.no_grad() + def __call__( + self, + conditional_embeds, + unconditional_embeds, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 50, + guidance_scale: float = 4.0, + latents: Optional[torch.Tensor] = None, + generator: Optional[torch.Generator] = None, + ref_latents=None, # edit/TI2I: List[List[(16, H, W)]] reference VAE latents + **kwargs, + ) -> List[Image.Image]: + model = self.model + device = model.device_torch + dtype = model.torch_dtype + transformer = model.transformer + patch = model.patch_size + ae_scale = model.vae_scale_factor # 8 + + latent_channels = transformer.config.in_channels + h_lat = height // ae_scale + w_lat = width // ae_scale + num_patch_tokens = (h_lat // patch) * (w_lat // patch) + + freqs_cis = model.get_freqs_cis() + + do_cfg = guidance_scale > 1.0 + + if latents is None: + shape = (1, latent_channels, h_lat, w_lat) + latents = randn_tensor( + shape, generator=generator, device=device, dtype=torch.float32 + ) + # In Boogu's domain t=0 is pure noise, so the initial latent IS the noise. + latents = latents.to(device, dtype=torch.float32) + + cond_feats, cond_mask = pad_instruction_features( + conditional_embeds.text_embeds, device, dtype + ) + if do_cfg: + uncond_feats, uncond_mask = pad_instruction_features( + unconditional_embeds.text_embeds, device, dtype + ) + + times = boogu_time_schedule(num_inference_steps, num_patch_tokens, device) + + for t, t_next in zip(times[:-1], times[1:]): + boogu_t = t.expand(latents.shape[0]) + v_cond = run_boogu_transformer( + transformer, + latents.to(dtype), + boogu_t, + cond_feats, + cond_mask, + freqs_cis, + ref_image_hidden_states=ref_latents, + ) + if do_cfg: + v_uncond = run_boogu_transformer( + transformer, + latents.to(dtype), + boogu_t, + uncond_feats, + uncond_mask, + freqs_cis, + ref_image_hidden_states=ref_latents, + ) + v = v_uncond + guidance_scale * (v_cond - v_uncond) + else: + v = v_cond + latents = latents + v.to(torch.float32) * (t_next - t) + + images = model.decode_latents(latents, device=device, dtype=dtype) + images = images.float().clamp(-1.0, 1.0) + images = ((images + 1.0) * 127.5).round().to(torch.uint8) + images = images.permute(0, 2, 3, 1).cpu().numpy() + return [Image.fromarray(arr) for arr in images] diff --git a/extensions_built_in/diffusion_models/boogu_image/src/rope.py b/extensions_built_in/diffusion_models/boogu_image/src/rope.py new file mode 100644 index 0000000..0184454 --- /dev/null +++ b/extensions_built_in/diffusion_models/boogu_image/src/rope.py @@ -0,0 +1,244 @@ +# Vendored from the Boogu-Image repository (boogu/models/transformers/rope.py). +# Original work: Copyright 2025 BAAI / OmniGen2 / HuggingFace. Apache-2.0. +# +# Only the double-stream rotary embedder (the one the transformer uses) and the +# ``get_freqs_cis`` precompute helper are kept. The MPS-specific branch is +# preserved verbatim. +from typing import List, Tuple + +import torch +import torch.nn as nn +from diffusers.models.embeddings import get_1d_rotary_pos_embed +from einops import repeat + + +def get_freqs_cis( + axes_dim: Tuple[int, int, int], axes_lens: Tuple[int, int, int], theta: int +) -> List[torch.Tensor]: + """Precompute the per-axis rotary frequency tables (done once per resolution).""" + freqs_cis = [] + freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 + for d, e in zip(axes_dim, axes_lens): + emb = get_1d_rotary_pos_embed(d, e, theta=theta, freqs_dtype=freqs_dtype) + freqs_cis.append(emb) + return freqs_cis + + +class BooguImageDoubleStreamRotaryPosEmbed(nn.Module): + def __init__( + self, + theta: int, + axes_dim: Tuple[int, int, int], + axes_lens: Tuple[int, int, int] = (300, 512, 512), + patch_size: int = 2, + ): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + self.axes_lens = axes_lens + self.patch_size = patch_size + + @staticmethod + def get_freqs_cis( + axes_dim: Tuple[int, int, int], axes_lens: Tuple[int, int, int], theta: int + ) -> List[torch.Tensor]: + return get_freqs_cis(axes_dim, axes_lens, theta) + + def _get_freqs_cis(self, freqs_cis, ids: torch.Tensor) -> torch.Tensor: + device = ids.device + if ids.device.type == "mps": + ids = ids.to("cpu") + + result = [] + for i in range(len(self.axes_dim)): + freqs = freqs_cis[i].to(ids.device) + index = ids[:, :, i : i + 1].repeat(1, 1, freqs.shape[-1]).to(torch.int64) + result.append( + torch.gather( + freqs.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index + ) + ) + return torch.cat(result, dim=-1).to(device) + + def forward( + self, + freqs_cis, + attention_mask, + l_effective_ref_img_len, + l_effective_img_len, + ref_img_sizes, + img_sizes, + device, + ): + batch_size = len(attention_mask) + p = self.patch_size + + encoder_seq_len = attention_mask.shape[1] + l_effective_cap_len = attention_mask.sum(dim=1).tolist() + + seq_lengths = [ + cap_len + sum(ref_img_len) + img_len + for cap_len, ref_img_len, img_len in zip( + l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len + ) + ] + + max_seq_len = max(seq_lengths) + max_ref_img_len = max( + [sum(ref_img_len) for ref_img_len in l_effective_ref_img_len] + ) + max_img_len = max(l_effective_img_len) + + # Create position IDs + position_ids = torch.zeros( + batch_size, max_seq_len, 3, dtype=torch.int32, device=device + ) + + for i, (cap_seq_len, seq_len) in enumerate( + zip(l_effective_cap_len, seq_lengths) + ): + # add text position ids + position_ids[i, :cap_seq_len] = repeat( + torch.arange(cap_seq_len, dtype=torch.int32, device=device), "l -> l 3" + ) + + pe_shift = cap_seq_len + pe_shift_len = cap_seq_len + + if ref_img_sizes[i] is not None: + for ref_img_size, ref_img_len in zip( + ref_img_sizes[i], l_effective_ref_img_len[i] + ): + H, W = ref_img_size + ref_H_tokens, ref_W_tokens = H // p, W // p + assert ref_H_tokens * ref_W_tokens == ref_img_len + + row_ids = repeat( + torch.arange(ref_H_tokens, dtype=torch.int32, device=device), + "h -> h w", + w=ref_W_tokens, + ).flatten() + col_ids = repeat( + torch.arange(ref_W_tokens, dtype=torch.int32, device=device), + "w -> h w", + h=ref_H_tokens, + ).flatten() + position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 0] = ( + pe_shift + ) + position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 1] = ( + row_ids + ) + position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 2] = ( + col_ids + ) + + pe_shift += max(ref_H_tokens, ref_W_tokens) + pe_shift_len += ref_img_len + + H, W = img_sizes[i] + H_tokens, W_tokens = H // p, W // p + assert H_tokens * W_tokens == l_effective_img_len[i] + + row_ids = repeat( + torch.arange(H_tokens, dtype=torch.int32, device=device), + "h -> h w", + w=W_tokens, + ).flatten() + col_ids = repeat( + torch.arange(W_tokens, dtype=torch.int32, device=device), + "w -> h w", + h=H_tokens, + ).flatten() + + assert pe_shift_len + l_effective_img_len[i] == seq_len + position_ids[i, pe_shift_len:seq_len, 0] = pe_shift + position_ids[i, pe_shift_len:seq_len, 1] = row_ids + position_ids[i, pe_shift_len:seq_len, 2] = col_ids + + # Get combined rotary embeddings + freqs_cis = self._get_freqs_cis(freqs_cis, position_ids) + + # create separate rotary embeddings for captions and images + cap_freqs_cis = torch.zeros( + batch_size, + encoder_seq_len, + freqs_cis.shape[-1], + device=device, + dtype=freqs_cis.dtype, + ) + ref_img_freqs_cis = torch.zeros( + batch_size, + max_ref_img_len, + freqs_cis.shape[-1], + device=device, + dtype=freqs_cis.dtype, + ) + img_freqs_cis = torch.zeros( + batch_size, + max_img_len, + freqs_cis.shape[-1], + device=device, + dtype=freqs_cis.dtype, + ) + + # Calculate combined image sequence lengths (ref_img + img) for each sample + combined_img_seq_lengths = [ + sum(ref_img_len) + img_len + for ref_img_len, img_len in zip( + l_effective_ref_img_len, l_effective_img_len + ) + ] + max_combined_img_len = max(combined_img_seq_lengths) + + # Create combined image rotary embeddings + combined_img_freqs_cis = torch.zeros( + batch_size, + max_combined_img_len, + freqs_cis.shape[-1], + device=device, + dtype=freqs_cis.dtype, + ) + + for i, (cap_seq_len, ref_img_len, img_len, seq_len) in enumerate( + zip( + l_effective_cap_len, + l_effective_ref_img_len, + l_effective_img_len, + seq_lengths, + ) + ): + cap_freqs_cis[i, :cap_seq_len] = freqs_cis[i, :cap_seq_len] + ref_img_freqs_cis[i, : sum(ref_img_len)] = freqs_cis[ + i, cap_seq_len : cap_seq_len + sum(ref_img_len) + ] + img_freqs_cis[i, :img_len] = freqs_cis[ + i, + cap_seq_len + sum(ref_img_len) : cap_seq_len + + sum(ref_img_len) + + img_len, + ] + + # Combined image rotary embeddings: ref_img + img (same order as img_patch_embed_and_refine) + combined_img_freqs_cis[i, : sum(ref_img_len)] = freqs_cis[ + i, cap_seq_len : cap_seq_len + sum(ref_img_len) + ] + combined_img_freqs_cis[i, sum(ref_img_len) : sum(ref_img_len) + img_len] = ( + freqs_cis[ + i, + cap_seq_len + sum(ref_img_len) : cap_seq_len + + sum(ref_img_len) + + img_len, + ] + ) + + return ( + cap_freqs_cis, + ref_img_freqs_cis, + img_freqs_cis, + freqs_cis, + l_effective_cap_len, + seq_lengths, + combined_img_freqs_cis, + combined_img_seq_lengths, + ) diff --git a/extensions_built_in/diffusion_models/boogu_image/src/transformer.py b/extensions_built_in/diffusion_models/boogu_image/src/transformer.py new file mode 100644 index 0000000..ce93406 --- /dev/null +++ b/extensions_built_in/diffusion_models/boogu_image/src/transformer.py @@ -0,0 +1,1174 @@ +"""Boogu-Image transformer (vendored & trimmed for ai-toolkit). + +Adapted from the Boogu-Image repository +(boogu/models/transformers/transformer_boogu.py). Apache-2.0. + +Differences from the upstream file: + * The TeaCache / TaylorSeer inference caches are removed (training/finetuning + never use them), along with the triton RMSNorm and flash-attn fast paths. + * Prompt-tuning (``PromptEmbedding``) is dropped -- the base model does not use it. + * Gradient checkpointing is wired through every heavy block stack (the refiner + loops as well as the double-/single-stream loops) and is gated on + ``torch.is_grad_enabled()`` so it is a no-op during sampling/inference. + +The mixed-stream topology, weight names and numerics are otherwise identical to +upstream so the released checkpoints load unchanged. +""" + +import itertools +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 diffusers.loaders import PeftAdapterMixin +from diffusers.loaders.single_file_model import FromOriginalModelMixin +from diffusers.models.attention_processor import Attention +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils import ( + USE_PEFT_BACKEND, + logging, + scale_lora_layers, + unscale_lora_layers, +) +from einops import rearrange +from torch.nn import RMSNorm + +from .attention_processor import ( + ATTENTION_BACKENDS, + _FLASH_ATTN_AVAILABLE, + BooguImageAttnProcessor, + BooguImageDoubleStreamSelfAttnProcessor, +) +from .block_lumina2 import ( + Lumina2CombinedTimestepCaptionEmbedding, + LuminaFeedForward, + LuminaLayerNormContinuous, + LuminaRMSNormZero, +) +from .rope import BooguImageDoubleStreamRotaryPosEmbed + +logger = logging.get_logger(__name__) + + +class BooguImageTransformerBlock(nn.Module): + """Basic Boogu-Image transformer block: attention + SwiGLU MLP + RMSNorm.""" + + def __init__( + self, + dim: int, + num_attention_heads: int, + num_kv_heads: int, + multiple_of: int, + ffn_dim_multiplier: float, + norm_eps: float, + modulation: bool = True, + ) -> None: + super().__init__() + self.head_dim = dim // num_attention_heads + self.modulation = modulation + + self.attn = Attention( + query_dim=dim, + cross_attention_dim=None, + dim_head=dim // num_attention_heads, + qk_norm="rms_norm", + heads=num_attention_heads, + kv_heads=num_kv_heads, + eps=1e-5, + bias=False, + out_bias=False, + processor=BooguImageAttnProcessor(), + ) + + self.feed_forward = LuminaFeedForward( + dim=dim, + inner_dim=4 * dim, + multiple_of=multiple_of, + ffn_dim_multiplier=ffn_dim_multiplier, + ) + + if modulation: + self.norm1 = LuminaRMSNormZero( + embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True + ) + else: + self.norm1 = RMSNorm(dim, eps=norm_eps) + + self.ffn_norm1 = RMSNorm(dim, eps=norm_eps) + self.norm2 = RMSNorm(dim, eps=norm_eps) + self.ffn_norm2 = RMSNorm(dim, eps=norm_eps) + + self.initialize_weights() + + def initialize_weights(self) -> None: + nn.init.xavier_uniform_(self.attn.to_q.weight) + nn.init.xavier_uniform_(self.attn.to_k.weight) + nn.init.xavier_uniform_(self.attn.to_v.weight) + nn.init.xavier_uniform_(self.attn.to_out[0].weight) + + nn.init.xavier_uniform_(self.feed_forward.linear_1.weight) + nn.init.xavier_uniform_(self.feed_forward.linear_2.weight) + nn.init.xavier_uniform_(self.feed_forward.linear_3.weight) + + if self.modulation: + nn.init.zeros_(self.norm1.linear.weight) + nn.init.zeros_(self.norm1.linear.bias) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + image_rotary_emb: torch.Tensor, + temb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if self.modulation: + if temb is None: + raise ValueError("temb must be provided when modulation is enabled") + norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1( + hidden_states, temb + ) + attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + hidden_states = hidden_states + gate_msa.unsqueeze(1).tanh() * self.norm2( + attn_output + ) + mlp_output = self.feed_forward( + self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1)) + ) + hidden_states = hidden_states + gate_mlp.unsqueeze( + 1 + ).tanh() * self.ffn_norm2(mlp_output) + else: + norm_hidden_states = self.norm1(hidden_states) + attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + hidden_states = hidden_states + self.norm2(attn_output) + mlp_output = self.feed_forward(self.ffn_norm1(hidden_states)) + hidden_states = hidden_states + self.ffn_norm2(mlp_output) + + return hidden_states + + +class BooguImageNoiseRefinerTransformerBlock(BooguImageTransformerBlock): + pass + + +class BooguImageRefImgRefinerTransformerBlock(BooguImageTransformerBlock): + pass + + +class BooguImageContextRefinerTransformerBlock(BooguImageTransformerBlock): + pass + + +class BooguImageSingleStreamTransformerBlock(BooguImageTransformerBlock): + pass + + +class BooguImageDoubleStreamTransformerBlock(nn.Module): + """Boogu-Image double-stream block: instruction & image tokens in parallel streams.""" + + def __init__( + self, + dim: int, + num_attention_heads: int, + num_kv_heads: int, + multiple_of: int, + ffn_dim_multiplier: float, + norm_eps: float, + modulation: bool = True, + ) -> None: + super().__init__() + self.head_dim = dim // num_attention_heads + self.num_attention_heads = num_attention_heads + self.modulation = modulation + self.hidden_size = dim + + double_stream_processor = BooguImageDoubleStreamSelfAttnProcessor( + head_dim=self.head_dim, + num_attention_heads=num_attention_heads, + num_kv_heads=num_kv_heads, + qkv_bias=False, + ) + + # Image stream components. + self.img_instruct_attn = Attention( + query_dim=dim, + cross_attention_dim=None, + dim_head=dim // num_attention_heads, + qk_norm="rms_norm", + heads=num_attention_heads, + kv_heads=num_kv_heads, + eps=1e-5, + bias=False, + out_bias=False, + processor=double_stream_processor, + ) + + self.img_self_attn = Attention( + query_dim=dim, + cross_attention_dim=None, + dim_head=dim // num_attention_heads, + qk_norm="rms_norm", + heads=num_attention_heads, + kv_heads=num_kv_heads, + eps=1e-5, + bias=False, + out_bias=False, + processor=BooguImageAttnProcessor(), + ) + + self.img_feed_forward = LuminaFeedForward( + dim=dim, + inner_dim=4 * dim, + multiple_of=multiple_of, + ffn_dim_multiplier=ffn_dim_multiplier, + ) + + if modulation: + self.img_norm1 = LuminaRMSNormZero( + embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True + ) + self.img_norm2 = LuminaRMSNormZero( + embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True + ) + self.img_norm3 = LuminaRMSNormZero( + embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True + ) + else: + self.img_norm1 = RMSNorm(dim, eps=norm_eps) + self.img_norm2 = RMSNorm(dim, eps=norm_eps) + self.img_norm3 = RMSNorm(dim, eps=norm_eps) + + self.img_ffn_norm1 = RMSNorm(dim, eps=norm_eps) + self.img_attn_norm = RMSNorm(dim, eps=norm_eps) + self.img_self_attn_norm = RMSNorm(dim, eps=norm_eps) + self.img_ffn_norm2 = RMSNorm(dim, eps=norm_eps) + + # Instruction stream components. + self.instruct_feed_forward = LuminaFeedForward( + dim=dim, + inner_dim=4 * dim, + multiple_of=multiple_of, + ffn_dim_multiplier=ffn_dim_multiplier, + ) + + if modulation: + self.instruct_norm1 = LuminaRMSNormZero( + embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True + ) + self.instruct_norm2 = LuminaRMSNormZero( + embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True + ) + else: + self.instruct_norm1 = RMSNorm(dim, eps=norm_eps) + self.instruct_norm2 = RMSNorm(dim, eps=norm_eps) + + self.instruct_ffn_norm1 = RMSNorm(dim, eps=norm_eps) + self.instruct_attn_norm = RMSNorm(dim, eps=norm_eps) + self.instruct_ffn_norm2 = RMSNorm(dim, eps=norm_eps) + + self.initialize_weights() + + # double_stream_processor owns its own q/k/v projections, so the wrapping + # Attention's q/k/v are unused -- drop them so they aren't saved/loaded. + for param in self.img_instruct_attn.to_q.parameters(): + param.requires_grad = False + for param in self.img_instruct_attn.to_k.parameters(): + param.requires_grad = False + for param in self.img_instruct_attn.to_v.parameters(): + param.requires_grad = False + + del self.img_instruct_attn.to_k + del self.img_instruct_attn.to_v + del self.img_instruct_attn.to_q + + def initialize_weights(self) -> None: + nn.init.xavier_uniform_(self.img_instruct_attn.to_out[0].weight) + + nn.init.xavier_uniform_(self.img_self_attn.to_q.weight) + nn.init.xavier_uniform_(self.img_self_attn.to_k.weight) + nn.init.xavier_uniform_(self.img_self_attn.to_v.weight) + nn.init.xavier_uniform_(self.img_self_attn.to_out[0].weight) + + nn.init.xavier_uniform_(self.img_feed_forward.linear_1.weight) + nn.init.xavier_uniform_(self.img_feed_forward.linear_2.weight) + nn.init.xavier_uniform_(self.img_feed_forward.linear_3.weight) + + nn.init.xavier_uniform_(self.instruct_feed_forward.linear_1.weight) + nn.init.xavier_uniform_(self.instruct_feed_forward.linear_2.weight) + nn.init.xavier_uniform_(self.instruct_feed_forward.linear_3.weight) + + if self.modulation: + nn.init.zeros_(self.img_norm1.linear.weight) + nn.init.zeros_(self.img_norm1.linear.bias) + nn.init.zeros_(self.img_norm2.linear.weight) + nn.init.zeros_(self.img_norm2.linear.bias) + nn.init.zeros_(self.img_norm3.linear.weight) + nn.init.zeros_(self.img_norm3.linear.bias) + + nn.init.zeros_(self.instruct_norm1.linear.weight) + nn.init.zeros_(self.instruct_norm1.linear.bias) + nn.init.zeros_(self.instruct_norm2.linear.weight) + nn.init.zeros_(self.instruct_norm2.linear.bias) + + def forward( + self, + img_hidden_states: torch.Tensor, # [B, L_img, D] + instruct_hidden_states: torch.Tensor, # [B, L_instruct, D] + img_attention_mask: torch.Tensor, # [B, L_img] + joint_attention_mask: torch.Tensor, # [B, L_total] + image_rotary_emb: torch.Tensor, # [B, L_img, head_dim] + rotary_emb: torch.Tensor, # [B, L_total, head_dim] + temb: Optional[torch.Tensor] = None, # [B, 1024] + encoder_seq_lengths: List[int] = None, # [B] + seq_lengths: List[int] = None, # [B] + ) -> Tuple[torch.Tensor, torch.Tensor]: + if self.modulation and temb is None: + raise ValueError("temb must be provided when modulation is enabled") + + batch_size = img_hidden_states.shape[0] + L_instruct = instruct_hidden_states.shape[1] + L_img = img_hidden_states.shape[1] + + if self.modulation: + # Step 1: modulation for both streams. + img_norm1_out, img_gate_msa, img_scale_mlp, img_gate_mlp = self.img_norm1( + img_hidden_states, temb + ) + img_norm2_out, img_shift_mlp, _, _ = self.img_norm2(img_hidden_states, temb) + img_norm3_out, img_gate_self, _, _ = self.img_norm3(img_hidden_states, temb) + + ( + instruct_norm1_out, + instruct_gate_msa, + instruct_scale_mlp, + instruct_gate_mlp, + ) = self.instruct_norm1(instruct_hidden_states, temb) + instruct_norm2_out, instruct_shift_mlp, _, _ = self.instruct_norm2( + instruct_hidden_states, temb + ) + + # Step 2: joint attention on [instruct + img]. + joint_attn_out = self.img_instruct_attn.processor( + attn=self.img_instruct_attn, + img_hidden_states=img_norm1_out, + instruct_hidden_states=instruct_norm1_out, + joint_attention_mask=joint_attention_mask, + rotary_emb=rotary_emb, + encoder_seq_lengths=encoder_seq_lengths, + seq_lengths=seq_lengths, + ) + + instruct_attn_out = instruct_hidden_states.new_zeros( + batch_size, L_instruct, self.hidden_size + ) + img_attn_out = img_hidden_states.new_zeros( + batch_size, L_img, self.hidden_size + ) + for i, (encoder_seq_len, seq_len) in enumerate( + zip(encoder_seq_lengths, seq_lengths) + ): + instruct_attn_out[i, :encoder_seq_len] = joint_attn_out[ + i, :encoder_seq_len + ] + img_attn_out[i, : seq_len - encoder_seq_len] = joint_attn_out[ + i, encoder_seq_len:seq_len + ] + + # Step 3: image self-attention. + img_self_attn_out = self.img_self_attn( + hidden_states=img_norm3_out, + encoder_hidden_states=img_norm3_out, + attention_mask=img_attention_mask, + image_rotary_emb=image_rotary_emb, + ) + + # Step 4: residual updates. + img_hidden_states = img_hidden_states + img_gate_msa.unsqueeze( + 1 + ).tanh() * self.img_attn_norm(img_attn_out) + img_hidden_states = img_hidden_states + img_gate_self.unsqueeze( + 1 + ).tanh() * self.img_self_attn_norm(img_self_attn_out) + + img_mlp_input = ( + 1 + img_scale_mlp.unsqueeze(1) + ) * img_norm2_out + img_shift_mlp.unsqueeze(1) + img_mlp_out = self.img_feed_forward(self.img_ffn_norm1(img_mlp_input)) + img_hidden_states = img_hidden_states + img_gate_mlp.unsqueeze( + 1 + ).tanh() * self.img_ffn_norm2(img_mlp_out) + + instruct_hidden_states = ( + instruct_hidden_states + + instruct_gate_msa.unsqueeze(1).tanh() + * self.instruct_attn_norm(instruct_attn_out) + ) + + instruct_mlp_input = ( + 1 + instruct_scale_mlp.unsqueeze(1) + ) * instruct_norm2_out + instruct_shift_mlp.unsqueeze(1) + instruct_mlp_out = self.instruct_feed_forward( + self.instruct_ffn_norm1(instruct_mlp_input) + ) + instruct_hidden_states = ( + instruct_hidden_states + + instruct_gate_mlp.unsqueeze(1).tanh() + * self.instruct_ffn_norm2(instruct_mlp_out) + ) + + else: + # Non-modulated branch used by context-style blocks. + img_norm1_out = self.img_norm1(img_hidden_states) + img_norm3_out = self.img_norm3(img_hidden_states) + instruct_norm1_out = self.instruct_norm1(instruct_hidden_states) + + joint_attn_out = self.img_instruct_attn.processor( + attn=self.img_instruct_attn, + img_hidden_states=img_norm1_out, + instruct_hidden_states=instruct_norm1_out, + joint_attention_mask=joint_attention_mask, + rotary_emb=rotary_emb, + encoder_seq_lengths=encoder_seq_lengths, + seq_lengths=seq_lengths, + ) + + instruct_attn_out = instruct_hidden_states.new_zeros( + batch_size, L_instruct, self.hidden_size + ) + img_attn_out = img_hidden_states.new_zeros( + batch_size, L_img, self.hidden_size + ) + for i, (encoder_seq_len, seq_len) in enumerate( + zip(encoder_seq_lengths, seq_lengths) + ): + instruct_attn_out[i, :encoder_seq_len] = joint_attn_out[ + i, :encoder_seq_len + ] + img_attn_out[i, : seq_len - encoder_seq_len] = joint_attn_out[ + i, encoder_seq_len:seq_len + ] + + img_self_attn_out = self.img_self_attn( + hidden_states=img_norm3_out, + encoder_hidden_states=img_norm3_out, + attention_mask=img_attention_mask, + image_rotary_emb=image_rotary_emb, + ) + + img_hidden_states = img_hidden_states + self.img_attn_norm(img_attn_out) + img_hidden_states = img_hidden_states + self.img_self_attn_norm( + img_self_attn_out + ) + img_norm2_out = self.img_norm2(img_hidden_states) + img_mlp_out = self.img_feed_forward(self.img_ffn_norm1(img_norm2_out)) + img_hidden_states = img_hidden_states + self.img_ffn_norm2(img_mlp_out) + + instruct_hidden_states = instruct_hidden_states + self.instruct_attn_norm( + instruct_attn_out + ) + instruct_norm2_out = self.instruct_norm2(instruct_hidden_states) + instruct_mlp_out = self.instruct_feed_forward( + self.instruct_ffn_norm1(instruct_norm2_out) + ) + instruct_hidden_states = instruct_hidden_states + self.instruct_ffn_norm2( + instruct_mlp_out + ) + + return img_hidden_states, instruct_hidden_states + + +class BooguImageTransformer2DModel( + ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin +): + """Boogu-Image transformer with mixed double-stream -> single-stream topology.""" + + _supports_gradient_checkpointing = True + _no_split_modules = [ + "BooguImageTransformerBlock", + "BooguImageNoiseRefinerTransformerBlock", + "BooguImageRefImgRefinerTransformerBlock", + "BooguImageContextRefinerTransformerBlock", + "BooguImageSingleStreamTransformerBlock", + "BooguImageDoubleStreamTransformerBlock", + ] + _skip_layerwise_casting_patterns = ["x_embedder", "norm", "embedding"] + + @register_to_config + def __init__( + self, + patch_size: int = 2, + in_channels: int = 16, + out_channels: Optional[int] = None, + hidden_size: int = 3360, + num_layers: int = 40, + num_double_stream_layers: int = 8, + num_refiner_layers: int = 2, + num_attention_heads: int = 28, + num_kv_heads: int = 7, + multiple_of: int = 256, + ffn_dim_multiplier: Optional[float] = None, + norm_eps: float = 1e-5, + axes_dim_rope: Tuple[int, int, int] = (40, 40, 40), + axes_lens: Tuple[int, int, int] = (2048, 1664, 1664), + instruction_feature_configs: Dict[str, Any] = dict( + instruction_feat_dim=4096, + reduce_type="mean", + num_instruction_feat_layers=1, + ), + prompt_tuning_configs: Dict[str, Any] = dict(use_prompt_tuning=False), + timestep_scale: float = 1000.0, + ) -> None: + super().__init__() + + if (hidden_size // num_attention_heads) != sum(axes_dim_rope): + raise ValueError( + f"hidden_size // num_attention_heads ({hidden_size // num_attention_heads}) " + f"must equal sum(axes_dim_rope) ({sum(axes_dim_rope)})" + ) + + if num_double_stream_layers > num_layers: + raise ValueError( + f"num_double_stream_layers ({num_double_stream_layers}) cannot be greater " + f"than num_layers ({num_layers})" + ) + + self.out_channels = out_channels or in_channels + self.num_double_stream_layers = num_double_stream_layers + self.num_single_stream_layers = num_layers - num_double_stream_layers + self.instruction_feature_configs = instruction_feature_configs + self.prompt_tuning_configs = prompt_tuning_configs + self.preprocessed_instruction_feat_dim = ( + self.cal_preprocessed_instruction_feat_dim(instruction_feature_configs) + ) + + self.rope_embedder = BooguImageDoubleStreamRotaryPosEmbed( + theta=10000, + axes_dim=axes_dim_rope, + axes_lens=axes_lens, + patch_size=patch_size, + ) + + self.x_embedder = nn.Linear( + in_features=patch_size * patch_size * in_channels, + out_features=hidden_size, + ) + + self.ref_image_patch_embedder = nn.Linear( + in_features=patch_size * patch_size * in_channels, + out_features=hidden_size, + ) + + self.time_caption_embed = Lumina2CombinedTimestepCaptionEmbedding( + hidden_size=hidden_size, + instruction_feat_dim=self.preprocessed_instruction_feat_dim, + norm_eps=norm_eps, + timestep_scale=timestep_scale, + ) + + self.noise_refiner = nn.ModuleList( + [ + BooguImageNoiseRefinerTransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=True, + ) + for _ in range(num_refiner_layers) + ] + ) + + self.ref_image_refiner = nn.ModuleList( + [ + BooguImageRefImgRefinerTransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=True, + ) + for _ in range(num_refiner_layers) + ] + ) + + self.context_refiner = nn.ModuleList( + [ + BooguImageContextRefinerTransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=False, + ) + for _ in range(num_refiner_layers) + ] + ) + + self.double_stream_layers = nn.ModuleList( + [ + BooguImageDoubleStreamTransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=True, + ) + for _ in range(num_double_stream_layers) + ] + ) + + self.single_stream_layers = nn.ModuleList( + [ + BooguImageSingleStreamTransformerBlock( + hidden_size, + num_attention_heads, + num_kv_heads, + multiple_of, + ffn_dim_multiplier, + norm_eps, + modulation=True, + ) + for _ in range(self.num_single_stream_layers) + ] + ) + + self.norm_out = LuminaLayerNormContinuous( + embedding_dim=hidden_size, + conditioning_embedding_dim=min(hidden_size, 1024), + elementwise_affine=False, + eps=1e-6, + bias=True, + out_dim=patch_size * patch_size * self.out_channels, + ) + + # Distinguish multiple reference images (supports up to 5 ref images). + self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size)) + + self.gradient_checkpointing = False + # Attention defaults to torch SDPA ("native"); flip to "flash" via + # set_attention_backend when flash-attn is installed and wanted. + self.attention_backend = "native" + + self.initialize_weights() + + self.layers = list(self.double_stream_layers) + list(self.single_stream_layers) + + def initialize_weights(self) -> None: + nn.init.xavier_uniform_(self.x_embedder.weight) + nn.init.constant_(self.x_embedder.bias, 0.0) + + nn.init.xavier_uniform_(self.ref_image_patch_embedder.weight) + nn.init.constant_(self.ref_image_patch_embedder.bias, 0.0) + + nn.init.zeros_(self.norm_out.linear_1.weight) + nn.init.zeros_(self.norm_out.linear_1.bias) + nn.init.zeros_(self.norm_out.linear_2.weight) + nn.init.zeros_(self.norm_out.linear_2.bias) + + nn.init.normal_(self.image_index_embedding, std=0.02) + + def set_attention_backend(self, backend: str) -> None: + """Select the attention implementation for every attention module. + + Args: + backend: "native" for ``F.scaled_dot_product_attention`` (the default, + no extra dependency) or "flash" for Flash Attention 2 + (``flash_attn_varlen_func``). Selecting "flash" requires the + ``flash_attn`` package to be installed. + """ + backend = backend.lower() + if backend not in ATTENTION_BACKENDS: + raise ValueError( + f"Unknown attention backend {backend!r}. " + f"Expected one of {ATTENTION_BACKENDS}." + ) + if backend == "flash" and not _FLASH_ATTN_AVAILABLE: + raise RuntimeError( + "Flash attention 2 backend requested but the `flash_attn` package " + "is not installed. Install it with `pip install flash-attn` or use " + "the 'native' backend." + ) + self.attention_backend = backend + # Processors live on the wrapping diffusers Attention modules. The single + # -stream processor is a plain object (not an nn.Module) so it isn't in + # self.modules(); reach every processor through its Attention instead. + for module in self.modules(): + if isinstance(module, Attention): + processor = getattr(module, "processor", None) + if hasattr(processor, "attention_backend"): + processor.attention_backend = backend + + def _ckpt(self, layer, *args): + """Run ``layer`` with activation checkpointing when training, else directly.""" + if torch.is_grad_enabled() and self.gradient_checkpointing: + return self._gradient_checkpointing_func(layer, *args) + return layer(*args) + + def img_patch_embed_and_refine( + self, + hidden_states, + ref_image_hidden_states, + padded_img_mask, + padded_ref_img_mask, + noise_rotary_emb, + ref_img_rotary_emb, + l_effective_ref_img_len, + l_effective_img_len, + temb, + ): + """Embed image patches and run the refiner blocks.""" + batch_size = len(hidden_states) + max_combined_img_len = max( + [ + img_len + sum(ref_img_len) + for img_len, ref_img_len in zip( + l_effective_img_len, l_effective_ref_img_len + ) + ] + ) + + hidden_states = self.x_embedder(hidden_states) + ref_image_hidden_states = self.ref_image_patch_embedder(ref_image_hidden_states) + + for i in range(batch_size): + shift = 0 + for j, ref_img_len in enumerate(l_effective_ref_img_len[i]): + ref_image_hidden_states[i, shift : shift + ref_img_len, :] = ( + ref_image_hidden_states[i, shift : shift + ref_img_len, :] + + self.image_index_embedding[j] + ) + shift += ref_img_len + + for layer in self.noise_refiner: + hidden_states = self._ckpt( + layer, hidden_states, padded_img_mask, noise_rotary_emb, temb + ) + + flat_l_effective_ref_img_len = list(itertools.chain(*l_effective_ref_img_len)) + num_ref_images = len(flat_l_effective_ref_img_len) + max_ref_img_len = max(flat_l_effective_ref_img_len) + + batch_ref_img_mask = ref_image_hidden_states.new_zeros( + num_ref_images, max_ref_img_len, dtype=torch.bool + ) + batch_ref_image_hidden_states = ref_image_hidden_states.new_zeros( + num_ref_images, max_ref_img_len, self.config.hidden_size + ) + batch_ref_img_rotary_emb = hidden_states.new_zeros( + num_ref_images, + max_ref_img_len, + ref_img_rotary_emb.shape[-1], + dtype=ref_img_rotary_emb.dtype, + ) + batch_temb = temb.new_zeros(num_ref_images, *temb.shape[1:], dtype=temb.dtype) + + # Flatten reference images into a temporary batch. + idx = 0 + for i in range(batch_size): + shift = 0 + for ref_img_len in l_effective_ref_img_len[i]: + batch_ref_img_mask[idx, :ref_img_len] = True + batch_ref_image_hidden_states[idx, :ref_img_len] = ( + ref_image_hidden_states[i, shift : shift + ref_img_len] + ) + batch_ref_img_rotary_emb[idx, :ref_img_len] = ref_img_rotary_emb[ + i, shift : shift + ref_img_len + ] + batch_temb[idx] = temb[i] + shift += ref_img_len + idx += 1 + + for layer in self.ref_image_refiner: + batch_ref_image_hidden_states = self._ckpt( + layer, + batch_ref_image_hidden_states, + batch_ref_img_mask, + batch_ref_img_rotary_emb, + batch_temb, + ) + + # Restore reference-image sequence layout. + idx = 0 + for i in range(batch_size): + shift = 0 + for ref_img_len in l_effective_ref_img_len[i]: + ref_image_hidden_states[i, shift : shift + ref_img_len] = ( + batch_ref_image_hidden_states[idx, :ref_img_len] + ) + shift += ref_img_len + idx += 1 + + combined_img_hidden_states = hidden_states.new_zeros( + batch_size, max_combined_img_len, self.config.hidden_size + ) + for i, (ref_img_len, img_len) in enumerate( + zip(l_effective_ref_img_len, l_effective_img_len) + ): + combined_img_hidden_states[i, : sum(ref_img_len)] = ref_image_hidden_states[ + i, : sum(ref_img_len) + ] + combined_img_hidden_states[ + i, sum(ref_img_len) : sum(ref_img_len) + img_len + ] = hidden_states[i, :img_len] + + return combined_img_hidden_states + + def flat_and_pad_to_seq(self, hidden_states, ref_image_hidden_states): + """Flatten patch tokens and pad to batched sequences.""" + batch_size = len(hidden_states) + p = self.config.patch_size + device = hidden_states[0].device + + img_sizes = [(img.size(1), img.size(2)) for img in hidden_states] + l_effective_img_len = [(H // p) * (W // p) for (H, W) in img_sizes] + + if ref_image_hidden_states is not None: + ref_img_sizes = [ + [(img.size(1), img.size(2)) for img in imgs] + if imgs is not None + else None + for imgs in ref_image_hidden_states + ] + l_effective_ref_img_len = [ + [ + (ref_img_size[0] // p) * (ref_img_size[1] // p) + for ref_img_size in _ref_img_sizes + ] + if _ref_img_sizes is not None + else [0] + for _ref_img_sizes in ref_img_sizes + ] + else: + ref_img_sizes = [None for _ in range(batch_size)] + l_effective_ref_img_len = [[0] for _ in range(batch_size)] + + max_ref_img_len = max( + [sum(ref_img_len) for ref_img_len in l_effective_ref_img_len] + ) + max_img_len = max(l_effective_img_len) + + # Reference-image patch embeddings. + flat_ref_img_hidden_states = [] + for i in range(batch_size): + if ref_img_sizes[i] is not None: + imgs = [] + for ref_img in ref_image_hidden_states[i]: + C, H, W = ref_img.size() + ref_img = rearrange( + ref_img, "c (h p1) (w p2) -> (h w) (p1 p2 c)", p1=p, p2=p + ) + imgs.append(ref_img) + + img = torch.cat(imgs, dim=0) + flat_ref_img_hidden_states.append(img) + else: + flat_ref_img_hidden_states.append(None) + + # Noise-image patch embeddings. + flat_hidden_states = [] + for i in range(batch_size): + img = hidden_states[i] + C, H, W = img.size() + + img = rearrange(img, "c (h p1) (w p2) -> (h w) (p1 p2 c)", p1=p, p2=p) + flat_hidden_states.append(img) + + padded_ref_img_hidden_states = torch.zeros( + batch_size, + max_ref_img_len, + flat_hidden_states[0].shape[-1], + device=device, + dtype=flat_hidden_states[0].dtype, + ) + padded_ref_img_mask = torch.zeros( + batch_size, max_ref_img_len, dtype=torch.bool, device=device + ) + for i in range(batch_size): + if ref_img_sizes[i] is not None: + padded_ref_img_hidden_states[i, : sum(l_effective_ref_img_len[i])] = ( + flat_ref_img_hidden_states[i] + ) + padded_ref_img_mask[i, : sum(l_effective_ref_img_len[i])] = True + + padded_hidden_states = torch.zeros( + batch_size, + max_img_len, + flat_hidden_states[0].shape[-1], + device=device, + dtype=flat_hidden_states[0].dtype, + ) + padded_img_mask = torch.zeros( + batch_size, max_img_len, dtype=torch.bool, device=device + ) + for i in range(batch_size): + padded_hidden_states[i, : l_effective_img_len[i]] = flat_hidden_states[i] + padded_img_mask[i, : l_effective_img_len[i]] = True + + return ( + padded_hidden_states, + padded_ref_img_hidden_states, + padded_img_mask, + padded_ref_img_mask, + l_effective_ref_img_len, + l_effective_img_len, + ref_img_sizes, + img_sizes, + ) + + def cal_preprocessed_instruction_feat_dim( + self, instruction_feature_configs: Dict[str, Any] + ): + num_instruction_feat_layers = max( + instruction_feature_configs.get("num_instruction_feat_layers", 1), 1 + ) + instruction_feat_dim = instruction_feature_configs.get( + "instruction_feat_dim", 4096 + ) + reduce_type = instruction_feature_configs.get("reduce_type", "concat") + if "cat" in reduce_type.lower(): + return num_instruction_feat_layers * instruction_feat_dim + elif "mean" in reduce_type.lower(): + return instruction_feat_dim + else: + raise ValueError(f"Invalid reduce_type: {reduce_type}") + + def preprocess_instruction_hidden_states( + self, raw_instruction_hidden_states, instruction_feature_configs: Dict[str, Any] + ): + num_instruction_feat_layers = max( + instruction_feature_configs.get("num_instruction_feat_layers", 1), 1 + ) + reduce_type = instruction_feature_configs.get("reduce_type", "concat") + + instruction_hidden_states = None + if isinstance(raw_instruction_hidden_states, torch.Tensor): + instruction_hidden_states = raw_instruction_hidden_states + elif isinstance(raw_instruction_hidden_states, (list, tuple)): + assert len(raw_instruction_hidden_states) == num_instruction_feat_layers + if "cat" in reduce_type.lower(): + instruction_hidden_states = torch.cat( + raw_instruction_hidden_states, dim=-1 + ) + elif "mean" in reduce_type.lower(): + instruction_hidden_states = torch.mean( + torch.stack(raw_instruction_hidden_states), dim=0 + ) + else: + raise ValueError(f"Invalid reduce_type: {reduce_type}") + else: + raise ValueError( + "Invalid type of raw_instruction_hidden_states, expected torch.Tensor " + f"or list, but got {type(raw_instruction_hidden_states)}" + ) + + assert ( + self.preprocessed_instruction_feat_dim + == instruction_hidden_states.shape[-1] + ) + + return instruction_hidden_states + + def forward( + self, + hidden_states: Union[torch.Tensor, List[torch.Tensor]], + timestep: torch.Tensor, + instruction_hidden_states: torch.Tensor, + freqs_cis: torch.Tensor, + instruction_attention_mask: torch.Tensor, + ref_image_hidden_states: Optional[List[List[torch.Tensor]]] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = False, + ) -> Union[torch.Tensor, Transformer2DModelOutput]: + """context/refiner -> double-stream -> fusion -> single-stream -> projection.""" + instruction_hidden_states = self.preprocess_instruction_hidden_states( + instruction_hidden_states, self.instruction_feature_configs + ) + + if attention_kwargs is not None: + attention_kwargs = attention_kwargs.copy() + lora_scale = attention_kwargs.pop("scale", 1.0) + else: + lora_scale = 1.0 + + if USE_PEFT_BACKEND: + scale_lora_layers(self, lora_scale) + elif attention_kwargs is not None and attention_kwargs.get("scale") is not None: + logger.warning( + "Passing `scale` via `attention_kwargs` when not using the PEFT " + "backend is ineffective." + ) + + batch_size = len(hidden_states) + is_hidden_states_tensor = isinstance(hidden_states, torch.Tensor) + + if is_hidden_states_tensor: + assert hidden_states.ndim == 4 + hidden_states = [_hidden_states for _hidden_states in hidden_states] + + device = hidden_states[0].device + + # Timestep and instruction embedding. + temb, instruction_hidden_states = self.time_caption_embed( + timestep, instruction_hidden_states, hidden_states[0].dtype + ) + + # Flatten and pad token sequences. + ( + hidden_states, + ref_image_hidden_states, + img_mask, + ref_img_mask, + l_effective_ref_img_len, + l_effective_img_len, + ref_img_sizes, + img_sizes, + ) = self.flat_and_pad_to_seq(hidden_states, ref_image_hidden_states) + + # Build rotary embeddings and sequence lengths. + ( + context_rotary_emb, + ref_img_rotary_emb, + noise_rotary_emb, + rotary_emb, + encoder_seq_lengths, + seq_lengths, + combined_img_rotary_emb, + combined_img_seq_lengths, + ) = self.rope_embedder( + freqs_cis, + instruction_attention_mask, + l_effective_ref_img_len, + l_effective_img_len, + ref_img_sizes, + img_sizes, + device, + ) + + # Context refinement (non-modulated, so no temb). + for layer in self.context_refiner: + instruction_hidden_states = self._ckpt( + layer, + instruction_hidden_states, + instruction_attention_mask, + context_rotary_emb, + ) + + # Image patch embedding and refinement. + combined_img_hidden_states = self.img_patch_embed_and_refine( + hidden_states, + ref_image_hidden_states, + img_mask, + ref_img_mask, + noise_rotary_emb, + ref_img_rotary_emb, + l_effective_ref_img_len, + l_effective_img_len, + temb, + ) + + instruct_hidden_states = instruction_hidden_states + img_hidden_states = combined_img_hidden_states + + # Joint mask for [instruct + image]. + max_seq_len = max(seq_lengths) + joint_attention_mask = hidden_states.new_zeros( + batch_size, max_seq_len, dtype=torch.bool + ) + for i, seq_len in enumerate(seq_lengths): + joint_attention_mask[i, :seq_len] = True + + # Double-stream stage. + if self.num_double_stream_layers > 0: + max_img_len = max(combined_img_seq_lengths) + img_attention_mask = hidden_states.new_zeros( + batch_size, max_img_len, dtype=torch.bool + ) + for i, img_seq_len in enumerate(combined_img_seq_lengths): + img_attention_mask[i, :img_seq_len] = True + + for layer in self.double_stream_layers: + img_hidden_states, instruct_hidden_states = self._ckpt( + layer, + img_hidden_states, + instruct_hidden_states, + img_attention_mask, + joint_attention_mask, + combined_img_rotary_emb, + rotary_emb, + temb, + encoder_seq_lengths, + seq_lengths, + ) + + # Fuse streams to joint sequence. + joint_hidden_states = hidden_states.new_zeros( + batch_size, max(seq_lengths), self.config.hidden_size + ) + for i, (encoder_seq_len, seq_len) in enumerate( + zip(encoder_seq_lengths, seq_lengths) + ): + joint_hidden_states[i, :encoder_seq_len] = instruct_hidden_states[ + i, :encoder_seq_len + ] + joint_hidden_states[i, encoder_seq_len:seq_len] = img_hidden_states[ + i, : seq_len - encoder_seq_len + ] + + hidden_states = joint_hidden_states + + # Single-stream stage. + for layer in self.single_stream_layers: + hidden_states = self._ckpt( + layer, hidden_states, joint_attention_mask, rotary_emb, temb + ) + + # Output projection. + hidden_states = self.norm_out(hidden_states, temb) + + # Reshape back to image format. + p = self.config.patch_size + output = [] + for i, (img_size, img_len, seq_len) in enumerate( + zip(img_sizes, l_effective_img_len, seq_lengths) + ): + height, width = img_size + img_tokens = hidden_states[i][seq_len - img_len : seq_len] + img_output = rearrange( + img_tokens, + "(h w) (p1 p2 c) -> c (h p1) (w p2)", + h=height // p, + w=width // p, + p1=p, + p2=p, + ) + output.append(img_output) + + if is_hidden_states_tensor: + output = torch.stack(output, dim=0) + + if USE_PEFT_BACKEND: + unscale_lora_layers(self, lora_scale) + + if not return_dict: + return output + return Transformer2DModelOutput(sample=output) diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 6ed0b38..c7eabeb 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -1402,9 +1402,9 @@ def validate_configs( raise ValueError("All datasets must have cache_text_embeddings set to True when caching text embeddings is enabled.") # qwen image edit cannot cache text embeddings - if model_config.arch == 'qwen_image_edit': + if model_config.arch in ['qwen_image_edit', 'boogu_image_edit']: if train_config.unload_text_encoder: - raise ValueError("Cannot cache unload text encoder with qwen_image_edit model. Control images are encoded with text embeddings. You can cache the text embeddings though") + raise ValueError(f"Cannot cache unload text encoder with {model_config.arch} model. Control images are encoded with text embeddings. You can cache the text embeddings though") if train_config.diff_output_preservation and train_config.blank_prompt_preservation: raise ValueError("Cannot use both differential output preservation and blank prompt preservation at the same time. Please set one of them to False.") diff --git a/ui/src/app/jobs/new/options.ts b/ui/src/app/jobs/new/options.ts index ad373a4..0cc5a6c 100644 --- a/ui/src/app/jobs/new/options.ts +++ b/ui/src/app/jobs/new/options.ts @@ -1031,6 +1031,58 @@ export const modelArchs: ModelArch[] = [ 'model.layer_offloading', ], }, + { + name: 'boogu_image', + label: 'Boogu Image', + group: 'image', + defaults: { + 'config.process[0].model.name_or_path': ['Boogu/Boogu-Image-0.1-Base', defaultNameOrPath], + 'config.process[0].model.quantize': [true, false], + 'config.process[0].model.quantize_te': [true, false], + 'config.process[0].train.timestep_type': ['linear', 'sigmoid'], + 'config.process[0].network.conv': [undefined, 16], + 'config.process[0].network.conv_alpha': [undefined, 16], + 'config.process[0].model.low_vram': [true, false], + }, + disableSections: [ + 'network.conv', + ], + additionalSections: [ + 'model.low_vram', + 'model.layer_offloading', + ], + }, + { + name: 'boogu_image_edit', + label: 'Boogu Image Edit', + group: 'instruction', + defaults: { + 'config.process[0].model.name_or_path': ['Boogu/Boogu-Image-0.1-Edit', defaultNameOrPath], + 'config.process[0].model.quantize': [true, false], + 'config.process[0].model.quantize_te': [true, false], + 'config.process[0].train.timestep_type': ['linear', 'sigmoid'], + 'config.process[0].network.conv': [undefined, 16], + 'config.process[0].network.conv_alpha': [undefined, 16], + 'config.process[0].model.low_vram': [true, false], + 'config.process[0].train.unload_text_encoder': [false, false], + 'config.process[0].model.model_kwargs': [ + { + match_target_res: false, + }, + {}, + ], + }, + disableSections: [ + 'network.conv', 'train.unload_text_encoder', + ], + additionalSections: [ + 'datasets.multi_control_paths', + 'sample.multi_ctrl_imgs', + 'model.low_vram', + 'model.layer_offloading', + 'model.qie.match_target_res', + ], + }, ].sort((a, b) => { // Sort by label, case-insensitive return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' }); diff --git a/version.py b/version.py index 9e35d6a..d6c3807 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.10.14" +VERSION = "0.10.15"