Skip to content

Commit e25c391

Browse files
authored
feat: Support Boogu-Image (CORE-308) (Comfy-Org#14523)
1 parent ca3dbe2 commit e25c391

9 files changed

Lines changed: 523 additions & 2 deletions

File tree

comfy/ldm/boogu/model.py

Lines changed: 321 additions & 0 deletions
Large diffs are not rendered by default.

comfy/ldm/omnigen/omnigen2.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ def apply_rotary_emb(x, freqs_cis):
2222

2323

2424
def swiglu(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
25-
return F.silu(x) * y
25+
return F.silu(x, inplace=True).mul_(y)
2626

2727

2828
class TimestepEmbedding(nn.Module):

comfy/model_base.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
import comfy.ldm.pixeldit.pid
5555
import comfy.ldm.ace.model
5656
import comfy.ldm.omnigen.omnigen2
57+
import comfy.ldm.boogu.model
5758
import comfy.ldm.qwen_image.model
5859
import comfy.ldm.ideogram4.model
5960
import comfy.ldm.kandinsky5.model
@@ -2103,6 +2104,11 @@ def extra_conds_shapes(self, **kwargs):
21032104
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
21042105
return out
21052106

2107+
class Boogu(Omnigen2):
2108+
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
2109+
super(Omnigen2, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.boogu.model.BooguTransformer2DModel)
2110+
self.memory_usage_factor_conds = ("ref_latents",)
2111+
21062112
class QwenImage(BaseModel):
21072113
def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
21082114
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.qwen_image.model.QwenImageTransformer2DModel)

comfy/model_detection.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -761,6 +761,16 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
761761

762762
return dit_config
763763

764+
if '{}double_stream_layers.0.img_instruct_attn.processor.img_to_q.weight'.format(key_prefix) in state_dict_keys: # Boogu-Image (OmniGen2 derivative + dual-stream stage)
765+
dit_config = {}
766+
dit_config["image_model"] = "boogu"
767+
dit_config["hidden_size"] = state_dict['{}x_embedder.weight'.format(key_prefix)].shape[0]
768+
dit_config["num_layers"] = count_blocks(state_dict_keys, '{}single_stream_layers.'.format(key_prefix) + '{}.')
769+
dit_config["num_double_stream_layers"] = count_blocks(state_dict_keys, '{}double_stream_layers.'.format(key_prefix) + '{}.')
770+
dit_config["num_refiner_layers"] = count_blocks(state_dict_keys, '{}noise_refiner.'.format(key_prefix) + '{}.')
771+
dit_config["instruction_feat_dim"] = state_dict['{}time_caption_embed.caption_embedder.0.weight'.format(key_prefix)].shape[0]
772+
return dit_config
773+
764774
if '{}time_caption_embed.timestep_embedder.linear_1.bias'.format(key_prefix) in state_dict_keys: # Omnigen2
765775
dit_config = {}
766776
dit_config["image_model"] = "omnigen2"

comfy/sd.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@
6868
import comfy.text_encoders.longcat_image
6969
import comfy.text_encoders.qwen35
7070
import comfy.text_encoders.qwen3vl
71+
import comfy.text_encoders.boogu
7172
import comfy.text_encoders.ernie
7273
import comfy.text_encoders.gemma4
7374
import comfy.text_encoders.cogvideo
@@ -1301,6 +1302,7 @@ class CLIPType(Enum):
13011302
LENS = 28
13021303
PIXELDIT = 29
13031304
IDEOGRAM4 = 30
1305+
BOOGU = 31
13041306

13051307

13061308

@@ -1622,6 +1624,10 @@ class EmptyClass:
16221624
clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
16231625
clip_target.clip = comfy.text_encoders.ideogram4.te_qwen3vl(**llama_detect(clip_data))
16241626
clip_target.tokenizer = comfy.text_encoders.ideogram4.Ideogram4Qwen3VLTokenizer
1627+
elif clip_type == CLIPType.BOOGU and te_model == TEModel.QWEN3VL_8B: # Boogu-Image: full Qwen3-VL-8B, last hidden state, no-think template.
1628+
clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
1629+
clip_target.clip = comfy.text_encoders.boogu.te(**llama_detect(clip_data))
1630+
clip_target.tokenizer = comfy.text_encoders.boogu.BooguTokenizer
16251631
elif clip_type in (CLIPType.FLUX, CLIPType.FLUX2): # Flux2 Klein reuses the Qwen3-VL LM (3-layer tap -> 12288); visual unused.
16261632
klein_model_type = "qwen3_8b" if te_model == TEModel.QWEN3VL_8B else "qwen3_4b"
16271633
clip_target.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type=klein_model_type)

comfy/supported_models.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import comfy.text_encoders.kandinsky5
2626
import comfy.text_encoders.z_image
2727
import comfy.text_encoders.ideogram4
28+
import comfy.text_encoders.boogu
2829
import comfy.text_encoders.anima
2930
import comfy.text_encoders.ace15
3031
import comfy.text_encoders.longcat_image
@@ -1758,6 +1759,27 @@ def clip_target(self, state_dict={}):
17581759
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_3b.transformer.".format(pref))
17591760
return supported_models_base.ClipTarget(comfy.text_encoders.omnigen2.Omnigen2Tokenizer, comfy.text_encoders.omnigen2.te(**hunyuan_detect))
17601761

