Add FunPackGemmaLoader, Gemma3 vision prompting, and i2v in-context conditioning

FunPackGemmaLoader: drop-in replacement for Load CLIP Model that additionally
loads the SigLIP vision tower and multi-modal projector from Google's
Gemma3-12B-IT checkpoint into the already-initialized Gemma3_12B module,
enabling reference-image vision prompting without a separate CLIP Vision model.

Vision prompting: _v2_encode_prompt accepts reference_image; when the clip
has _vision_loaded=True, tokenizes with image=reference_image and
skip_template=False so Gemma3 sees the reference frame before generating
conditioning. Cache key includes image fingerprint to avoid stale hits.
source_image from Studio/refine_v2 flows through to _v2_conditioning_source.

K/V injection revamp (Technique 5): identity blocks in i2v mode now use
in-context conditioning instead of lerp injection. Pre-hook prepends reference
hidden-state tokens to vx and extends v_pe with neutral-rotation entries
(cos=1, sin=0) so video tokens self-attend to reference tokens in attn1.
Post-hook crops the reference prefix from the block output. Strength raised
from 0.02 → 1.0 (sigma-gated: active at 0.35–0.80). _extend_v_pe helper
handles 2-D/3-D/4-D PE tensors. Anchor blocks keep the existing lerp path.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

__init__.py

@@ -19,7 +19,7 @@ if __package__:
         FunPackStoryMemLastFrameExtractor,
         FunPackVideoStitch,
     )
-    from .model_management import FunPackApplyLoraWeights, FunPackLoraLoader
+    from .model_management import FunPackApplyLoraWeights, FunPackGemmaLoader, FunPackLoraLoader
     from .samplers import FunPackHybridEuler2SSampler, FunPackDistilledFlowSampler, FunPackLTXAVSceneChainSampler
     from .templates import FunPackRefinementKeyLoader, FunPackSceneBuilder
 else:
@@ -50,9 +50,10 @@ else:
         FunPackStoryMemLastFrameExtractor = None
         FunPackVideoStitch = None
     try:
-        from model_management import FunPackApplyLoraWeights, FunPackLoraLoader
+        from model_management import FunPackApplyLoraWeights, FunPackGemmaLoader, FunPackLoraLoader
     except Exception:
         FunPackApplyLoraWeights = None
+        FunPackGemmaLoader = None
         FunPackLoraLoader = None
     try:
         from samplers import FunPackHybridEuler2SSampler, FunPackDistilledFlowSampler, FunPackLTXAVSceneChainSampler
