What happened to Boogu 2K model?

#5
by lvladikov - opened

When you launched Boogu the notes said that late June there would be a Turbo 2K version released. Than this note was removed. Currently Boogu Turbo mostly duplicates subjects if 2K is attempted (typical issue for models trained on 1024x1024). Even now in your notes you suggest it is more stable at 1K than 2K. Is there plan for proper 2K and 4K versions of Boogu (Base/Turbo)? Those would be amazing addition!

for those like me trying to get the best out of Boogu at 2K+ (for example "Extreme macro of a spiderweb covered in tiny dew droplets, each droplet refracting light, dark green blurred background", by default would become two spiders with 2K), I have found a workaround that doesn't cost much:

NTK-aware RoPE scaling for above-native resolutions. Reduces duplicated subjects at 2K/3K+ by stretching the model's rope theta by s**(d/(d-2)), mapping the larger token grid back onto the coordinate range the model trained on. No-op at or below native resolution (bit-identical output), so it is safe to leave on. Speed cost is ~1.3 ms per generation (unmeasurable).

The whole idea in one line:

theta_ntk = theta * (s ** (d / (d - 2))) # s = current_tokens_per_side / native_tokens_per_side

Portable, dependency-free — this is the shareable part:

import numpy as np

def ntk_theta(theta: float, cur_tokens: int, native_tokens: int, axis_dim: int) -> float:
"""NTK-aware RoPE base for a canvas LARGER than the training resolution.

theta         : the model's own rope_theta (Boogu 10000, Krea2 1000, Flux.2 2000, Z-Image 256)
cur_tokens    : tokens per side you are generating at (px / vae_downsample / patch_size)
native_tokens : tokens per side the model trained at
axis_dim      : per-axis head dim (d) for THIS axis
"""
s = cur_tokens / float(native_tokens)
if s <= 1.0:
    return theta          # at or below native: no-op, output stays bit-identical
return theta * (s ** (axis_dim / (axis_dim - 2)))