1762+
class Boogu(Omnigen2):
1763+
unet_config = {
1764+
"image_model": "boogu",
1765+
}
1766+
1767+
sampling_settings = {
1768+
"multiplier": 1.0,
1769+
"shift": 3.16,
1770+
}
1771+
1772+
memory_usage_factor = 1.95 #TODO
1773+
1774+
def get_model(self, state_dict, prefix="", device=None):
1775+
out = model_base.Boogu(self, device=device)
1776+
return out
1777+
1778+
def clip_target(self, state_dict={}):
1779+
pref = self.text_encoder_key_prefix[0]
1780+
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl_8b.transformer.".format(pref))
1781+
return supported_models_base.ClipTarget(comfy.text_encoders.boogu.BooguTokenizer, comfy.text_encoders.boogu.te(**hunyuan_detect))
1782+
17611783
class Ideogram4(supported_models_base.BASE):
17621784
unet_config = {
17631785
"image_model": "ideogram4",
@@ -2300,6 +2322,7 @@ def get_model(self, state_dict, prefix="", device=None):
23002322
ACEStep,
23012323
ACEStep15,
23022324
Omnigen2,
2325+
Boogu,
23032326
QwenImage,
23042327
Ideogram4,
23052328
Flux2,

comfy/text_encoders/boogu.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
"""Boogu-Image text encoder: full Qwen3-VL-8B, last hidden state (4096-dim).
2+
3+
Boogu uses the final hidden state of Qwen3-VL as the per-token instruction feature
4+
(num_instruction_feature_layers=1, reduce_type=mean -> just the last layer).
5+
The model itself is the standard Qwen3-VL TE, only the chat template differs
6+
(a fixed system prompt and no <think> block).
7+
"""
8+
9+
import comfy.text_encoders.qwen3vl
10+
from comfy import sd1_clip
11+
12+
13+
# System prompts from the reference pipeline (pipeline_boogu.py).
14+
# T2I (non-empty instruction, no image) uses the helpful-assistant prompt
15+
# everything else (the CFG negative / "drop" condition, and any image case) uses the TI2I "describe" prompt.
16+
BOOGU_T2I_SYSTEM = "You are a helpful assistant that generates high-quality images based on user instructions. The instructions are as follows."
17+
BOOGU_DROP_SYSTEM = "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."
18+
19+
20+
class BooguTokenizer(comfy.text_encoders.qwen3vl.Qwen3VLTokenizer):
21+
def __init__(self, embedding_directory=None, tokenizer_data={}):
22+
super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, model_type="qwen3vl_8b")
23+
# apply_chat_template without add_generation_prompt
24+
self.llama_template = "<|im_start|>system\n" + BOOGU_T2I_SYSTEM + "<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n"
25+
self.llama_template_images = "<|im_start|>system\n" + BOOGU_DROP_SYSTEM + "<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n"
26+
# Reference SYSTEM_PROMPT_DROP: used for the empty negative/uncond instruction.
27+
self.llama_template_drop = "<|im_start|>system\n" + BOOGU_DROP_SYSTEM + "<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n"
28+
29+
def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], prevent_empty_text=False, thinking=True, **kwargs):
30+
if llama_template is None and len(images) == 0 and text.strip() == "":
31+
llama_template = self.llama_template_drop
32+
# Boogu conditions on the no-think template; thinking=True drops the empty <think> block qwen3vl adds by default.
33+
return super().tokenize_with_weights(text, return_word_ids=return_word_ids, llama_template=llama_template, images=images, prevent_empty_text=prevent_empty_text, thinking=thinking, **kwargs)
34+
35+
36+
class BooguQwen3VLClipModel(comfy.text_encoders.qwen3vl.Qwen3VLClipModel):
37+
def __init__(self, device="cpu", dtype=None, attention_mask=True, model_options={}, model_type="qwen3vl_8b"):
38+
super().__init__(device=device, dtype=dtype, attention_mask=attention_mask, model_options=model_options, model_type=model_type)
39+
# apply the final RMSNorm to the tapped last layer
40+
self.layer_norm_hidden_state = True
41+
42+
43+
class BooguTEModel(sd1_clip.SD1ClipModel):
44+
def __init__(self, device="cpu", dtype=None, model_options={}):
45+
clip_model = lambda **kw: BooguQwen3VLClipModel(**kw, model_type="qwen3vl_8b")
46+
super().__init__(device=device, dtype=dtype, name="qwen3vl_8b", clip_model=clip_model, model_options=model_options)
47+
48+
49+
def te(dtype_llama=None, llama_quantization_metadata=None):
50+
class BooguTEModel_(BooguTEModel):
51+
def __init__(self, device="cpu", dtype=None, model_options={}):
52+
if dtype_llama is not None:
53+
dtype = dtype_llama
54+
if llama_quantization_metadata is not None:
55+
model_options = model_options.copy()
56+
model_options["quantization_metadata"] = llama_quantization_metadata
57+
super().__init__(device=device, dtype=dtype, model_options=model_options)
58+
return BooguTEModel_

