Add support for Boogu Image and Boogu Image Edit

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,
 ]

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"]

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()
+        }

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"

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

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

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)

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]

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,
+        )

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.")

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' });

version.py

@@ -1 +1 @@
-VERSION = "0.10.14"
+VERSION = "0.10.15"