def rope_cos_sin(dim, end, theta):
inv = 1.0 / (theta ** (np.arange(0, dim, 2, dtype=np.float64)[: dim // 2] / dim))
ang = np.outer(np.arange(end, dtype=np.float64), inv)
return np.cos(ang), np.sin(ang)

Notes:

  1. Spatial axes only. ax == 0 is the caption/pe_shift axis — it doesn't grow with the canvas, and scaling it corrupts text conditioning for no benefit.
  2. Guard on s <= 1. Below native you're in-distribution; changing θ there perturbs a working operator. The guard makes the flag safe to leave on permanently.
  3. Use the model's own θ and d. Not 10,000. The exponent shifts too: d/(d-2) is 1.0526 at d=40, 1.0667 at d=32, 1.0435 at d=48.
  4. Where to apply it depends on the implementation. If positions are integer indices into precomputed tables (Boogu, Z-Image), you must rebuild the tables. If angles are computed from float positions (Krea2, Flux.2, Ideogram-4), you can set θ on the instance directly.

θ values across the models we surveyed, so people can sanity-check their own:

model θ d (spatial) 2× (2048) 3× (3072)
Boogu 10,000 40 20,743 31,786
Krea2 1,000 48 2,061 3,147
Flux.2-klein 2,000 32 4,189 6,456
Z-Image 256 32 536 826

The above was researched/generated with the help of Claude Opus 5. I hope it helps others.

my dear, i have same problem and i wrote this before and they said put examples .. like as they didn't know what their model doing when we edit on 2k ...
but how to run the solution you wrote here ?are you adding special node ? can you please explain more how to do that fix ?

I am not running this via comfyui so can't tell, I run inference via python code (what comfyui does behind the scenes). So all I can say is find out how to use NTK with ComfyUI if that is how you render (I assume since you mentioned node)

For ComfyUI, I checked again with Claude Opus, and this is what it suggests:

"""
ComfyUI custom node: NTK-aware RoPE scaling for Boogu-Image at above-native resolutions.

Install:  ComfyUI/custom_nodes/comfyui_boogu_ntk_rope.py   (single file, no deps)
Use:      Loader -> [Boogu NTK RoPE (high-res)] -> KSampler
          Set width/height to match your empty latent. Leave native_px at 1024.

WHAT IT FIXES
Boogu was trained at 1024px = 64 latent tokens per side. It learns object scale in TOKEN
units, so on a 2048px canvas (128 tokens/side) there is simply room for a second copy of the
subject -- you get two birds, two spiders, two faces. Measured: at 2048 the same prompt+seed
produces two fused kingfishers, and with this node, one.

HOW
NTK-aware scaling stretches the rope base theta by s**(d/(d-2)), where s = tokens_now /
tokens_native. Because wavelength scales as theta**(d/D), this lengthens mostly the LOW
frequency dims -- the ones actually running off the end of a bigger canvas -- while leaving
the HIGH frequency dims (which separate neighbouring tokens and carry fine detail) as trained.
Boogu: theta 10,000 -> 12,648 @1280, 15,324 @1536, 20,743 @2048.

WHY A PER-AXIS EmbedND IS REQUIRED (this is the part that is easy to get wrong)
ComfyUI's stock EmbedND holds ONE theta shared by all axes:
rope(ids[..., i], self.axes_dim[i], self.theta) # comfy/ldm/flux/layers.py
On Flux.1 you can get away with scaling that single value, because its non-spatial axis
positions are all zero and theta is inert at pos=0. BOOGU IS NOT LIKE THAT. Its axis 0 is
pe_shift: caption tokens get arange(cap_len) and image tokens get the constant cap_len -- all
NONZERO. Scaling a global theta would therefore rescale text conditioning too. This node swaps
in a per-axis EmbedND so ONLY axes 1 (row) and 2 (col) are touched.

Boogu also has axes_dim = (40, 40, 40) -- every axis is the same width -- so you cannot
identify the spatial axes by picking the largest dim, as generic Flux snippets do. The axis
roles must be known, not inferred.

NOTES
* No-op at or below native (s <= 1): returns the model untouched, so it is safe to leave
permanently in a workflow.
* Free: the tables are the same size, so token count, attention and every matmul are
unchanged. Only the numbers inside them differ.
* NOT a complete fix. On a 6-scene probe set it removed subject duplication in 2/2 cases but
ADDED a duplicated word to a shop sign. Check text-heavy scenes.
* Do NOT use linear position interpolation instead (dividing the ids by s). It also removes
duplication, but it makes the 16px token grid visible as a mesh across the whole image --
measured 13.9x excess energy at exactly the token pitch vs 1.0x baseline. NTK avoids this
precisely because it leaves the high-frequency dims alone.

Verified against ComfyUI master: comfy/ldm/boogu/model.py uses OmniGen2RotaryPosEmbed
(comfy/ldm/omnigen/omnigen2.py) with theta=10000, axes_dim_rope=(40,40,40),
axes_lens=(2048,1664,1664), which holds an inner EmbedND at .rope_embedder.
"""

import torch
from torch import Tensor, nn

from comfy.ldm.flux.math import rope

class PerAxisEmbedND(nn.Module):
"""EmbedND with a per-axis theta. Drop-in for comfy.ldm.flux.layers.EmbedND."""

def __init__(self, dim: int, thetas, axes_dim):
    super().__init__()
    self.dim = dim
    self.thetas = [float(t) for t in thetas]
    self.theta = self.thetas[0]        # kept so anything reading .theta still works
    self.axes_dim = list(axes_dim)

def forward(self, ids: Tensor) -> Tensor:
    n_axes = ids.shape[-1]
    emb = torch.cat(
        [rope(ids[..., i], self.axes_dim[i], self.thetas[i]) for i in range(n_axes)],
        dim=-3,
    )
    return emb.unsqueeze(1)

class BooguNTKRoPE:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("MODEL",),
"width": ("INT", {"default": 2048, "min": 256, "max": 8192, "step": 16}),
"height": ("INT", {"default": 2048, "min": 256, "max": 8192, "step": 16}),
"native_px": ("INT", {"default": 1024, "min": 256, "max": 4096, "step": 64}),
}
}

RETURN_TYPES = ("MODEL",)
FUNCTION = "patch"
CATEGORY = "advanced/model"
DESCRIPTION = "NTK-aware RoPE scaling for Boogu-Image above its native 1024px. Reduces duplicated subjects at 2K. No-op at or below native."

def patch(self, model, width, height, native_px):
    m = model.clone()

    # OmniGen2RotaryPosEmbed -> holds the inner EmbedND that actually builds the tables.
    omni = m.get_model_object("diffusion_model.rope_embedder")
    inner = omni.rope_embedder

    base = float(getattr(inner, "theta", 10000.0))
    axes_dim = list(inner.axes_dim)

    # tokens/side = px / 8 (VAE) / 2 (patch) = px/16, so the ratio is just px/native_px.
    s = max(int(width), int(height)) / float(native_px)
    if s <= 1.0:
        print(f"[Boogu NTK] {max(width, height)}px <= native {native_px}px (s={s:.2f}) — no-op.")
        return (m,)

    # Axis 0 = pe_shift/caption (NONZERO on Boogu) -> must keep the original theta.
    # Axes 1/2 = spatial row/col -> NTK-stretched.
    thetas = [base if ax == 0 else base * (s ** (d / (d - 2)))
              for ax, d in enumerate(axes_dim)]

    m.add_object_patch(
        "diffusion_model.rope_embedder.rope_embedder",
        PerAxisEmbedND(inner.dim, thetas, axes_dim),
    )
    print(f"[Boogu NTK] {max(width, height)}px = {max(width, height)//16} tokens/side "
          f"(native {native_px//16}), s={s:.2f}, axes_dim={axes_dim}")
    print(f"[Boogu NTK] theta per axis: {[round(t, 1) for t in thetas]}  "
          f"(axis 0 = caption, unchanged)")
    return (m,)

NODE_CLASS_MAPPINGS = {"BooguNTKRoPE": BooguNTKRoPE}
NODE_DISPLAY_NAME_MAPPINGS = {"BooguNTKRoPE": "Boogu NTK RoPE (high-res)"}

Save as ComfyUI/custom_nodes/comfyui_boogu_ntk_rope.py, restart, then wire Loader → Boogu NTK RoPE → KSampler with width/height matching the empty latent.

θ it produces 1024 → no-op · 1280 → 12,647.7 · 1536 → 15,323.5 · 2048 → 20,743.1, with axis 0 held at 10,000 throughout.

... again, can't verify as I have my own python bespoke inference script, and don't use ComfyUI

Sign up or log in to comment