"""HydraGemma4: one model, many heads. model = HydraGemma4.from_pretrained( base_model="google/gemma-4-E2B-it", adapter_path="./colgemma4_adapter", ) # Retrieval -- LoRA ON, bidirectional attention embeddings = model.embed(images) # [B, N, 128] # Generation -- LoRA OFF, causal attention text = model.generate(image, "Describe this document.") One ColGemma4 + LoRA + lm_head. Attention mode and LoRA toggled per call. Gemma 4 E2B layer structure: 35 layers total: 28 sliding_attention + 7 full_attention full_attention at indices: [4, 9, 14, 19, 24, 29, 34] Only full_attention layers are patched for bidirectional mode. Sliding layers stay causal always. Key Gemma 4 differences from Qwen3.5: - Layer forward signature includes `per_layer_input` (Per-Layer Embeddings) - Layer type via config.text_config.layer_types list (not layer.layer_type attribute) - Vision tower at self.vision_tower (not self.visual) - final_logit_softcapping for logit capping - image_position_ids (not image_grid_thw) """ import re from typing import List, Union import torch import torch.nn as nn from PIL import Image from peft import PeftModel from transformers import AutoProcessor, Gemma4ForConditionalGeneration from transformers.models.gemma4.modeling_gemma4 import Gemma4Model as _Gemma4Base from colgemma4 import ColGemma4, ColGemma4Processor class HydraGemma4: """One ColGemma4 model with LoRA + lm_head. Two faces via toggling. Retrieval face: LoRA ON, bidirectional attention -> custom_text_proj -> 128-dim Generation face: LoRA OFF, causal attention -> lm_head -> autoregressive text """ def __init__(self, model, lm_head, emb_processor, gen_processor, attn_fns, config): self.model = model self.lm_head = lm_head self.emb_processor = emb_processor self.gen_processor = gen_processor self._attn_fns = attn_fns # {layer_idx: (causal_fn, bidir_fn)} self._config = config self._base = self._unwrap(model) @staticmethod def _unwrap(model): m = model if hasattr(m, "module"): m = m.module if hasattr(m, "base_model"): m = m.base_model if hasattr(m, "model"): m = m.model return m def _set_bidirectional(self): for idx, (_, bidir_fn) in self._attn_fns.items(): self._base.language_model.layers[idx].forward = bidir_fn def _set_causal(self): for idx, (causal_fn, _) in self._attn_fns.items(): self._base.language_model.layers[idx].forward = causal_fn @classmethod def from_pretrained( cls, base_model: str = "google/gemma-4-E2B-it", adapter_path: str = "/tmp/colgemma4_adapter", torch_dtype=torch.bfloat16, max_visual_tokens: int = 560, device: str = "cuda", ): """Load single model from base + adapter + lm_head.""" from pathlib import Path print(f"Loading ColGemma4 from {base_model}...", flush=True) config = Gemma4ForConditionalGeneration.config_class.from_pretrained( base_model, trust_remote_code=True ) config.text_config.use_cache = False model = ColGemma4.from_pretrained( base_model, config=config, torch_dtype=torch_dtype, attn_implementation="sdpa", ignore_mismatched_sizes=True, ) # Store causal/bidirectional forwards for full-attention layers layer_types = config.text_config.layer_types attn_fns = {} for idx, layer in enumerate(model.language_model.layers): if layer_types[idx] == "full_attention": causal_fn = layer.forward def make_bidir(orig): def bidir( hidden_states, per_layer_input=None, position_embeddings=None, attention_mask=None, **kw, ): if ( attention_mask is not None and attention_mask.ndim == 4 and attention_mask.dtype.is_floating_point ): min_dtype = torch.finfo(attention_mask.dtype).min diag = torch.diagonal(attention_mask, dim1=-2, dim2=-1) is_valid = diag > (min_dtype / 2) bidir_mask = torch.where( is_valid.unsqueeze(-1) & is_valid.unsqueeze(-2), attention_mask.new_zeros(1), attention_mask.new_full((1,), min_dtype), ) attention_mask = bidir_mask return orig( hidden_states, per_layer_input=per_layer_input, position_embeddings=position_embeddings, attention_mask=attention_mask, **kw, ) return bidir bidir_fn = make_bidir(causal_fn) attn_fns[idx] = (causal_fn, bidir_fn) layer.forward = bidir_fn print(f"Patched {len(attn_fns)} full-attention layers", flush=True) print(f"Loading LoRA adapter from {adapter_path}...", flush=True) model = PeftModel.from_pretrained(model, adapter_path) model = model.to(device).eval() # Load lm_head print("Loading lm_head...", flush=True) lm_head_path = Path(adapter_path) / "lm_head.pt" if lm_head_path.exists(): base_cfg = Gemma4ForConditionalGeneration.config_class.from_pretrained(base_model) lm_head = nn.Linear( base_cfg.text_config.hidden_size, base_cfg.text_config.vocab_size, bias=False, ) lm_head.load_state_dict(torch.load(lm_head_path, map_location="cpu")) else: base = Gemma4ForConditionalGeneration.from_pretrained( base_model, torch_dtype=torch_dtype ) lm_head = base.lm_head del base torch.cuda.empty_cache() lm_head = lm_head.to(device).to(torch_dtype) emb_processor = ColGemma4Processor.from_pretrained( base_model, max_num_visual_tokens=max_visual_tokens ) gen_processor = AutoProcessor.from_pretrained(base_model, trust_remote_code=True) params = sum(p.numel() for p in model.parameters()) / 1e6 lm_params = sum(p.numel() for p in lm_head.parameters()) / 1e6 print( f"Ready. Single model: {params:.0f}M + lm_head: {lm_params:.0f}M", flush=True, ) return cls(model, lm_head, emb_processor, gen_processor, attn_fns, config) def embed(self, images: Union[Image.Image, List[Image.Image]]) -> torch.Tensor: """Embed images for retrieval. LoRA ON, bidirectional attention.""" if isinstance(images, Image.Image): images = [images] self.model.enable_adapter_layers() self._set_bidirectional() inputs = self.emb_processor.process_images(images) device = next(self.model.parameters()).device inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16): return self.model(**inputs) def embed_queries(self, queries: Union[str, List[str]]) -> torch.Tensor: """Embed queries for retrieval. LoRA ON, bidirectional attention.""" if isinstance(queries, str): queries = [queries] self.model.enable_adapter_layers() self._set_bidirectional() inputs = self.emb_processor.process_queries(queries) device = next(self.model.parameters()).device inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16): return self.model(**inputs) @torch.no_grad() def generate( self, image: Image.Image, prompt: str, max_new_tokens: int = 4096, system_prompt: str = None, temperature: float = 1.0, top_p: float = 0.95, top_k: int = 64, ) -> str: """Generate text given image + prompt. LoRA OFF, causal attention.""" self.model.disable_adapter_layers() self._set_causal() try: return self._generate_with_kv_cache( image, prompt, max_new_tokens, system_prompt, temperature=temperature, top_p=top_p, top_k=top_k, ) except Exception: return self._generate_no_cache( image, prompt, max_new_tokens, system_prompt, temperature=temperature, top_p=top_p, top_k=top_k, ) finally: self.model.enable_adapter_layers() self._set_bidirectional() def _prepare_inputs(self, image, prompt, system_prompt=None): msgs = [] if system_prompt: msgs.append({"role": "system", "content": system_prompt}) msgs.append({ "role": "user", "content": [ {"type": "image", "image": image}, {"type": "text", "text": prompt}, ], }) txt = self.gen_processor.apply_chat_template( msgs, tokenize=False, add_generation_prompt=True, ) inp = self.gen_processor( text=[txt], images=[image], return_tensors="pt", padding=True ) device = next(self.model.parameters()).device return {k: v.to(device) for k, v in inp.items()} def _decode(self, generated): if not generated: return "" gen_ids = torch.cat(generated, dim=1) return self.gen_processor.tokenizer.decode(gen_ids[0], skip_special_tokens=True).strip() def _is_eos(self, token_id): eos = self.gen_processor.tokenizer.eos_token_id if isinstance(eos, list): return token_id in eos return token_id == eos @staticmethod def _sample_token(logits, temperature, top_p, top_k): if temperature <= 0: return logits.argmax(dim=-1) logits_1d = logits[0, 0] / temperature if top_k > 0: topk_vals, _ = torch.topk(logits_1d, min(top_k, logits_1d.size(-1))) logits_1d[logits_1d < topk_vals[-1]] = float("-inf") probs = torch.softmax(logits_1d, dim=-1) if 0 < top_p < 1.0: sorted_probs, sorted_idx = torch.sort(probs, descending=True) cumsum = torch.cumsum(sorted_probs, dim=-1) mask = cumsum - sorted_probs > top_p sorted_probs[mask] = 0 probs = torch.zeros_like(probs).scatter_(0, sorted_idx, sorted_probs) probs = probs / probs.sum() token_id = torch.multinomial(probs, num_samples=1) return token_id.reshape(1, 1) def _generate_with_kv_cache( self, image, prompt, max_new_tokens, system_prompt=None, temperature=1.0, top_p=0.95, top_k=64, ): """KV-cache generation. Pixel values only on first step.""" inp = self._prepare_inputs(image, prompt, system_prompt) input_ids = inp["input_ids"] attn_mask = inp["attention_mask"] mm_ids = inp.get("mm_token_type_ids") past_key_values = None generated = [] final_softcap = self._config.text_config.final_logit_softcapping for step in range(max_new_tokens): if step == 0: kw = { "input_ids": input_ids, "attention_mask": attn_mask, "pixel_values": inp.get("pixel_values"), "image_position_ids": inp.get("image_position_ids"), "use_cache": True, "output_hidden_states": True, "return_dict": True, } if mm_ids is not None: kw["mm_token_type_ids"] = mm_ids else: kw = { "input_ids": next_token, "attention_mask": attn_mask, "past_key_values": past_key_values, "use_cache": True, "output_hidden_states": True, "return_dict": True, } outputs = _Gemma4Base.forward(self._base, **kw) past_key_values = outputs.past_key_values logits = self.lm_head(outputs.last_hidden_state[:, -1:, :]) if final_softcap is not None: logits = logits / final_softcap logits = torch.tanh(logits) logits = logits * final_softcap next_token = self._sample_token(logits, temperature, top_p, top_k) generated.append(next_token) if self._is_eos(next_token.item()): break attn_mask = torch.cat([attn_mask, torch.ones_like(next_token)], dim=1) return self._decode(generated) def _generate_no_cache( self, image, prompt, max_new_tokens, system_prompt=None, temperature=1.0, top_p=0.95, top_k=64, ): """No-cache fallback. Passes all inputs every step.""" inp = self._prepare_inputs(image, prompt, system_prompt) input_ids = inp["input_ids"] attn_mask = inp["attention_mask"] mm_ids = inp.get("mm_token_type_ids") generated = [] final_softcap = self._config.text_config.final_logit_softcapping for step in range(max_new_tokens): kw = { "input_ids": input_ids, "attention_mask": attn_mask, "pixel_values": inp.get("pixel_values"), "image_position_ids": inp.get("image_position_ids"), } if mm_ids is not None: kw["mm_token_type_ids"] = mm_ids hidden = _Gemma4Base.forward( self._base, **kw, use_cache=False, output_hidden_states=True, return_dict=True, ).last_hidden_state logits = self.lm_head(hidden[:, -1:, :]) if final_softcap is not None: logits = logits / final_softcap logits = torch.tanh(logits) logits = logits * final_softcap next_token = self._sample_token(logits, temperature, top_p, top_k) generated.append(next_token) if self._is_eos(next_token.item()): break input_ids = torch.cat([input_ids, next_token], dim=1) attn_mask = torch.cat([attn_mask, torch.ones_like(next_token)], dim=1) if mm_ids is not None: mm_ids = torch.cat([mm_ids, torch.zeros_like(next_token)], dim=1) return self._decode(generated)