Merge branch 'master' into matt/disambiguate-deprecated-partner-nodes

comfy/ldm/krea2/model.py

@@ -15,6 +15,7 @@ from einops import rearrange
 import comfy.model_management
 import comfy.patcher_extension
 import comfy.ldm.common_dit
+import comfy.utils
 from comfy.ldm.flux.layers import EmbedND, timestep_embedding
 from comfy.ldm.flux.math import apply_rope
 from comfy.ldm.modules.attention import optimized_attention_masked
@@ -73,11 +74,20 @@ class Attention(nn.Module):
         self.wo = operations.Linear(dim, dim, bias=bias, device=device, dtype=dtype)
 
     def forward(self, x, freqs=None, mask=None, transformer_options={}):
+        transformer_patches = transformer_options.get("patches", {})
+        extra_options = transformer_options.copy()
         q, k, v, gate = self.wq(x), self.wk(x), self.wv(x), self.gate(x)
         q = rearrange(q, "B L (H D) -> B H L D", H=self.heads)
         k = rearrange(k, "B L (H D) -> B H L D", H=self.kvheads)
         v = rearrange(v, "B L (H D) -> B H L D", H=self.kvheads)
         q, k = self.qknorm(q, k)
+
+        if "block_index" in transformer_options and "attn1_patch" in transformer_patches:
+            for p in transformer_patches["attn1_patch"]:
+                out = p(q, k, v, pe=freqs, attn_mask=mask, extra_options=extra_options)
+                q, k, v = out.get("q", q), out.get("k", k), out.get("v", v)
+                freqs, mask = out.get("pe", freqs), out.get("attn_mask", mask)
+
         if freqs is not None:
             q, k = apply_rope(q, k, freqs)
         if self.kvheads != self.heads:
@@ -86,6 +96,11 @@ class Attention(nn.Module):
             v = v.repeat_interleave(rep, dim=1)
         out = optimized_attention_masked(q, k, v, self.heads, mask=mask, skip_reshape=True,
                                          transformer_options=transformer_options)
+
+        if "block_index" in transformer_options and "attn1_output_patch" in transformer_patches:
+            for p in transformer_patches["attn1_output_patch"]:
+                out = p(out, extra_options)
+
         return self.wo(out * F.sigmoid(gate))
 
 
@@ -158,8 +173,44 @@ class SingleStreamBlock(nn.Module):
         self.attn = Attention(features, heads, kvheads=kvheads, bias=bias, device=device, dtype=dtype, operations=operations)
         self.mlp = SwiGLU(features, multiplier, bias, device=device, dtype=dtype, operations=operations)
 
-    def forward(self, x, vec, freqs, mask=None, transformer_options={}):
+    def forward(self, x, vec, freqs, mask=None, timestep_zero_index=None, transformer_options={}):
         prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec)
+        if timestep_zero_index is not None:
+            bs = x.shape[0]
+            ref_prescale = prescale[bs:]
+            ref_preshift = preshift[bs:]
+            ref_pregate = pregate[bs:]
+            ref_postscale = postscale[bs:]
+            ref_postshift = postshift[bs:]
+            ref_postgate = postgate[bs:]
+            prescale = prescale[:bs]
+            preshift = preshift[:bs]
+            pregate = pregate[:bs]
+            postscale = postscale[:bs]
+            postshift = postshift[:bs]
+            postgate = postgate[:bs]
+
+            pre = self.prenorm(x)
+            pre[:, :timestep_zero_index].mul_(1 + prescale).add_(preshift)
+            pre[:, timestep_zero_index:].mul_(1 + ref_prescale).add_(ref_preshift)
+            attn = self.attn(pre, freqs, mask, transformer_options=transformer_options)
+            del pre
+            attn[:, :timestep_zero_index].mul_(pregate)
+            attn[:, timestep_zero_index:].mul_(ref_pregate)
+            x = x + attn
+            del attn
+
+            post = self.postnorm(x)
+            post[:, :timestep_zero_index].mul_(1 + postscale).add_(postshift)
+            post[:, timestep_zero_index:].mul_(1 + ref_postscale).add_(ref_postshift)
+            mlp = self.mlp(post)
+            del post
+            mlp[:, :timestep_zero_index].mul_(postgate)
+            mlp[:, timestep_zero_index:].mul_(ref_postgate)
+            x = x + mlp
+            del mlp
+            return x
+
         x = x + pregate * self.attn((1 + prescale) * self.prenorm(x) + preshift, freqs, mask, transformer_options=transformer_options)
         x = x + postgate * self.mlp((1 + postscale) * self.postnorm(x) + postshift)
         return x
@@ -181,7 +232,7 @@ class LastLayer(nn.Module):
 class SingleStreamDiT(nn.Module):
     def __init__(self, features=6144, tdim=256, txtdim=2560, heads=48, kvheads=12, multiplier=4,
                  layers=28, patch=2, channels=16, bias=False, theta=1e3, txtlayers=12,
-                 txtheads=20, txtkvheads=20, image_model=None,
+                 txtheads=20, txtkvheads=20, default_ref_method=None, image_model=None,
                  device=None, dtype=None, operations=None, **kwargs):
         super().__init__()
         self.dtype = dtype
@@ -191,6 +242,7 @@ class SingleStreamDiT(nn.Module):
         self.heads = heads
         self.txtdim = txtdim
         self.txtlayers = txtlayers
+        self.default_ref_method = default_ref_method
 
         headdim = features // heads
         axes = [headdim - 12 * (headdim // 16), 6 * (headdim // 16), 6 * (headdim // 16)]
@@ -221,61 +273,110 @@ class SingleStreamDiT(nn.Module):
             operations.Linear(features, features * 6, device=device, dtype=dtype),
         )
 
-    def forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs):
+    def forward(self, x, timesteps, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs):
         return comfy.patcher_extension.WrapperExecutor.new_class_executor(
             self._forward,
             self,
             comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options),
-        ).execute(x, timesteps, context, attention_mask, transformer_options, **kwargs)
+        ).execute(x, timesteps, context, attention_mask, ref_latents, transformer_options, **kwargs)
 
-    def _forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs):
+    def process_img(self, x, index=0):
+        patch = self.patch
+        x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch))
+        h, w = x.shape[-2] // patch, x.shape[-1] // patch
+        img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch)
+
+        img_ids = torch.zeros(h, w, 3, device=x.device, dtype=torch.float32)
+        img_ids[..., 0] = index
+        img_ids[..., 1] = torch.arange(h, device=x.device, dtype=torch.float32)[:, None]
+        img_ids[..., 2] = torch.arange(w, device=x.device, dtype=torch.float32)[None, :]
+        return img, img_ids.reshape(1, h * w, 3).repeat(x.shape[0], 1, 1), h, w
+
+    def _forward(self, x, timesteps, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs):
+        transformer_options = transformer_options.copy()
         temporal = x.ndim == 5
         if temporal:
             b5, c5, t5, h5, w5 = x.shape
             x = x.reshape(b5 * t5, c5, h5, w5)
-        bs, c, H_orig, W_orig = x.shape
+        bs, _, h_orig, w_orig = x.shape
         patch = self.patch
-        # Pad the latent up to a multiple of patch (as Flux/Lumina/QwenImage do); crop back at the end.
-        x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch))
-        H, W = x.shape[-2], x.shape[-1]
-        h_, w_ = H // patch, W // patch
 
         # context arrives as (B, seq, txtlayers*txtdim); reshape to (B, txtlayers, seq, txtdim).
         context = self._unpack_context(context)
 