@@ -87,6 +88,7 @@ NODE_CLASS_MAPPINGS = {
     "FunPackDistilledFlowSampler": FunPackDistilledFlowSampler,
     "FunPackLTXAVSceneChainSampler": FunPackLTXAVSceneChainSampler,
     "FunPackApplyLoraWeights": FunPackApplyLoraWeights,
+    "FunPackGemmaLoader": FunPackGemmaLoader,
     "FunPackLoraLoader": FunPackLoraLoader,
     "FunPackRefinementKeyLoader": FunPackRefinementKeyLoader,
     "FunPackSceneBuilder": FunPackSceneBuilder,
@@ -112,6 +114,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "FunPackDistilledFlowSampler": "FunPack Distilled Flow Sampler",
     "FunPackLTXAVSceneChainSampler": "FunPack LTXAV Scene Chain Sampler",
     "FunPackApplyLoraWeights": "FunPack Apply LoRA Weights",
+    "FunPackGemmaLoader": "FunPack Gemma Loader",
     "FunPackLoraLoader": "FunPack LoRA Loader",
     "FunPackRefinementKeyLoader": "FunPack Refinement Key Loader",
     "FunPackSceneBuilder": "FunPack Scene Builder",

conditioning.py

@@ -6154,13 +6154,32 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
             meta = {"pooled_output": None}
         return cond, meta
 
-    def _v2_encode_prompt(self, clip, prompt_text, encode_cache=None):
+    @staticmethod
+    def _gemma3_has_vision(clip):
+        try:
+            return getattr(clip.cond_stage_model.gemma3_12b.transformer, "_vision_loaded", False)
+        except AttributeError:
+            return False
+
+    @staticmethod
+    def _image_fingerprint(image):
+        try:
+            flat = image.float().reshape(-1)[:16]
+            return hash(tuple(flat.tolist()))
+        except Exception:
+            return None
+
+    def _v2_encode_prompt(self, clip, prompt_text, encode_cache=None, reference_image=None):
         if clip is None:
             return None, {"pooled_output": None}, "CLIP missing"
         prompt_text = str(prompt_text or "").strip()
         if not prompt_text:
             return None, {"pooled_output": None}, "prompt empty"
-        cache_key = (id(clip), prompt_text)
+
+        use_vision = reference_image is not None and self._gemma3_has_vision(clip)
+        img_fp = self._image_fingerprint(reference_image) if use_vision else None
+
+        cache_key = (id(clip), prompt_text, img_fp)
         if isinstance(encode_cache, dict):
             cached = encode_cache.get(cache_key)
             if cached is not None:
@@ -6171,13 +6190,18 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
                 encode_cache[cache_key] = cached
             return cached
         try:
-            encoded = clip.encode_from_tokens_scheduled(clip.tokenize(prompt_text))
+            if use_vision:
+                tokens = clip.tokenize(prompt_text, image=reference_image, skip_template=False)
+            else:
+                tokens = clip.tokenize(prompt_text)
+            encoded = clip.encode_from_tokens_scheduled(tokens)
         except Exception as error:
             return None, {"pooled_output": None}, f"encode failed: {error}"
         cond, meta = self._v2_extract_conditioning(encoded)
         if not isinstance(cond, torch.Tensor):
             return None, {"pooled_output": None}, "encode returned invalid conditioning"
-        result = (cond, meta, f"encoded {self._get_conditioning_seq_len(cond)} positions")
+        vision_tag = " +vision" if use_vision else ""
+        result = (cond, meta, f"encoded {self._get_conditioning_seq_len(cond)} positions{vision_tag}")
         if isinstance(encode_cache, dict):
             encode_cache[cache_key] = result
         if len(_V2_PERSISTENT_ENCODE_CACHE) < _V2_PERSISTENT_CACHE_MAX:
@@ -6193,9 +6217,11 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
             return f"Gemma3 tokenizer loaded: {source}"
         return "Gemma3 tokenizer loaded"
 
-    def _v2_conditioning_source(self, clip, prompt_text, positive_conditioning, encode_cache=None):
+    def _v2_conditioning_source(self, clip, prompt_text, positive_conditioning, encode_cache=None, reference_image=None):
         if clip is not None:
-            cond, meta, encode_status = self._v2_encode_prompt(clip, prompt_text, encode_cache=encode_cache)
+            cond, meta, encode_status = self._v2_encode_prompt(
+                clip, prompt_text, encode_cache=encode_cache, reference_image=reference_image
+            )
             return cond, meta, encode_status, "CLIP-owned"
 
         if positive_conditioning is None:
@@ -11408,6 +11434,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
             prompt_to_encode,
             positive_conditioning,
             encode_cache=encode_cache,
+            reference_image=source_image,
         )
         fallback_graph = render_refinement_loss_graph(refinement_key, "v2", "clip", 0, 0.0, [])
         if not isinstance(cond, torch.Tensor):

ltx_enhancements.py

@@ -433,6 +433,38 @@ def _sigma_gated_strength(base_strength, sigma, sigma_high=0.5, sigma_low=0.2):
     return base_strength * (sigma - sigma_low) / (sigma_high - sigma_low)
 
 