comfy_extras/nodes_boogu.py

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
import math
2+
3+
import node_helpers
4+
import comfy.utils
5+
from typing_extensions import override
6+
from comfy_api.latest import ComfyExtension, io
7+
8+
9+
class TextEncodeBooguEdit(io.ComfyNode):
10+
"""Boogu-Image Edit conditioning.
11+
12+
The edit image is used twice, matching the reference pipeline:
13+
- Qwen3-VL vision tokens (instruction understanding) -> positive only
14+
- VAE reference latent (image identity) -> positive and negative
15+
The ref latent is in both conds so it cancels under CFG (identity preserved);
16+
the vision tokens are only in the positive so CFG amplifies the instruction.
17+
The tokenizer selects the right system prompt automatically (image -> TI2I,
18+
empty negative -> DROP), so no template plumbing is needed here.
19+
"""
20+
21+
@classmethod
22+
def define_schema(cls):
23+
return io.Schema(
24+
node_id="TextEncodeBooguEdit",
25+
category="model/conditioning/boogu",
26+
inputs=[
27+
io.Clip.Input("clip"),
28+
io.String.Input("prompt", multiline=True, dynamic_prompts=True),
29+
io.Vae.Input("vae"),
30+
io.Autogrow.Input(
31+
"images",
32+
template=io.Autogrow.TemplateNames(
33+
io.Image.Input("image"),
34+
names=[f"image_{i}" for i in range(1, 17)],
35+
min=1,
36+
),
37+
tooltip="Reference image(s) to edit. Boogu focuses on one reference per sample; more are allowed.",
38+
),
39+
],
40+
outputs=[
41+
io.Conditioning.Output(display_name="positive"),
42+
io.Conditioning.Output(display_name="negative"),
43+
],
44+
)
45+
46+
@classmethod
47+
def execute(cls, clip, prompt, vae=None, images: io.Autogrow.Type = None) -> io.NodeOutput:
48+
ref_latents = []
49+
images_vl = []
50+
51+
images = images or {}
52+
for name in sorted(images, key=lambda n: int(n.rsplit("_", 1)[-1])):
53+
image = images[name]
54+
if image is None:
55+
continue
56+
samples = image.movedim(-1, 1)
57+
58+
# Vision tower input: the reference caps the VLM image at 384x384
59+
# (max_vlm_input_pil_pixels in pipeline_boogu.py).
60+
total = int(384 * 384)
61+
scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
62+
width = round(samples.shape[3] * scale_by)
63+
height = round(samples.shape[2] * scale_by)
64+
s = comfy.utils.common_upscale(samples, width, height, "area", "disabled")
65+
images_vl.append(s.movedim(1, -1)[:, :, :, :3])
66+
67+
# Reference latent: align to 16 px (VAE /8 * patch_size 2).
68+
if vae is not None:
69+
total = int(1024 * 1024)
70+
scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
71+
width = round(samples.shape[3] * scale_by / 16.0) * 16
72+
height = round(samples.shape[2] * scale_by / 16.0) * 16
73+
s = comfy.utils.common_upscale(samples, width, height, "area", "disabled")
74+
ref_latents.append(vae.encode(s.movedim(1, -1)[:, :, :, :3]))
75+
76+
# positive: instruction + vision tokens; negative: empty (no vision). Ref latent on both.
77+
positive = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=images_vl))
78+
negative = clip.encode_from_tokens_scheduled(clip.tokenize(""))
79+
80+
if len(ref_latents) > 0:
81+
positive = node_helpers.conditioning_set_values(positive, {"reference_latents": ref_latents}, append=True)
82+
negative = node_helpers.conditioning_set_values(negative, {"reference_latents": ref_latents}, append=True)
83+
84+
return io.NodeOutput(positive, negative)
85+
86+
87+
class BooguExtension(ComfyExtension):
88+
@override
89+
async def get_node_list(self) -> list[type[io.ComfyNode]]:
90+
return [
91+
TextEncodeBooguEdit,
92+
]
93+
94+
95+
async def comfy_entrypoint() -> BooguExtension:
96+
return BooguExtension()

nodes.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -969,7 +969,7 @@ class CLIPLoader:
969969
@classmethod
970970
def INPUT_TYPES(s):
971971
return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
972-
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4"], ),
972+
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu"], ),
973973
},
974974
"optional": {
975975
"device": (["default", "cpu"], {"advanced": True}),
@@ -2425,6 +2425,7 @@ async def init_builtin_extra_nodes():
24252425
"nodes_tcfg.py",
24262426
"nodes_context_windows.py",
24272427
"nodes_qwen.py",
2428+
"nodes_boogu.py",
24282429
"nodes_chroma_radiance.py",
24292430
"nodes_pid.py",
24302431
"nodes_model_patch.py",

0 commit comments

Comments
 (0)