-        img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch)
+        img, imgpos, h_, w_ = self.process_img(x)
+        img_tokens = img.shape[1]
+        timestep_zero_index = None
+        ref_method = kwargs.get("ref_latents_method", self.default_ref_method)
+        if ref_method is not None and ref_latents is not None and len(ref_latents) > 0:
+            ref_tokens = []
+            ref_pos = []
+            ref_num_tokens = []
+            for index, ref in enumerate(ref_latents, 1):
+                if ref.ndim == 5:
+                    rb, rc, rt, rh5, rw5 = ref.shape
+                    ref = ref.reshape(rb * rt, rc, rh5, rw5)
+                ref = comfy.utils.repeat_to_batch_size(ref, bs)
+                kontext, kontext_ids, _, _ = self.process_img(ref, index=index)
+                ref_tokens.append(kontext)
+                ref_pos.append(kontext_ids)
+                ref_num_tokens.append(kontext.shape[1])
+            img = torch.cat([img] + ref_tokens, dim=1)
+            imgpos = torch.cat([imgpos] + ref_pos, dim=1)
+            del ref_tokens, ref_pos
+            if ref_method == "index_timestep_zero":
+                timestep_zero_index = img_tokens
+            transformer_options["reference_image_num_tokens"] = ref_num_tokens
+
         img = self.first(img)
 
         t = self.tmlp(timestep_embedding(timesteps, self.tdim).unsqueeze(1).to(img.dtype))
         tvec = self.tproj(t)
+        if timestep_zero_index is not None:
+            t0 = self.tmlp(timestep_embedding(torch.zeros_like(timesteps), self.tdim).unsqueeze(1).to(img.dtype))
+            tvec = torch.cat((tvec, self.tproj(t0)), dim=0)
 
         context = self.txtfusion(context, mask=None, transformer_options=transformer_options)
         context = self.txtmlp(context)
 
-        txtlen, imglen = context.shape[1], img.shape[1]
+        txtlen = context.shape[1]
+        device = context.device
+        txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)
+
+        patches = transformer_options.get("patches", {})
+        if "post_input" in patches:
+            for p in patches["post_input"]:
+                out = p({"img": img, "txt": context, "img_ids": imgpos, "txt_ids": txtpos, "transformer_options": transformer_options})
+                img, context = out["img"], out["txt"]
+                imgpos, txtpos = out["img_ids"], out["txt_ids"]
+
         combined = torch.cat((context, img), dim=1)
+        del context, img
+        if timestep_zero_index is not None:
+            timestep_zero_index += txtlen
 
         # Position ids: text at 0, image at (0, h_idx, w_idx).
-        device = combined.device
-        txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)
-        imgids = torch.zeros(h_, w_, 3, device=device, dtype=torch.float32)
-        imgids[..., 1] = torch.arange(h_, device=device, dtype=torch.float32)[:, None]
-        imgids[..., 2] = torch.arange(w_, device=device, dtype=torch.float32)[None, :]
-        imgpos = imgids.reshape(1, h_ * w_, 3).repeat(bs, 1, 1)
         pos = torch.cat((txtpos, imgpos), dim=1)
+        del txtpos, imgpos
 
         freqs = self.pe_embedder(pos)
+        del pos
 
-        for block in self.blocks:
-            combined = block(combined, tvec, freqs, None, transformer_options=transformer_options)
+        transformer_options["total_blocks"] = len(self.blocks)
+        transformer_options["block_type"] = "single"
+        transformer_options["img_slice"] = [txtlen, combined.shape[1]]
+        for i, block in enumerate(self.blocks):
+            transformer_options["block_index"] = i
+            combined = block(combined, tvec, freqs, None, timestep_zero_index=timestep_zero_index, transformer_options=transformer_options)
 
         final = self.last(combined, t)
-        out = final[:, txtlen:txtlen + imglen, :]
+        del combined
+        out = final[:, txtlen:txtlen + img_tokens, :]
         out = rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)",
                         h=h_, w=w_, ph=patch, pw=patch, c=self.channels)
-        out = out[:, :, :H_orig, :W_orig]  # crop padding back off
+        out = out[:, :, :h_orig, :w_orig]  # crop padding back off
         if temporal:
-            out = out.reshape(b5, t5, self.channels, H_orig, W_orig).movedim(1, 2)
+            out = out.reshape(b5, t5, self.channels, h_orig, w_orig).movedim(1, 2)
         return out
 
     def _unpack_context(self, context):

comfy/model_base.py

@@ -2024,11 +2024,11 @@ class WAN22_WanDancer(WAN21):
 
         fps = kwargs.get("fps", None)
         if fps is not None:
-            out['fps'] = comfy.conds.CONDRegular(torch.FloatTensor([fps]))
+            out['fps'] = comfy.conds.CONDConstant(fps)
 
         audio_inject_scale = kwargs.get("audio_inject_scale", None)
         if audio_inject_scale is not None:
-            out['audio_inject_scale'] = comfy.conds.CONDRegular(torch.FloatTensor([audio_inject_scale]))
+            out['audio_inject_scale'] = comfy.conds.CONDConstant(audio_inject_scale)
         return out
 
 class Hunyuan3Dv2(BaseModel):
@@ -2227,10 +2227,7 @@ class Omnigen2(BaseModel):
             out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
         ref_latents = kwargs.get("reference_latents", None)
         if ref_latents is not None:
-            latents = []
-            for lat in ref_latents:
-                latents.append(self.process_latent_in(lat))
-            out['ref_latents'] = comfy.conds.CONDList(latents)
+            out['ref_latents'] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in ref_latents])
         return out
 
     def extra_conds_shapes(self, **kwargs):
@@ -2317,12 +2314,30 @@ class Ideogram4(BaseModel):
 class Krea2(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.krea2.model.SingleStreamDiT)
+        self.memory_usage_factor_conds = ("ref_latents",)
 
     def extra_conds(self, **kwargs):
         out = super().extra_conds(**kwargs)
         cross_attn = kwargs.get("cross_attn", None)
         if cross_attn is not None:
             out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
+        ref_latents = kwargs.get("reference_latents", None)
+        if ref_latents is not None:
+            latents = []
+            for lat in ref_latents:
+                latents.append(self.process_latent_in(lat))
+            out['ref_latents'] = comfy.conds.CONDList(latents)
+
+            ref_latents_method = kwargs.get("reference_latents_method", None)
+            if ref_latents_method is not None:
+                out['ref_latents_method'] = comfy.conds.CONDConstant(ref_latents_method)
+        return out
+
+    def extra_conds_shapes(self, **kwargs):
+        out = {}
+        ref_latents = kwargs.get("reference_latents", None)
+        if ref_latents is not None:
+            out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
         return out
 
 class HunyuanImage21(BaseModel):

comfy/sd.py

@@ -1434,6 +1434,7 @@ class TEModel(Enum):
     GPT_OSS_20B = 33
     QWEN3VL_4B = 34
     QWEN3VL_8B = 35
+    GEMMA_4_12B = 36
 
 
 def detect_te_model(sd):
@@ -1463,6 +1464,9 @@ def detect_te_model(sd):
     if 'model.layers.0.post_feedforward_layernorm.weight' in sd:
         if 'model.layers.59.self_attn.q_norm.weight' in sd:
             return TEModel.GEMMA_4_31B
