diff --git a/__init__.py b/__init__.py index ec72ca82..4d58fb1f 100644 --- a/__init__.py +++ b/__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", diff --git a/conditioning.py b/conditioning.py index a16abb02..314b8629 100644 --- a/conditioning.py +++ b/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): diff --git a/ltx_enhancements.py b/ltx_enhancements.py index b7ad7770..92877f84 100644 --- a/ltx_enhancements.py +++ b/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: diff --git a/model_management.py b/model_management.py index 06ad323f..36d8395f 100644 --- a/model_management.py +++ b/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