+def _extend_v_pe(kwargs, n_ref):
+    """Extend v_pe in kwargs by prepending n_ref neutral-rotation entries (cos=1, sin=0)."""
+    v_pe = kwargs.get("v_pe")
+    if v_pe is None:
+        return kwargs
+    try:
+        cos, sin = v_pe[0], v_pe[1]
+        dev, dt = cos.device, cos.dtype
+        ndim = cos.ndim
+        if ndim == 4:        # [B, H, T, d]
+            r = (cos.shape[0], cos.shape[1], n_ref, cos.shape[3])
+            dim = 2
+        elif ndim == 3:      # [B, T, d]
+            r = (cos.shape[0], n_ref, cos.shape[2])
+            dim = 1
+        elif ndim == 2:      # [T, d]
+            r = (n_ref, cos.shape[1])
+            dim = 0
+        else:
+            return kwargs
+        ref_cos = torch.ones(r, device=dev, dtype=dt)
+        ref_sin = torch.zeros(r, device=dev, dtype=dt)
+        ext_cos = torch.cat([ref_cos, cos], dim=dim)
+        ext_sin = torch.cat([ref_sin, sin], dim=dim)
+        tail = tuple(v_pe[2:]) if len(v_pe) > 2 else ()
+        new_kwargs = dict(kwargs)
+        new_kwargs["v_pe"] = (ext_cos, ext_sin) + tail
+        return new_kwargs
+    except Exception:
+        return kwargs
+
+
 def _build_block_replacement(block_idx, temp_scale, capture_buf, inject_tensor, inject_strength,
                               sigma_state=None, sigma_gate=(0.2, 0.5)):
     """
@@ -588,7 +620,7 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
         ref_x = _get_reference_video_tensor(reference_latent)
         lazy_injects = {idx: _LazyInject() for idx in list(IDENTITY_BLOCKS)}
         anchor_strength = 0.0   # no anchor injection for i2v - would bleed scene content
-        identity_strength = 0.02  # very light - hidden states encode full scene context
+        identity_strength = 1.0  # in-context gate: inject when sigma in gate range
     else:
         blessed_maps = _load_blessed_maps(refinement_key) if refinement_key else None
         ref_x = None
@@ -739,6 +771,7 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
             bl = _block_list() if _block_list is not None else None
             if not bl or not _hook_block_params:
                 return
+
             def _extract_tensor(out):
                 if isinstance(out, dict):
                     return out.get("img") or out.get("hidden_states") or next(iter(out.values()), None), out
@@ -778,11 +811,71 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
                         return None
                 return _hook
 
+            def _make_incontext_hooks(block_idx, buf, lazy_ref, strength, gate, s_state):
+                """Pre+post hook pair for in-context conditioning.
+
+                Pre-hook: captures block input at sigma (0.85-0.95); prepends ref tokens
+                          to vx and extends v_pe with neutral-rotation entries.
+                Post-hook: crops the ref-token prefix from the block output's vx.
+                """
+                _n_prepended = [0]
+
+                def _pre(module, args, kwargs):
+                    sigma = s_state[0]
+                    # Capture block INPUT during mid-sigma window
+                    if 0.85 <= sigma <= 0.95:
+                        try:
+                            vx = args[0][0]
+                            buf[block_idx] = vx.detach().cpu()
+                        except Exception:
+                            pass
+                    # Inject: prepend ref tokens to extend self-attention context
+                    if lazy_ref is None or lazy_ref.tensor is None:
+                        _n_prepended[0] = 0
+                        return None
+                    effective = _sigma_gated_strength(strength, sigma, gate[1], gate[0])
+                    if effective <= 0.0:
+                        _n_prepended[0] = 0
+                        return None
+                    try:
+                        vx, ax = args[0]
+                        ref = lazy_ref.tensor.to(vx.device, vx.dtype)
+                        ref_exp = ref.expand(vx.shape[0], -1, -1)
+                        new_vx = torch.cat([ref_exp, vx], dim=1)
+                        n_ref = ref.shape[1]
+                        _n_prepended[0] = n_ref
+                        new_kwargs = _extend_v_pe(kwargs, n_ref)
+                        return ((new_vx, ax),) + args[1:], new_kwargs
+                    except Exception:
+                        _n_prepended[0] = 0
+                        return None
+
+                def _post(module, inp, out):
+                    n = _n_prepended[0]
+                    if n == 0:
+                        return None
+                    try:
+                        vx, ax = out
+                        return (vx[:, n:], ax)
+                    except Exception:
+                        return None
+
+                return _pre, _post
+
             for idx, (buf, lazy_ref, strength, gate) in _hook_block_params.items():
-                h = bl[idx].register_forward_hook(
-                    _make_hook(idx, buf, lazy_ref, strength, gate, _state)
-                )
-                _active_handles.append(h)
+                is_identity = idx in IDENTITY_BLOCKS
+                if is_identity and has_i2v:
+                    # In-context conditioning: video tokens attend to reference tokens
+                    pre, post = _make_incontext_hooks(idx, buf, lazy_ref, strength, gate, _state)
+                    h1 = bl[idx].register_forward_pre_hook(pre, with_kwargs=True)
+                    h2 = bl[idx].register_forward_hook(post)
+                    _active_handles.append(h1)
+                    _active_handles.append(h2)
+                else:
+                    h = bl[idx].register_forward_hook(
+                        _make_hook(idx, buf, lazy_ref, strength, gate, _state)
+                    )
+                    _active_handles.append(h)
 
         def _remove_hooks():
             for h in _active_handles:

model_management.py

@@ -16,6 +16,9 @@ from aiohttp import web
 from server import PromptServer
 
 
+import comfy.model_management
+import comfy.clip_model
+
 LORA_TYPES = ["general", "action", "style", "quality", "character"]
 LORA_STACK_TYPE = "FUNPACK_LORA_STACK"
 TRANSFORMER_BLOCK_PATTERN = re.compile(r"(?:^|\.)transformer_blocks\.(\d+)\.")
@@ -184,6 +187,70 @@ def patch_energy(value):
     return 0.0
 
 
+def _load_gemma3_vision_weights(clip, path):
+    """Load vision_tower and multi_modal_projector from a Google Gemma3-12B-IT checkpoint.
+
+    Google's HuggingFace keys:
+      model.vision_tower.*          → vision_model.*   (into Gemma3_12B.vision_model)
+      model.multi_modal_projector.* → multi_modal_projector.*   (into Gemma3_12B.multi_modal_projector)
+    """
+    try:
+        sd = comfy.utils.load_torch_file(path, safe_load=True)
+    except Exception as e:
+        print(f"[FunPackGemmaLoader] Could not open checkpoint for vision weights: {e}")
+        return
+
+    vision_sd = {}
+    for k, v in sd.items():
+        if k.startswith("model.vision_tower."):
+            vision_sd[k[len("model.vision_tower."):]] = v
+        elif k.startswith("model.multi_modal_projector."):
+            vision_sd[k[len("model."):]] = v
+
+    if not vision_sd:
+        print("[FunPackGemmaLoader] No vision_tower / multi_modal_projector keys found — vision prompting unavailable")
+        return
+
+    try:
+        transformer = clip.cond_stage_model.gemma3_12b.transformer
+        missing, unexpected = transformer.load_state_dict(vision_sd, strict=False)
+        loaded = len(vision_sd) - len(unexpected)
+        print(f"[FunPackGemmaLoader] Vision weights loaded: {loaded}/{len(vision_sd)} tensors "
+              f"({len(missing)} missing, {len(unexpected)} unexpected)")
+        transformer._vision_loaded = True
+    except Exception as e:
+        print(f"[FunPackGemmaLoader] Vision weight injection failed: {e}")
+
+
+class FunPackGemmaLoader:
+    """Drop-in replacement for 'Load CLIP Model' that also loads the SigLIP vision tower
+    bundled in Google's Gemma3-12B-IT checkpoint, enabling reference-image vision prompting
+    inside Studio without a separate CLIP Vision model.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model_name": (folder_paths.get_filename_list("text_encoders"),),
+            }
+        }
+
+    RETURN_TYPES = ("CLIP",)
+    RETURN_NAMES = ("clip",)
+    FUNCTION = "load"
+    CATEGORY = "FunPack"
+
+    def load(self, model_name):
+        path = folder_paths.get_full_path("text_encoders", model_name)
+        clip = comfy.sd.load_clip(
+            ckpt_paths=[path],
+            clip_type=comfy.sd.CLIPType.LTXV,
+        )
+        _load_gemma3_vision_weights(clip, path)
+        return (clip,)
+
+
 class FunPackApplyLoraWeights:
     """
     Builds a LoRA stack from user base weights, then applies prompt-specific