+        # Gemma4 12B Unified: 48 layers, encoder-free; global layers drop v_proj (attention_k_eq_v).
+        if 'model.layers.47.self_attn.q_norm.weight' in sd and 'model.layers.5.self_attn.v_proj.weight' not in sd:
+            return TEModel.GEMMA_4_12B
         if 'model.layers.41.self_attn.q_norm.weight' in sd and 'model.layers.47.self_attn.q_norm.weight' not in sd:
             return TEModel.GEMMA_4_E4B
         if 'model.layers.34.self_attn.q_norm.weight' in sd and 'model.layers.41.self_attn.q_norm.weight' not in sd:
@@ -1618,10 +1622,11 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
             clip_target.clip = comfy.text_encoders.sa3.SAT5GemmaModel
             clip_target.tokenizer = comfy.text_encoders.sa3.SAT5GemmaTokenizer
             tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
-        elif te_model in (TEModel.GEMMA_4_E4B, TEModel.GEMMA_4_E2B, TEModel.GEMMA_4_31B):
+        elif te_model in (TEModel.GEMMA_4_E4B, TEModel.GEMMA_4_E2B, TEModel.GEMMA_4_31B, TEModel.GEMMA_4_12B):
             variant = {TEModel.GEMMA_4_E4B: comfy.text_encoders.gemma4.Gemma4_E4B,
                        TEModel.GEMMA_4_E2B: comfy.text_encoders.gemma4.Gemma4_E2B,
-                       TEModel.GEMMA_4_31B: comfy.text_encoders.gemma4.Gemma4_31B}[te_model]
+                       TEModel.GEMMA_4_31B: comfy.text_encoders.gemma4.Gemma4_31B,
+                       TEModel.GEMMA_4_12B: comfy.text_encoders.gemma4.Gemma4_12B}[te_model]
             clip_target.clip = comfy.text_encoders.gemma4.gemma4_te(**llama_detect(clip_data), model_class=variant)
             clip_target.tokenizer = variant.tokenizer
             tokenizer_data["tokenizer_json"] = clip_data[0].get("tokenizer_json", None)

comfy/text_encoders/gemma4.py

@@ -1,11 +1,15 @@
 import torch
 import torch.nn as nn
+import torchaudio.functional as AF
+import torchvision.transforms.functional as TVF
 import numpy as np
+from tokenizers import Tokenizer
 from dataclasses import dataclass
 import math
 
 from comfy import sd1_clip
 import comfy.model_management
+import comfy.ops
 from comfy.ldm.modules.attention import optimized_attention_for_device
 from comfy.rmsnorm import rms_norm
 from comfy.text_encoders.llama import RMSNorm, MLP, BaseLlama, BaseGenerate, _make_scaled_embedding
@@ -21,6 +25,10 @@ GEMMA4_VISION_CONFIG = {"hidden_size": 768, "image_size": 896, "intermediate_siz
 GEMMA4_VISION_31B_CONFIG = {"hidden_size": 1152, "image_size": 896, "intermediate_size": 4304, "num_attention_heads": 16, "num_hidden_layers": 27, "patch_size": 16, "head_dim": 72, "rms_norm_eps": 1e-6, "position_embedding_size": 10240, "pooling_kernel_size": 3}
 GEMMA4_AUDIO_CONFIG = {"hidden_size": 1024, "num_hidden_layers": 12, "num_attention_heads": 8, "intermediate_size": 4096, "conv_kernel_size": 5, "attention_chunk_size": 12, "attention_context_left": 13, "attention_context_right": 0, "attention_logit_cap": 50.0, "output_proj_dims": 1536, "rms_norm_eps": 1e-6, "residual_weight": 0.5}
 
+# Encoder-free (gemma4_unified) multimodal embedders: raw patches/waveform projected directly into LM space.
+GEMMA4_UNIFIED_VISION_CONFIG = {"model_patch_size": 48, "patch_size": 16, "pooling_kernel_size": 3, "mm_embed_dim": 3840, "mm_posemb_size": 1120, "output_proj_dims": 3840, "rms_norm_eps": 1e-6}
+GEMMA4_UNIFIED_AUDIO_CONFIG = {"audio_samples_per_token": 640, "output_proj_dims": 640, "rms_norm_eps": 1e-6}
+
 @dataclass
 class Gemma4Config:
     vocab_size: int = 262144
@@ -35,6 +43,9 @@ class Gemma4Config:
     transformer_type: str = "gemma4"
     head_dim = 256
     global_head_dim = 512
+    num_global_key_value_heads = None
+    attention_k_eq_v = False
+    vision_bidirectional = False
     rms_norm_add = False
     mlp_activation = "gelu_pytorch_tanh"
     qkv_bias = False
@@ -51,6 +62,7 @@ class Gemma4Config:
     num_kv_shared_layers: int = 18
     use_double_wide_mlp: bool = False
     stop_tokens = [1, 50, 106]
+    suppress_tokens = []
     vision_config = GEMMA4_VISION_CONFIG
     audio_config = GEMMA4_AUDIO_CONFIG
     mm_tokens_per_image = 280
@@ -72,12 +84,30 @@ class Gemma4_31B_Config(Gemma4Config):
     num_hidden_layers: int = 60
     num_attention_heads: int = 32
     num_key_value_heads: int = 16
+    vision_bidirectional = True
     sliding_attention = [1024, 1024, 1024, 1024, 1024, False]
     hidden_size_per_layer_input: int = 0
     num_kv_shared_layers: int = 0
     audio_config = None
     vision_config = GEMMA4_VISION_31B_CONFIG
 
+@dataclass
+class Gemma4_12B_Config(Gemma4Config):
+    hidden_size: int = 3840
+    intermediate_size: int = 15360
+    num_hidden_layers: int = 48
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 8
+    num_global_key_value_heads = 1
+    attention_k_eq_v = True
+    vision_bidirectional = True
+    sliding_attention = [1024, 1024, 1024, 1024, 1024, False]
+    hidden_size_per_layer_input: int = 0
+    num_kv_shared_layers: int = 0
+    audio_config = GEMMA4_UNIFIED_AUDIO_CONFIG
+    vision_config = GEMMA4_UNIFIED_VISION_CONFIG
+    suppress_tokens = [258883, 258882]
+
 
 # unfused RoPE as addcmul_ RoPE diverges from reference code
 def _apply_rotary_pos_emb(x, freqs_cis):
@@ -89,17 +119,18 @@ def _apply_rotary_pos_emb(x, freqs_cis):
     return out
 
 class Gemma4Attention(nn.Module):
-    def __init__(self, config, head_dim, device=None, dtype=None, ops=None):
+    def __init__(self, config, head_dim, num_kv_heads=None, k_eq_v=False, device=None, dtype=None, ops=None):
         super().__init__()
         self.num_heads = config.num_attention_heads
-        self.num_kv_heads = config.num_key_value_heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else config.num_key_value_heads
         self.hidden_size = config.hidden_size
         self.head_dim = head_dim
         self.inner_size = self.num_heads * head_dim
 
         self.q_proj = ops.Linear(config.hidden_size, self.inner_size, bias=config.qkv_bias, device=device, dtype=dtype)
         self.k_proj = ops.Linear(config.hidden_size, self.num_kv_heads * head_dim, bias=config.qkv_bias, device=device, dtype=dtype)
-        self.v_proj = ops.Linear(config.hidden_size, self.num_kv_heads * head_dim, bias=config.qkv_bias, device=device, dtype=dtype)
+        # k_eq_v: V reuses the K projection (no separate v_proj weight)
+        self.v_proj = None if k_eq_v else ops.Linear(config.hidden_size, self.num_kv_heads * head_dim, bias=config.qkv_bias, device=device, dtype=dtype)
         self.o_proj = ops.Linear(self.inner_size, config.hidden_size, bias=False, device=device, dtype=dtype)
 
         self.q_norm = None
@@ -133,7 +164,10 @@ class Gemma4Attention(nn.Module):
             shareable_kv = None
         else:
             xk = self.k_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
-            xv = self.v_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
+            if self.v_proj is not None:
+                xv = self.v_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
+            else:
+                xv = xk  # k_eq_v: V is the raw K projection (before k_norm/RoPE)
             if self.k_norm is not None:
                 xk = self.k_norm(xk)
             xv = rms_norm(xv)
@@ -186,7 +220,10 @@ class TransformerBlockGemma4(nn.Module):
 
         head_dim = config.head_dim if self.sliding_attention else config.global_head_dim
 
-        self.self_attn = Gemma4Attention(config, head_dim=head_dim, device=device, dtype=dtype, ops=ops)
+        # k_eq_v only on global layers, which then use num_global_key_value_heads
+        k_eq_v = config.attention_k_eq_v and not self.sliding_attention
+        num_kv_heads = config.num_global_key_value_heads if k_eq_v else config.num_key_value_heads
+        self.self_attn = Gemma4Attention(config, head_dim=head_dim, num_kv_heads=num_kv_heads, k_eq_v=k_eq_v, device=device, dtype=dtype, ops=ops)
 
         num_kv_shared = config.num_kv_shared_layers
         first_kv_shared = config.num_hidden_layers - num_kv_shared
@@ -203,9 +240,9 @@ class TransformerBlockGemma4(nn.Module):
             self.per_layer_input_gate = ops.Linear(config.hidden_size, self.hidden_size_per_layer_input, bias=False, device=device, dtype=dtype)
             self.per_layer_projection = ops.Linear(self.hidden_size_per_layer_input, config.hidden_size, bias=False, device=device, dtype=dtype)
             self.post_per_layer_input_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, device=device, dtype=dtype)
