from __future__ import annotations from dataclasses import dataclass import torch import torch.nn.functional as F from torch import nn from .configuration import HrmPenguinConfig from .recurrent.hrm_multimodal import MultimodalHrmTextModel @dataclass class HrmPenguinOutput: logits: torch.Tensor | None past_key_values: object | None = None class HrmPenguinForConditionalGeneration(nn.Module): def __init__(self, vision_encoder, language_model, projector, config: HrmPenguinConfig): super().__init__() self.vision_encoder = vision_encoder self.language_model = language_model self.projector = projector self.config = config self.last_effective_depth_metrics: dict[str, torch.Tensor] = {} @property def recurrent_model(self) -> MultimodalHrmTextModel: return self.language_model.model @staticmethod def _replace_visual_tokens( text_embeddings: torch.Tensor, visual_embeddings: torch.Tensor, visual_mask: torch.Tensor, image_batch_indices: torch.LongTensor | None = None, ) -> torch.Tensor: counts = visual_mask.sum(dim=1) visual_rows = (counts > 0).nonzero(as_tuple=False).flatten() if image_batch_indices is None: image_batch_indices = visual_rows image_batch_indices = image_batch_indices.to(device=visual_mask.device, dtype=torch.long) if not torch.equal(visual_rows, image_batch_indices): raise ValueError( f"visual rows {visual_rows.tolist()} do not match image batch indices " f"{image_batch_indices.tolist()}" ) if visual_embeddings.ndim == 2: if int(counts[visual_rows].sum()) != visual_embeddings.shape[0]: raise ValueError( f"visual_mask counts {counts.tolist()} do not match packed visual " f"embeddings {tuple(visual_embeddings.shape)}" ) elif visual_embeddings.ndim == 3: if visual_rows.numel() != visual_embeddings.shape[0] or not torch.all( counts[visual_rows] == visual_embeddings.shape[1] ): raise ValueError( f"visual_mask counts {counts.tolist()} do not match visual embeddings " f"{tuple(visual_embeddings.shape)}" ) else: raise ValueError( "visual_embeddings must be packed [sum(N_i), D] or dense [B, N, D]" ) expanded_mask = visual_mask.unsqueeze(-1).expand_as(text_embeddings) replacement = torch.zeros_like(text_embeddings).masked_scatter( expanded_mask, visual_embeddings.to(text_embeddings.dtype).reshape(-1) ) return torch.where(expanded_mask, replacement, text_embeddings) def encode_images(self, pixel_values, grid_sizes, merge_sizes) -> torch.Tensor: vision_parameter = next(self.vision_encoder.parameters()) pixel_values = pixel_values.to(device=vision_parameter.device, dtype=vision_parameter.dtype) grid_sizes = grid_sizes.to(vision_parameter.device) merge_sizes = merge_sizes.to(vision_parameter.device) with torch.no_grad(): features = self.vision_encoder(pixel_values, grid_sizes, merge_sizes) return self.projector(features) def forward( self, input_ids: torch.LongTensor, attention_mask: torch.Tensor, token_type_ids: torch.LongTensor, position_ids: torch.LongTensor, rope_position_ids: torch.LongTensor, visual_mask: torch.BoolTensor, instruction_mask: torch.BoolTensor, image_batch_indices: torch.LongTensor | None = None, visual_token_counts: torch.LongTensor | None = None, grid_heights: torch.LongTensor | None = None, grid_widths: torch.LongTensor | None = None, pixel_values: torch.Tensor | None = None, grid_sizes: torch.Tensor | None = None, merge_sizes: torch.Tensor | None = None, vision_features: torch.Tensor | None = None, collect_visual_diagnostics: bool = False, collect_depth_diagnostics: bool = False, ) -> HrmPenguinOutput: mask_counts = visual_mask.sum(dim=1) if visual_token_counts is not None and not torch.equal( mask_counts, visual_token_counts.to(mask_counts.device) ): raise ValueError("visual_mask does not match visual_token_counts") if grid_heights is not None and grid_widths is not None: expected_counts = grid_heights.to(mask_counts.device) * grid_widths.to(mask_counts.device) if not torch.equal(mask_counts, expected_counts): raise ValueError("visual token counts do not match dynamic grid areas") text_embeddings = self.language_model.get_input_embeddings()(input_ids) if bool(visual_mask.any()): if vision_features is None: if pixel_values is None or grid_sizes is None or merge_sizes is None: raise ValueError( "pixel_values/grid_sizes/merge_sizes or vision_features are required" ) vision_features = self.encode_images(pixel_values, grid_sizes, merge_sizes) # encode_images returns projected values for the regular path. visual_embeddings = vision_features else: visual_embeddings = self.projector(vision_features) inputs_embeds = self._replace_visual_tokens( text_embeddings, visual_embeddings, visual_mask, image_batch_indices=image_batch_indices, ) else: inputs_embeds = text_embeddings outputs = self.recurrent_model( inputs_embeds=inputs_embeds, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, rope_position_ids=rope_position_ids, visual_mask=visual_mask, instruction_mask=instruction_mask, use_cache=False, collect_visual_diagnostics=collect_visual_diagnostics, collect_depth_diagnostics=collect_depth_diagnostics, ) logits = None self.last_effective_depth_metrics = {} diagnostics = self.recurrent_model.last_visual_diagnostics if ( collect_depth_diagnostics and diagnostics is not None and diagnostics.logit_lens_response_hidden and diagnostics.logit_lens_response_hidden[-1].numel() ): with torch.no_grad(): final_hidden = diagnostics.logit_lens_response_hidden[-1] final_log_probs = F.log_softmax( self.language_model.lm_head(final_hidden).float(), dim=-1 ) final_probs = final_log_probs.exp() for name, hidden in zip( diagnostics.logit_lens_names, diagnostics.logit_lens_response_hidden, strict=True, ): if hidden.shape != final_hidden.shape: continue intermediate_log_probs = F.log_softmax( self.language_model.lm_head(hidden).float(), dim=-1 ) self.last_effective_depth_metrics[f"{name}_kl_to_final"] = ( F.kl_div( intermediate_log_probs, final_probs, reduction="batchmean", log_target=False, ) ) logits = self.language_model.lm_head(outputs.last_hidden_state) return HrmPenguinOutput( logits=logits, past_key_values=outputs.past_key_values, )