-            self.register_buffer("layer_scalar", torch.ones(1, device=device, dtype=dtype))
-        else:
-            self.layer_scalar = None
+
+        # layer_scalar exists on every gemma4 variant, independent of per-layer input
+        self.register_buffer("layer_scalar", torch.empty(1, device=device, dtype=dtype))
 
     def forward(self, x, attention_mask=None, freqs_cis=None, past_key_value=None, per_layer_input=None, shared_kv=None):
         sliding_window = None
@@ -244,8 +281,7 @@ class TransformerBlockGemma4(nn.Module):
             x = self.post_per_layer_input_norm(x)
             x = residual + x
 
-        if self.layer_scalar is not None:
-            x = x * self.layer_scalar
+        x = x * comfy.ops.cast_to_input(self.layer_scalar, x)
 
         return x, present_key_value, shareable_kv
 
@@ -334,6 +370,19 @@ class Gemma4Transformer(nn.Module):
             causal_mask.masked_fill_(torch.ones_like(causal_mask, dtype=torch.bool).triu_(1), min_val)
             mask = mask + causal_mask if mask is not None else causal_mask
 
+            # Bidirectional attention within each image soft-token block (prefill only; text/audio stay causal).
+            if self.config.vision_bidirectional and past_len == 0 and embeds_info:
+                block_ids = torch.full((seq_len,), -1, dtype=torch.long, device=x.device)
+                group = 0
+                for info in embeds_info:
+                    if info.get("type") == "image":
+                        start = info["index"]
+                        block_ids[start:start + info["size"]] = group
+                        group += 1
+                if group > 0:
+                    same_block = (block_ids[:, None] == block_ids[None, :]) & (block_ids[:, None] >= 0)
+                    mask = mask.masked_fill(same_block, 0.0)
+
         # Per-layer inputs
         per_layer_inputs = None
         if self.hidden_size_per_layer_input:
@@ -354,8 +403,24 @@ class Gemma4Transformer(nn.Module):
         shared_global_kv = None   # KV from last non-shared global layer
 
         intermediate = None
+        all_intermediate = None
+        only_layers = None
+        if intermediate_output is not None:
+            if isinstance(intermediate_output, list):
+                all_intermediate = []
+                only_layers = {len(self.layers) + layer if layer < 0 else layer for layer in intermediate_output}
+            elif intermediate_output == "all":
+                all_intermediate = []
+                intermediate_output = None
+            elif intermediate_output < 0:
+                intermediate_output = len(self.layers) + intermediate_output
+
         next_key_values = []
         for i, layer in enumerate(self.layers):
+            if all_intermediate is not None:
+                if only_layers is None or (i in only_layers):
+                    all_intermediate.append(x.unsqueeze(1).clone())
+
             past_kv = past_key_values[i] if past_key_values is not None and len(past_key_values) > 0 else None
 
             layer_kwargs = {}
@@ -385,7 +450,18 @@ class Gemma4Transformer(nn.Module):
         if self.norm is not None:
             x = self.norm(x)
 
-        if len(next_key_values) > 0:
+        if all_intermediate is not None:
+            if only_layers is None or (len(self.layers) in only_layers):
+                all_intermediate.append(x.unsqueeze(1).clone())
+            if len(all_intermediate) > 0:
+                intermediate = torch.cat(all_intermediate, dim=1)
+
+        if intermediate is not None and final_layer_norm_intermediate and self.norm is not None:
+            intermediate = self.norm(intermediate)
+
+        # Only hand back the KV cache when caching was actually requested; SDClipModel reads
+        # outputs[2] as the pooled output.
+        if past_key_values is not None and len(next_key_values) > 0:
             return x, intermediate, next_key_values
         return x, intermediate
 
@@ -404,6 +480,8 @@ class Gemma4Base(BaseLlama, BaseGenerate, torch.nn.Module):
         cap = self.model.config.final_logit_softcapping
         if cap:
             logits = cap * torch.tanh(logits / cap)
+        if self.model.config.suppress_tokens:
+            logits[..., self.model.config.suppress_tokens] = torch.finfo(logits.dtype).min
         return logits
 
     def init_kv_cache(self, batch, max_cache_len, device, execution_dtype):
@@ -441,6 +519,28 @@ class Gemma4AudioMixin:
         return None, None
 
 
+class Gemma4UnifiedBase(Gemma4Base):
+    """Encoder-free multimodal Gemma4 (gemma4_unified, e.g. 12B): raw image patches and audio frames projected directly into LM space."""
+    def _init_model(self, config, dtype, device, operations):
+        self.num_layers = config.num_hidden_layers
+        self.model = Gemma4Transformer(config, device=device, dtype=dtype, ops=operations)
+        self.dtype = dtype
+        self.vision_model = Gemma4UnifiedVisionEmbedder(config.vision_config, device=device, dtype=dtype, ops=operations)
+        self.multi_modal_projector = Gemma4RMSNormProjector(config.vision_config["output_proj_dims"], config.hidden_size, dtype=dtype, device=device, ops=operations)
+        self.audio_projector = Gemma4RMSNormProjector(config.audio_config["output_proj_dims"], config.hidden_size, dtype=dtype, device=device, ops=operations)
+
+    def preprocess_embed(self, embed, device):
+        if embed["type"] == "image":
+            pixels = embed.pop("data").movedim(-1, 1).to(device, dtype=self.dtype)  # [B, H, W, C] -> [B, C, H, W], [0,1]
+            patches, positions = self.vision_model.patchify(pixels)
+            vision_out = self.vision_model(patches, positions)
+            return self.multi_modal_projector(vision_out), None
+        if embed["type"] == "audio":
+            audio = embed.pop("data").to(device, dtype=self.dtype)  # [1, T, audio_samples_per_token]
+            return self.audio_projector(audio), None
+        return None, None
+
+
 # Vision Encoder
 
 def _compute_vision_2d_rope(head_dim, pixel_position_ids, theta=100.0, device=None):
@@ -713,6 +813,73 @@ class Gemma4MultiModalProjector(Gemma4RMSNormProjector):
         super().__init__(config.vision_config["hidden_size"], config.hidden_size, dtype=dtype, device=device, ops=ops)
 
 
+# Encoder-free vision (gemma4_unified): raw merged pixel patches projected directly into LM space.
+
+def _patches_merge(patches, positions_xy, length):
+    patch_size = math.isqrt(patches.shape[-1] // 3)
+    k = math.isqrt(patches.shape[-2] // length)
+    batch = patches.shape[:-2]
+
+    max_x = positions_xy[..., 0].max(dim=-1, keepdim=True)[0] + 1
+    kidx = torch.div(positions_xy, k, rounding_mode="floor")
+    rem = torch.remainder(positions_xy, k)
+    order = rem[..., 0] + rem[..., 1] * k + k * k * kidx[..., 0] + k * max_x * kidx[..., 1]
+    perm = order.long().argsort(dim=-1)
+
+    merged = patches.gather(-2, perm.unsqueeze(-1).expand_as(patches))
+    merged = merged.reshape(*batch, length, k, k, patch_size, patch_size, 3)
+    merged = merged.permute(*range(len(batch)), -6, -5, -3, -4, -2, -1).reshape(*batch, length, (k * patch_size) ** 2 * 3)
+
+    pos = positions_xy.gather(-2, perm.unsqueeze(-1).expand_as(positions_xy))
+    pad = (positions_xy == -1).all(dim=-1, keepdim=True)
+    pos = torch.where(pad, positions_xy, pos).reshape(*batch, length, k * k, 2)
+    pos = torch.div(pos, k, rounding_mode="floor").min(dim=-2)[0]
+    return merged, pos
+
+
+class Gemma4UnifiedVisionEmbedder(nn.Module):
+    """Encoder-free patch embedder (LN -> Dense -> LN -> +2D posemb -> LN); projection to text space is the separate multi_modal_projector."""
+    def __init__(self, config, device=None, dtype=None, ops=None):
+        super().__init__()
+        self.patch_size = config["patch_size"]
+        self.pooling_kernel_size = config["pooling_kernel_size"]
+        patch_dim = config["model_patch_size"] ** 2 * 3
+        mm_embed_dim = config["mm_embed_dim"]
+        self.patch_ln1 = ops.LayerNorm(patch_dim, device=device, dtype=dtype)
+        self.patch_dense = ops.Linear(patch_dim, mm_embed_dim, device=device, dtype=dtype)
+        self.patch_ln2 = ops.LayerNorm(mm_embed_dim, device=device, dtype=dtype)
+        self.pos_embedding = nn.Parameter(torch.empty(config["mm_posemb_size"], 2, mm_embed_dim, device=device, dtype=dtype))
+        self.pos_norm = ops.LayerNorm(mm_embed_dim, device=device, dtype=dtype)
+
+    def patchify(self, pixels):
+        """pixels: [B, C, H, W] in [0,1] -> merged patches [B, N, 6912], positions [B, N, 2]."""
+        ps, k = self.patch_size, self.pooling_kernel_size
+        out_patches, out_positions = [], []
+        for img in pixels:
+            ph, pw = img.shape[-2] // ps, img.shape[-1] // ps
+            teacher = img.reshape(img.shape[0], ph, ps, pw, ps).permute(1, 3, 2, 4, 0).reshape(ph * pw, -1)
+            grid = torch.meshgrid(torch.arange(pw, device=img.device), torch.arange(ph, device=img.device), indexing="xy")
+            tpos = torch.stack(grid, dim=-1).reshape(teacher.shape[0], 2)
+            n_model = teacher.shape[0] // (k * k)
+            mp, mpos = _patches_merge(teacher.unsqueeze(0), tpos.unsqueeze(0), n_model)
+            out_patches.append(mp.squeeze(0))
+            out_positions.append(mpos.squeeze(0))
+        return torch.stack(out_patches), torch.stack(out_positions)
+
+    def forward(self, pixel_values, image_position_ids):
+        x = self.patch_ln1(pixel_values)
+        x = self.patch_dense(x)
+        x = self.patch_ln2(x)
+
+        clamped = image_position_ids.clamp(min=0).long()
+        valid = (image_position_ids != -1).to(x.dtype).unsqueeze(-1)
+        axes = torch.arange(2, device=image_position_ids.device)
+        pos = comfy.model_management.cast_to_device(self.pos_embedding, x.device, x.dtype)
+        pos_embs = (pos[clamped, axes] * valid).sum(-2)
+        x = x + pos_embs
+        return self.pos_norm(x)
+
+
 # Audio Encoder
 
 class Gemma4AudioConvSubsampler(nn.Module):
@@ -990,6 +1157,30 @@ class Gemma4AudioProjector(Gemma4RMSNormProjector):
 
 # Tokenizer and Wrappers
 
+def _get_aspect_ratio_preserving_size(height, width, patch_size, max_patches, pooling_kernel_size):
+    target_px = max_patches * patch_size ** 2
+    factor = math.sqrt(target_px / (height * width))
+    side_mult = pooling_kernel_size * patch_size
+    target_height = math.floor(factor * height / side_mult) * side_mult
+    target_width = math.floor(factor * width / side_mult) * side_mult
+
+    if target_height == 0 and target_width == 0:
+        raise ValueError(f"Attempting to resize to a 0 x 0 image. Resized height should be divisible by {side_mult}.")
+
+    max_side_length = (max_patches // pooling_kernel_size ** 2) * side_mult
+    if target_height == 0:
+        target_height = side_mult
+        target_width = min(math.floor(width / height) * side_mult, max_side_length)
+    elif target_width == 0:
+        target_width = side_mult
+        target_height = min(math.floor(height / width) * side_mult, max_side_length)
+
+    if target_height * target_width > target_px:
+        raise ValueError(f"Resizing [{height}x{width}] to [{target_height}x{target_width}] exceeds the patch budget.")
+
+    return target_height, target_width
+
+
 class Gemma4_Tokenizer():
     tokenizer_json_data = None
 
@@ -998,25 +1189,35 @@ class Gemma4_Tokenizer():
             return {"tokenizer_json": self.tokenizer_json_data}
         return {}
 
-    def _extract_mel_spectrogram(self, waveform, sample_rate):
-        """Extract 128-bin log mel spectrogram.
-        Uses numpy for FFT/matmul/log to produce bit-identical results with reference code.
-        """
-        # Mix to mono first, then resample to 16kHz
+    def _audio_token_count(self, num_samples):
+        # Default (E2B/E4B): mel frames after two stride-2 conv subsamples.
+        _fl = 320  # int(round(16000 * 20.0 / 1000.0))
+        _hl = 160  # int(round(16000 * 10.0 / 1000.0))
+        _nmel = (num_samples + _fl // 2 - (_fl + 1)) // _hl + 1
+        _t = _nmel
+        for _ in range(2):
+            _t = (_t + 2 - 3) // 2 + 1
+        return min(_t, 750)
+
+    @staticmethod
+    def _resample_16k(waveform, sample_rate):
+        """Mix to mono and resample to 16kHz. Kaiser params reproduce the reference (transformers
+        load_audio -> librosa/soxr_hq) to ~1e-12 MSE using only torchaudio."""
         if waveform.dim() > 1 and waveform.shape[0] > 1:
             waveform = waveform.mean(dim=0, keepdim=True)
         if waveform.dim() == 1:
             waveform = waveform.unsqueeze(0)
-        audio = waveform.squeeze(0).float().numpy()
+        audio = waveform.float()
         if sample_rate != 16000:
-            # Use scipy's resample_poly with a high-quality FIR filter to get as close as possible to librosa's resampling (while still not full match)
-            from scipy.signal import resample_poly, firwin
-            from math import gcd
-            g = gcd(sample_rate, 16000)
-            up, down = 16000 // g, sample_rate // g
-            L = max(up, down)
-            h = firwin(160 * L + 1, 0.96 / L, window=('kaiser', 6.5))
-            audio = resample_poly(audio, up, down, window=h).astype(np.float32)
+            audio = AF.resample(audio, sample_rate, 16000, resampling_method="sinc_interp_kaiser",
+                                lowpass_filter_width=121, rolloff=0.9568384289091556, beta=21.01531462440614)
+        return audio.squeeze(0).contiguous()
+
+    def _extract_audio_features(self, waveform, sample_rate):
+        """Default (E2B/E4B): 128-bin log mel spectrogram for the conformer audio encoder.
+        Uses numpy for FFT/matmul/log to produce bit-identical results with reference code.
+        """
+        audio = self._resample_16k(waveform, sample_rate).numpy()
         n = len(audio)
 
         # Pad to multiple of 128, build sample-level mask
@@ -1064,8 +1265,8 @@ class Gemma4_Tokenizer():
         if audio is not None:
             waveform = audio["waveform"].squeeze(0) if hasattr(audio, "__getitem__") else audio
             sample_rate = audio.get("sample_rate", 16000) if hasattr(audio, "get") else 16000
-            mel, mel_mask = self._extract_mel_spectrogram(waveform, sample_rate)
-            audio_features = [(mel.unsqueeze(0), mel_mask.unsqueeze(0))]  # ([1, T, 128], [1, T])
+            feat, feat_mask = self._extract_audio_features(waveform, sample_rate)
+            audio_features = [(feat.unsqueeze(0), feat_mask.unsqueeze(0))]  # ([1, T, D], [1, T])
 
         # Process image/video frames
         is_video = video is not None
@@ -1090,13 +1291,8 @@ class Gemma4_Tokenizer():
             pooling_k = 3
             max_soft_tokens = kwargs.get("max_soft_tokens", 70 if is_video else 280)
             max_patches = max_soft_tokens * pooling_k * pooling_k
-            target_px = max_patches * patch_size * patch_size
-            factor = (target_px / (h * w)) ** 0.5
-            side_mult = pooling_k * patch_size
-            target_h = max(int(factor * h // side_mult) * side_mult, side_mult)
-            target_w = max(int(factor * w // side_mult) * side_mult, side_mult)
+            target_h, target_w = _get_aspect_ratio_preserving_size(h, w, patch_size, max_patches, pooling_k)
 
-            import torchvision.transforms.functional as TVF
             for i in range(num_frames):
                 # rescaling to match reference code
                 s = (samples[i].clamp(0, 1) * 255).to(torch.uint8)  # [C, H, W] uint8
@@ -1115,7 +1311,7 @@ class Gemma4_Tokenizer():
                 llama_text = llama_template.format(text)
             else:
                 # Build template from modalities present
-                system = "<|turn>system\n<|think|><turn|>\n" if thinking else ""
+                system = "<|turn>system\n<|think|>\n<turn|>\n" if thinking else ""
                 media = ""
                 if len(images) > 0:
                     if is_video:
@@ -1135,15 +1331,11 @@ class Gemma4_Tokenizer():
                 if len(audio_features) > 0:
                     # Compute audio token count (always at 16kHz)
                     num_samples = int(waveform.shape[-1] * 16000 / sample_rate) if sample_rate != 16000 else waveform.shape[-1]
-                    _fl = 320  # int(round(16000 * 20.0 / 1000.0))
-                    _hl = 160  # int(round(16000 * 10.0 / 1000.0))
-                    _nmel = (num_samples + _fl // 2 - (_fl + 1)) // _hl + 1
-                    _t = _nmel
-                    for _ in range(2):
-                        _t = (_t + 2 - 3) // 2 + 1
-                    n_audio_tokens = min(_t, 750)
+                    n_audio_tokens = self._audio_token_count(num_samples)
                     media += "<|audio>" + "<|audio|>" * n_audio_tokens + "<audio|>"
-                llama_text = f"{system}<|turn>user\n{media}{text}<turn|>\n<|turn>model\n"
+                # Non-thinking mode primes an empty thought channel so the model answers directly.
+                model_open = "" if thinking else "<|channel>thought\n<channel|>"
+                llama_text = f"{system}<|turn>user\n{text}{media}<turn|>\n<|turn>model\n{model_open}"
 
         text_tokens = super().tokenize_with_weights(llama_text, return_word_ids)
 
@@ -1178,7 +1370,6 @@ class Gemma4_Tokenizer():
 class _Gemma4Tokenizer:
     """Tokenizer using the tokenizers (Gemma4 doesn't come with sentencepiece model)"""
     def __init__(self, tokenizer_json_bytes=None, **kwargs):
-        from tokenizers import Tokenizer
         if isinstance(tokenizer_json_bytes, torch.Tensor):
             tokenizer_json_bytes = bytes(tokenizer_json_bytes.tolist())
         self.tokenizer = Tokenizer.from_str(tokenizer_json_bytes.decode("utf-8"))
@@ -1224,6 +1415,30 @@ class Gemma4Tokenizer(sd1_clip.SD1Tokenizer):
         super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="gemma4", tokenizer=self.tokenizer_class)
 
 
+class Gemma4UnifiedSDTokenizer(Gemma4SDTokenizer):
+    """Encoder-free (gemma4_unified) audio: raw 16kHz waveform frames instead of mel spectrogram."""
+    embedding_size = 3840
+
+    def _extract_audio_features(self, waveform, sample_rate):
+        audio = self._resample_16k(waveform, sample_rate)
+        spt = 640  # audio_samples_per_token (40ms at 16kHz)
+        pad = (-audio.shape[0]) % spt
+        if pad:
+            audio = torch.nn.functional.pad(audio, (0, pad))
+        num_tokens = audio.shape[0] // spt
+        feats = audio[:num_tokens * spt].reshape(num_tokens, spt)
+        feats = feats[:750]  # audio_seq_length cap (matches reference truncation, ~30s)
+        mask = torch.ones(feats.shape[0], dtype=torch.bool)
+        return feats, mask
+
+    def _audio_token_count(self, num_samples):
+        return min((num_samples + 639) // 640, 750)
+
+
+class Gemma4UnifiedTokenizer(Gemma4Tokenizer):
+    tokenizer_class = Gemma4UnifiedSDTokenizer
+
+
 # Model wrappers
 class Gemma4Model(sd1_clip.SDClipModel):
     model_class = None
@@ -1256,7 +1471,7 @@ class Gemma4Model(sd1_clip.SDClipModel):
                 expanded_idx += 1
         initial_token_ids = [ids]
         input_ids = torch.tensor(initial_token_ids, device=self.execution_device)
-        return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, initial_tokens=initial_token_ids[0], presence_penalty=presence_penalty, initial_input_ids=input_ids)
+        return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, initial_tokens=initial_token_ids[0], presence_penalty=presence_penalty, initial_input_ids=input_ids, embeds_info=embeds_info)
 
 
 def gemma4_te(dtype_llama=None, llama_quantization_metadata=None, model_class=None):
@@ -1296,3 +1511,11 @@ def _make_variant(config_cls):
 Gemma4_E4B = _make_variant(Gemma4Config)
 Gemma4_E2B = _make_variant(Gemma4_E2B_Config)
 Gemma4_31B = _make_variant(Gemma4_31B_Config)
+
+
+# Gemma4 12B Unified: encoder-free multimodal, distinct base/tokenizer (not via _make_variant).
+class Gemma4_12B(Gemma4UnifiedBase):
+    def __init__(self, config_dict, dtype, device, operations):
+        super().__init__()
+        self._init_model(Gemma4_12B_Config(**config_dict), dtype, device, operations)
+Gemma4_12B.tokenizer = Gemma4UnifiedTokenizer

comfy/text_encoders/llama.py

@@ -876,7 +876,7 @@ class BaseGenerate:
                                     torch.empty([batch, model_config.num_key_value_heads, max_cache_len, model_config.head_dim], device=device, dtype=execution_dtype), 0))
         return past_key_values
 
-    def generate(self, embeds=None, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.9, min_p=0.0, repetition_penalty=1.0, seed=42, stop_tokens=None, initial_tokens=[], execution_dtype=None, min_tokens=0, presence_penalty=0.0, initial_input_ids=None, position_ids=None, deepstack_embeds=None, visual_pos_masks=None):
+    def generate(self, embeds=None, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.9, min_p=0.0, repetition_penalty=1.0, seed=42, stop_tokens=None, initial_tokens=[], execution_dtype=None, min_tokens=0, presence_penalty=0.0, initial_input_ids=None, position_ids=None, deepstack_embeds=None, visual_pos_masks=None, embeds_info=None):
         device = embeds.device
 
         if stop_tokens is None:
@@ -911,7 +911,7 @@ class BaseGenerate:
             if step == 0 and deepstack_embeds is not None:
                 extra["deepstack_embeds"] = deepstack_embeds
                 extra["visual_pos_masks"] = visual_pos_masks
-            x, _, past_key_values = self.model.forward(None, embeds=embeds, attention_mask=None, past_key_values=past_key_values, input_ids=current_input_ids, position_ids=position_ids, **extra)
+            x, _, past_key_values = self.model.forward(None, embeds=embeds, attention_mask=None, past_key_values=past_key_values, input_ids=current_input_ids, position_ids=position_ids, **extra, embeds_info=(embeds_info if step == 0 else None))
             logits = self.logits(x)[:, -1]
             next_token = self.sample_token(logits, temperature, top_k, top_p, min_p, repetition_penalty, initial_tokens + generated_token_ids, generator, do_sample=do_sample, presence_penalty=presence_penalty)
             token_id = next_token[0].item()

comfy_api_nodes/apis/gemini.py

@@ -1,6 +1,6 @@
 from datetime import date
 from enum import Enum
-from typing import Any
+from typing import Any, Literal
 
 from pydantic import BaseModel, Field
 
@@ -242,3 +242,60 @@ class GeminiGenerateContentResponse(BaseModel):
     promptFeedback: GeminiPromptFeedback | None = Field(None)
     usageMetadata: GeminiUsageMetadata | None = Field(None)
     modelVersion: str | None = Field(None)
+
+
+class GeminiInteractionTextPart(BaseModel):
+    type: Literal["text"] = "text"
+    text: str = Field(...)
+
+
+class GeminiInteractionMediaPart(BaseModel):
+    type: str = Field(..., description="One of: image, video, audio, document.")
+    data: str | None = Field(None, description="Base64-encoded media bytes.")
+    uri: str | None = Field(None, description="URI of the media, as an alternative to inline data.")
+    mime_type: str | None = Field(None)
+
+
+class GeminiInteractionGenerationConfig(BaseModel):
+    temperature: float | None = Field(None, ge=0.0, le=2.0)
+    top_p: float | None = Field(None, ge=0.0, le=1.0)
+
+
+class GeminiInteractionRequest(BaseModel):
+    model: str = Field(...)
+    input: list[GeminiInteractionTextPart | GeminiInteractionMediaPart] = Field(...)
+    generation_config: GeminiInteractionGenerationConfig | None = Field(None)
+
+
+class GeminiInteractionModalityTokens(BaseModel):
+    modality: str | None = Field(None, description="One of: text, image, audio, video, document.")
+    tokens: int | None = Field(None)
+
+
+class GeminiInteractionUsage(BaseModel):
+    input_tokens_by_modality: list[GeminiInteractionModalityTokens] | None = Field(None)
+    output_tokens_by_modality: list[GeminiInteractionModalityTokens] | None = Field(None)
+    total_thought_tokens: int | None = Field(None)
+
+
+class GeminiInteractionContent(BaseModel):
+    type: str | None = Field(None)
+    text: str | None = Field(None)
+    data: str | None = Field(None)
+    uri: str | None = Field(None)
+    mime_type: str | None = Field(None)
+
+
+class GeminiInteractionStep(BaseModel):
+    type: str | None = Field(None)
+    content: list[GeminiInteractionContent] | None = Field(None)
+
+
+class GeminiInteraction(BaseModel):
+    id: str | None = Field(None)
+    status: str | None = Field(
+        None,
+        description="One of: in_progress, requires_action, completed, failed, cancelled, incomplete.",
+    )
+    steps: list[GeminiInteractionStep] | None = Field(None)
+    usage: GeminiInteractionUsage | None = Field(None)

comfy_api_nodes/nodes_gemini.py

@@ -24,6 +24,11 @@ from comfy_api_nodes.apis.gemini import (
     GeminiImageGenerateContentRequest,
     GeminiImageGenerationConfig,
     GeminiInlineData,
+    GeminiInteraction,
+    GeminiInteractionGenerationConfig,
+    GeminiInteractionMediaPart,
+    GeminiInteractionRequest,
+    GeminiInteractionTextPart,
     GeminiMimeType,
     GeminiPart,
     GeminiRole,
@@ -51,6 +56,7 @@ from comfy_api_nodes.util import (
 )
 
 GEMINI_BASE_ENDPOINT = "/proxy/vertexai/gemini"
+GEMINI_INTERACTIONS_ENDPOINT = "/proxy/gemini-interactions"
 GEMINI_MAX_INPUT_FILE_SIZE = 20 * 1024 * 1024  # 20 MB
 GEMINI_URL_INPUT_BUDGET = 10
 GEMINI_MAX_INLINE_BYTES = 18 * 1024 * 1024
@@ -231,29 +237,10 @@ async def get_image_from_response(response: GeminiGenerateContentResponse, thoug
     return torch.cat(image_tensors, dim=0)
 
 
-async def get_video_from_response(
-    response: GeminiGenerateContentResponse, cls: type[IO.ComfyNode] | None = None
-) -> InputImpl.VideoFromFile:
-    parts = get_parts_by_type(response, "video/*")
-    for part in parts:
-        if part.inlineData and part.inlineData.data:
-            return InputImpl.VideoFromFile(BytesIO(base64.b64decode(part.inlineData.data)))
-        if part.fileData and part.fileData.fileUri:
-            return await download_url_to_video_output(part.fileData.fileUri, cls=cls)
-    model_message = get_text_from_response(response).strip()
-    if model_message:
-        raise ValueError(f"Gemini did not generate a video. Model response: {model_message}")
-    raise ValueError(
-        "Gemini did not generate a video. Try rephrasing your prompt, "
-        "shortening the requested duration, or reducing the number of input images/videos."
-    )
-
-
 def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | None:
     if not response.modelVersion:
         return None
     # Define prices (Cost per 1,000,000 tokens), see https://cloud.google.com/vertex-ai/generative-ai/pricing
-    output_video_tokens_price = 0.0
     if response.modelVersion == "gemini-2.5-pro":
         input_tokens_price = 1.25
         output_text_tokens_price = 10.0
@@ -290,11 +277,6 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
         input_tokens_price = 0.25
         output_text_tokens_price = 1.50
         output_image_tokens_price = 30.0
-    elif response.modelVersion == "gemini-omni-flash-preview":
-        input_tokens_price = 2.145
-        output_text_tokens_price = 12.87
-        output_image_tokens_price = 0.0
-        output_video_tokens_price = 25.025
     else:
         return None
     final_price = response.usageMetadata.promptTokenCount * input_tokens_price
@@ -302,8 +284,6 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
         for i in response.usageMetadata.candidatesTokensDetails:
             if i.modality == Modality.IMAGE:
                 final_price += output_image_tokens_price * i.tokenCount  # for Nano Banana models
-            elif i.modality == Modality.VIDEO:
-                final_price += output_video_tokens_price * i.tokenCount  # for Omni Flash
             else:
                 final_price += output_text_tokens_price * i.tokenCount
     if response.usageMetadata.thoughtsTokenCount:
@@ -311,6 +291,58 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
     return final_price / 1_000_000.0
 
 
+def get_text_from_interaction(interaction: GeminiInteraction) -> str:
+    """Extract and concatenate all model output text from an Interactions API response."""
+    texts = []
+    for step in interaction.steps or []:
+        if step.type != "model_output":
+            continue
+        for content in step.content or []:
+            if content.type == "text" and content.text:
+                texts.append(content.text)
+    return "\n".join(texts)
+
+
+async def get_video_from_interaction(
+    interaction: GeminiInteraction, cls: type[IO.ComfyNode] | None = None
+) -> InputImpl.VideoFromFile:
+    for step in interaction.steps or []:
+        if step.type != "model_output":
+            continue
+        for content in step.content or []:
+            if content.type != "video":
+                continue
+            if content.data:
+                return InputImpl.VideoFromFile(BytesIO(base64.b64decode(content.data)))
+            if content.uri:
+                return await download_url_to_video_output(content.uri, cls=cls)
+    model_message = get_text_from_interaction(interaction).strip()
+    if model_message:
+        raise ValueError(f"Gemini did not generate a video. Model response: {model_message}")
+    raise ValueError(
+        "Gemini did not generate a video. Try rephrasing your prompt, "
+        "shortening the requested duration, or reducing the number of input images/videos."
+    )
+
+
+def calculate_interaction_tokens_price(interaction: GeminiInteraction) -> float | None:
+    if interaction.usage is None:
+        return None
+    input_tokens_price = 1.5
+    output_tokens_prices = {"text": 9.0, "video": 17.5}
+    thoughts_tokens_price = 9.0
+    final_price = 0.0
+    for i in interaction.usage.input_tokens_by_modality or []:
+        if i.tokens:
+            final_price += input_tokens_price * i.tokens
+    for i in interaction.usage.output_tokens_by_modality or []:
+        if i.tokens and i.modality in output_tokens_prices:
+            final_price += output_tokens_prices[i.modality] * i.tokens
+    if interaction.usage.total_thought_tokens:
+        final_price += thoughts_tokens_price * interaction.usage.total_thought_tokens
+    return final_price / 1_000_000.0
+
+
 def create_video_parts(video_input: Input.Video) -> list[GeminiPart]:
     """Convert a single video input to Gemini API compatible parts (inline MP4/H.264)."""
     base_64_string = video_to_base64_string(
@@ -445,6 +477,15 @@ async def build_gemini_media_parts(
     return parts
 
 
+def to_interaction_media_part(part: GeminiPart) -> GeminiInteractionMediaPart:
+    """Convert a fileData/inlineData GeminiPart into an Interactions API media part."""
+    if part.fileData:
+        mime = part.fileData.mimeType.value
+        return GeminiInteractionMediaPart(type=mime.split("/")[0], uri=part.fileData.fileUri, mime_type=mime)
+    mime = part.inlineData.mimeType.value
+    return GeminiInteractionMediaPart(type=mime.split("/")[0], data=part.inlineData.data, mime_type=mime)
+
+
 class GeminiNode(IO.ComfyNode):
     """
     Node to generate text responses from a Gemini model.
@@ -1684,7 +1725,7 @@ class GeminiVideoOmni(IO.ComfyNode):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                expr='{"type":"usd","usd":0.146,"format":{"suffix":"/second","approximate":true}}'
+                expr='{"type":"usd","usd":0.101,"format":{"suffix":"/second","approximate":true}}'
             ),
         )
 
@@ -1703,27 +1744,34 @@ class GeminiVideoOmni(IO.ComfyNode):
         for video in videos:
             validate_video_duration(video, max_duration=10)
 
-        parts: list[GeminiPart] = []
+        parts: list[GeminiInteractionTextPart | GeminiInteractionMediaPart] = []
         if images or videos:
-            parts.extend(await build_gemini_media_parts(cls, images, [], videos))
-        parts.append(GeminiPart(text=prompt))
-        response = await sync_op(
+            media_parts = await build_gemini_media_parts(cls, images, [], videos)
+            parts.extend(to_interaction_media_part(p) for p in media_parts)
+        parts.append(GeminiInteractionTextPart(text=prompt))
+        interaction = await sync_op(
             cls,
-            ApiEndpoint(path=f"{GEMINI_BASE_ENDPOINT}/{model_id}", method="POST"),
-            data=GeminiGenerateContentRequest(
-                contents=[GeminiContent(role=GeminiRole.user, parts=parts)],
-                generationConfig=GeminiGenerationConfig(
-                    responseModalities=["TEXT", "VIDEO"],
+            ApiEndpoint(path=GEMINI_INTERACTIONS_ENDPOINT, method="POST"),
+            data=GeminiInteractionRequest(
+                model=model_id,
+                input=parts,
+                generation_config=GeminiInteractionGenerationConfig(
                     temperature=model.get("temperature", 1.0),
-                    topP=model.get("top_p", 0.95),
+                    top_p=model.get("top_p", 0.95),
                 ),
             ),
-            response_model=GeminiGenerateContentResponse,
-            price_extractor=calculate_tokens_price,
+            response_model=GeminiInteraction,
+            price_extractor=calculate_interaction_tokens_price,
         )
+        if interaction.status != "completed":
+            model_message = get_text_from_interaction(interaction).strip()
+            raise ValueError(
+                f"Gemini interaction did not complete (status: {interaction.status})."
+                + (f" Model response: {model_message}" if model_message else "")
+            )
         return IO.NodeOutput(
-            await get_video_from_response(response, cls=cls),
-            get_text_from_response(response),
+            await get_video_from_interaction(interaction, cls=cls),
+            get_text_from_interaction(interaction),
         )
 
 

requirements.txt

@@ -1,5 +1,5 @@
 comfyui-frontend-package==1.45.21
-comfyui-workflow-templates==0.11.11
+comfyui-workflow-templates==0.11.12
 comfyui-embedded-docs==0.5.8
 torch
 torchsde
@@ -22,7 +22,7 @@ alembic
 SQLAlchemy>=2.0.0
 filelock
 av>=16.0.0
-comfy-kitchen==0.2.21
+comfy-kitchen==0.2.22
 comfy-aimdo==0.4.10
 requests
 simpleeval>=1.0.0