Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/cat_tabby.jpg filter=lfs diff=lfs merge=lfs -text
+examples/dog.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,51 @@
 ---
-title: Loopvl
-emoji: 🐠
-colorFrom: pink
-colorTo: gray
+title: LoopVL
+emoji: 🔁
+colorFrom: purple
+colorTo: yellow
 sdk: gradio
 sdk_version: 6.30.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Recurrent VLM with a controllable compute budget
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# LoopVL: Recurrent Vision-Language Model
+
+**LoopVL** ([TierFlow/LoopVL](https://huggingface.co/TierFlow/LoopVL)) is a ~1.6B
+vision-language model that reuses its 16-layer transformer stack in nested loops
+instead of stacking 128 layers. The distinctive control this demo exposes is the
+**recurrent inference budget**: the H (high) and L (low) cycle counts decide how
+many times the stack is applied per forward pass — unrolled depth
+`16 × H × (L+1)` effective layers. The schedule can be changed at runtime without
+reloading weights.
+
+Ask any question about an image, pick a schedule, and compare answers across
+compute budgets.
+
+- **Model**: [TierFlow/LoopVL](https://huggingface.co/TierFlow/LoopVL) (BF16
+  vision/language cores + FP32 projector, visual gate and anchor)
+- **Code**: [Tier-Flow/LoopVL](https://github.com/Tier-Flow/LoopVL) — the
+  inference runtime (`hrm_penguin`, `penguin_encoder`, `modeling_loopvl.py`,
+  `metadata.py`) is vendored here from that Apache-2.0 repository; the Penguin
+  vision encoder keeps its Tencent `LICENSE.txt` alongside.
+- **Examples**: images from
+  [linoyts/repo-to-space-example-inputs](https://huggingface.co/datasets/linoyts/repo-to-space-example-inputs)
+  (free to use, CC0-style).
+
+## How it works
+
+1. The image is resized under a dynamic-resolution policy (256–2048 visual
+   tokens) and patched by the PenguinVL vision encoder.
+2. Visual tokens are injected into the text sequence; the HRM-Text recurrent
+   LM loops its stack `H × (L+1)` times per forward pass, re-injecting an
+   anchored version of the visual tokens after each high cycle, gated by a
+   learned saturation-aware gate.
+3. Decoding is greedy, recomputing the full sequence per token (no KV cache).
+
+## Reproducing the reference
+
+The `answer` endpoint follows the authors' `runtime/infer.py` one-to-one:
+same prompt template (`condition\nImage:\n` + placeholders + `\nInstruction:\n…`),
+same greedy decode, same `decode(skip_special_tokens=True).strip()` postprocess.

app.py

@@ -0,0 +1,253 @@
+"""LoopVL demo: a recurrent vision-language model with a user-controlled visual compute budget.
+
+Loads the public TierFlow/LoopVL checkpoint (~1.6B parameters) with the authors'
+single-file loader, then answers questions about an image. The distinctive control
+is the recurrent schedule: the model reuses its 16-layer transformer stack, and the
+H (high) and L (low) cycle counts decide how many times the stack is applied per
+forward pass — the "inference budget" the paper sweeps across.
+"""
+
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # MUST come before any torch / CUDA-touching import
+
+import time
+from PIL import Image, ImageFile
+
+import gradio as gr
+import torch
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+MODEL_ID = "TierFlow/LoopVL"
+DEVICE = "cuda"
+
+# Vendored from https://github.com/Tier-Flow/LoopVL (Apache-2.0; Penguin encoder
+# keeps its Tencent LICENSE.txt alongside).
+from modeling_loopvl import LoopVLForConditionalGeneration, load_tokenizer_and_processor
+from hrm_penguin.generation import greedy_generate
+from hrm_penguin.inference import build_inference_batch
+from hrm_penguin.recurrent.visual_anchor import VisualAnchor
+
+TRIGGER = "\nExplain your reasoning, then give the final answer."
+
+
+def _download_weights() -> str:
+    """Fetch config + tokenizer + single safetensors checkpoint into a local dir."""
+    from huggingface_hub import snapshot_download
+
+    return snapshot_download(
+        MODEL_ID,
+        allow_patterns=[
+            "config.json",
+            "generation_config.json",
+            "preprocessor_config.json",
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "model.safetensors",
+        ],
+    )
+
+
+# --- Load everything at module scope, eagerly on "cuda" (string device). ------
+_model_dir = _download_weights()
+model = LoopVLForConditionalGeneration.from_pretrained(_model_dir, device=DEVICE)
+tokenizer, processor = load_tokenizer_and_processor(_model_dir, model.config)
+
+# Base visual-anchor coefficients (H_cycles=2 in the shipped checkpoint). Any
+# schedule change rebuilds VisualAnchor, re-copying these values (clamped).
+# Kept on CPU: schedule changes run inside the ZeroGPU worker, where a plain
+# CPU tensor is trivially inherited across forks and copied cross-device.
+recurrent = model.recurrent_model
+BASE_ANCHOR_VALUES = recurrent.visual_anchor.raw_alpha.detach().to("cpu").clone()
+ANCHOR_MAX_SCALE = recurrent.visual_anchor.max_scale
+LAST_SCHEDULE = {"h": int(recurrent.config.H_cycles), "l": int(recurrent.config.L_cycles)}
+
+
+def _apply_schedule(h_cycles: int, l_cycles: int) -> None:
+    """Set the recurrent inference schedule (the authors' configure() pattern).
+
+    The forward pass reads config.H_cycles / config.L_cycles each call, and the
+    visual anchor holds one learned coefficient per high cycle, so changing H
+    rebuilds it and re-copies the base coefficients (clamped, as in the paper).
+    """
+    h_cycles, l_cycles = int(h_cycles), int(l_cycles)
+    if h_cycles < 1 or l_cycles < 1:
+        raise ValueError("H and L cycle counts must be positive")
+    if LAST_SCHEDULE["h"] == h_cycles and LAST_SCHEDULE["l"] == l_cycles:
+        return  # already configured; avoids rebuilding the anchor needlessly
+    recurrent.config.H_cycles = h_cycles
+    recurrent.config.L_cycles = l_cycles
+    old = recurrent.visual_anchor
+    new = VisualAnchor(h_cycles, max_scale=ANCHOR_MAX_SCALE).to(
+        device=old.raw_alpha.device, dtype=old.raw_alpha.dtype
+    )
+    with torch.no_grad():
+        for i in range(h_cycles):
+            new.raw_alpha[i].copy_(
+                BASE_ANCHOR_VALUES[min(i, len(BASE_ANCHOR_VALUES) - 1)]
+            )
+    recurrent.visual_anchor = new
+    LAST_SCHEDULE["h"], LAST_SCHEDULE["l"] = h_cycles, l_cycles
+
+
+# Preset schedules from the paper's Table 4 sweep: (H, L) -> unrolled depth
+# 16 * H * (L + 1). H2L3 (the trained schedule) is the default.
+SCHEDULES = [
+    ("H1L1 (32 layers, fastest)", 1, 1),
+    ("H1L3 (64 layers)", 1, 3),
+    ("H2L1 (64 layers)", 2, 1),
+    ("H2L3 (128 layers, default)", 2, 3),
+    ("H2L5 (192 layers)", 2, 5),
+    ("H4L3 (256 layers)", 4, 3),
+]
+_SCHEDULE_LOOKUP = {name: (h, l) for name, h, l in SCHEDULES}
+
+
+def _estimate_duration(*args, **kwargs) -> int:
+    """Callable GPU duration: base time grows with unrolled depth and token budget."""
+    schedule = kwargs.get("schedule") or (args[2] if len(args) > 2 else None)
+    budget = kwargs.get("max_tokens") or (args[3] if len(args) > 3 else 32)
+    try:
+        h, l = _SCHEDULE_LOOKUP.get(str(schedule), (2, 3))
+    except Exception:
+        h, l = 2, 3
+    depth = 16 * h * (l + 1)
+    return min(300, 20 + int(depth * 0.09) + int(budget * 0.35))
+
+
+@spaces.GPU(duration=_estimate_duration)
+def answer(
+    image: Image.Image,
+    prompt: str,
+    schedule: str = "H2L3 (128 layers, default)",
+    max_tokens: int = 32,
+    use_trigger: bool = False,
+    progress=gr.Progress(track_tqdm=True),
+) -> str:
+    """Answer a question about an image using LoopVL, a recurrent vision-language model.
+
+    Args:
+        image: the input image to reason about.
+        prompt: the question or instruction about the image.
+        schedule: recurrent compute schedule (H and L cycle counts). Higher means
+            more effective transformer layers per forward pass.
+        max_tokens: maximum number of tokens to generate.
+        use_trigger: append the authors' reasoning trigger to the instruction.
+        progress: Gradio progress bar.
+
+    Returns:
+        The model's answer text.
+    """
+    if image is None:
+        return "Please provide an image."
+    prompt = (prompt or "").strip()
+    if not prompt:
+        return "Please provide a question or instruction about the image."
+
+    h, l = _SCHEDULE_LOOKUP.get(schedule, (2, 3))
+    _apply_schedule(h, l)
+
+    instruction = prompt + (TRIGGER if use_trigger else "")
+    pil_image = image if isinstance(image, Image.Image) else Image.fromarray(image)
+    batch = build_inference_batch(
+        tokenizer,
+        processor,
+        pil_image.convert("RGB"),
+        instruction,
+        model.config,
+        torch.device(DEVICE),
+    )
+    max_tokens = max(1, min(int(max_tokens), 96))
+    start = time.perf_counter()
+    with torch.inference_mode():
+        generated = greedy_generate(
+            model,
+            batch,
+            max_new_tokens=max_tokens,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+    elapsed = time.perf_counter() - start
+    prediction = tokenizer.decode(generated[0], skip_special_tokens=True).strip()
+    depth = 16 * h * (l + 1)
+    label = f"Answer (schedule {schedule.split(' ')[0]} → {depth} effective layers, {elapsed:.1f}s):"
+    return f"{label}\n{prediction}" if prediction else f"{label}\n(no tokens generated)"
+
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """
+            # LoopVL: Recurrent Vision-Language Model
+
+            **LoopVL** (~1.6B params) reuses its 16-layer transformer stack in loops
+            instead of stacking 128 layers. You control the compute budget: pick an
+            **H/L schedule** to decide how many times the stack runs per forward pass,
+            then ask anything about the image.
+
+            [Model](https://huggingface.co/TierFlow/LoopVL) ·
+            [Code](https://github.com/Tier-Flow/LoopVL)
+            """
+        )
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_in = gr.Image(
+                    type="pil", label="Image", sources=["upload", "clipboard"]
+                )
+                prompt_in = gr.Textbox(
+                    label="Question / instruction",
+                    placeholder="What is in the image?",
+                    lines=2,
+                )
+                schedule_in = gr.Dropdown(
+                    choices=[name for name, _, _ in SCHEDULES],
+                    value="H2L3 (128 layers, default)",
+                    label="Recurrent schedule (compute budget)",
+                    info="Unrolled depth = 16 × H × (L+1) effective layers per forward pass",
+                )
+                run_btn = gr.Button("Answer", variant="primary")
+                with gr.Accordion("Advanced", open=False):
+                    max_tokens_in = gr.Slider(
+                        minimum=8, maximum=96, step=8, value=32,
+                        label="Max new tokens",
+                    )
+                    trigger_in = gr.Checkbox(
+                        value=False,
+                        label="Append reasoning trigger",
+                        info='Adds "Explain your reasoning, then give the final answer." '
+                             "(the authors' math-reasoning prompt)",
+                    )
+            with gr.Column(scale=1):
+                answer_out = gr.Textbox(label="Answer", lines=8, show_copy_button=True)
+
+        gr.Examples(
+            examples=[
+                ["examples/cat_tabby.jpg", "What animal is this, and what is it doing?"],
+                ["examples/rubber_ducks.jpg", "How many objects are in this image? Describe them."],
+                ["examples/landscape.jpg", "Describe the scene and the weather."],
+                ["examples/dog.jpg", "What breed of dog does this look like?"],
+            ],
+            inputs=[image_in, prompt_in],
+            outputs=answer_out,
+            fn=answer,
+            cache_examples=True,
+            cache_mode="lazy",
+            label="Examples (click to load, then press Answer)",
+        )
+
+    run_btn.click(
+        answer,
+        inputs=[image_in, prompt_in, schedule_in, max_tokens_in, trigger_in],
+        outputs=answer_out,
+        api_name="generate",
+        concurrency_limit=1,
+    )
+
+demo.launch(mcp_server=True)

hrm_penguin/__init__.py

@@ -0,0 +1,7 @@
+"""Penguin-Encoder x HRM-Text recurrent VLM."""
+
+from .configuration import HrmPenguinConfig
+
+__all__ = ["HrmPenguinConfig"]
+__version__ = "0.1.0"
+

hrm_penguin/configuration.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field, fields
+from typing import Any
+
+@dataclass
+class VisionConfig:
+    model_path: str = "runtime/models/Penguin-Encoder"
+    hidden_size: int = 1024
+    patch_size: int = 14
+    input_size: int = 336
+    grid_height: int = 24
+    grid_width: int = 24
+    visual_tokens: int = 576
+    dynamic_resolution: bool = False
+    min_visual_tokens: int = 256
+    max_visual_tokens: int = 1024
+    attn_implementation: str = "sdpa"
+
+
+@dataclass
+class LanguageConfig:
+    model_path: str = "runtime/models/HRM-Text-1B"
+    hidden_size: int = 1536
+    num_attention_heads: int = 12
+    head_dim: int = 128
+    H_cycles: int = 2
+    L_cycles: int = 3
+    rope_theta: float = 10000.0
+    rope_parameters: dict[str, Any] = field(default_factory=dict)
+    max_position_embeddings: int = 4096
+    attn_implementation: str = "sdpa"
+
+
+@dataclass
+class RecurrentVisualConfig:
+    rope_mode: str = "hybrid_2d"
+    use_visual_anchor: bool = True
+    use_visual_gate: bool = True
+    gate_dim: int = 256
+    anchor_raw_init: float = -2.0
+    anchor_max_scale: float = 0.5
+    gate_saturation_metric_lower: float = 0.05
+    gate_saturation_metric_upper: float = 1.95
+
+    def validate(self) -> None:
+        if self.rope_mode not in {"original_1d", "hybrid_2d"}:
+            raise ValueError(f"Unsupported rope_mode: {self.rope_mode}")
+        if not 0 <= self.gate_saturation_metric_lower < 1:
+            raise ValueError("Invalid lower Gate saturation metric bound")
+        if not 1 < self.gate_saturation_metric_upper <= 2:
+            raise ValueError("Invalid upper Gate saturation metric bound")
+
+
+@dataclass
+class InferenceConfig:
+    max_seq_len: int = 1024
+    use_cache: bool = False
+
+
+@dataclass
+class HrmPenguinConfig:
+    vision: VisionConfig = field(default_factory=VisionConfig)
+    language: LanguageConfig = field(default_factory=LanguageConfig)
+    recurrent_visual: RecurrentVisualConfig = field(default_factory=RecurrentVisualConfig)
+    inference: InferenceConfig = field(default_factory=InferenceConfig)
+
+    def validate(self) -> None:
+        self.recurrent_visual.validate()
+        if self.vision.input_size % self.vision.patch_size:
+            raise ValueError("input_size must be divisible by patch_size")
+        expected = self.vision.grid_height * self.vision.grid_width
+        if expected != self.vision.visual_tokens:
+            raise ValueError("visual token count does not match grid")
+        if self.vision.min_visual_tokens < 1:
+            raise ValueError("min_visual_tokens must be positive")
+        if self.vision.max_visual_tokens < self.vision.min_visual_tokens:
+            raise ValueError("max_visual_tokens must be >= min_visual_tokens")
+        if (
+            self.vision.dynamic_resolution
+            and self.inference.max_seq_len <= self.vision.max_visual_tokens + 2
+        ):
+            raise ValueError(
+                "Dynamic resolution max_seq_len must leave room for instruction and response"
+            )
+        if self.inference.max_seq_len > self.language.max_position_embeddings:
+            raise ValueError(
+                "max_seq_len exceeds the HRM language model position limit"
+            )
+        if self.language.head_dim % 2:
+            raise ValueError("head_dim must be even for hybrid RoPE")
+
+    @classmethod
+    def from_dict(cls, raw: dict[str, Any]) -> "HrmPenguinConfig":
+        """Read the model architecture and inference settings from embedded metadata."""
+        def component(kind, name, defaults=None):
+            values = dict(defaults or {})
+            values.update(raw.get(name, {}))
+            allowed = {item.name for item in fields(kind)}
+            return kind(**{key: value for key, value in values.items() if key in allowed})
+
+        language = component(LanguageConfig, "language")
+        cfg = cls(
+            vision=component(VisionConfig, "vision"),
+            language=language,
+            recurrent_visual=component(RecurrentVisualConfig, "recurrent_visual"),
+            inference=component(
+                InferenceConfig, "inference",
+                {"max_seq_len": language.max_position_embeddings},
+            ),
+        )
+        cfg.validate()
+        return cfg
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)

hrm_penguin/generation.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import torch
+
+from .recurrent.positions import build_compact_positions
+
+
+@torch.no_grad()
+def greedy_generate(
+    model,
+    batch: dict[str, torch.Tensor],
+    max_new_tokens: int = 64,
+    eos_token_id: int | None = None,
+) -> torch.LongTensor:
+    """Correctness-first generation that recomputes the recurrent sequence each step."""
+    if batch["input_ids"].shape[0] != 1:
+        raise ValueError("greedy_generate currently supports batch_size=1")
+    model.eval()
+    input_ids = batch["input_ids"]
+    header_length = int(batch["header_length"])
+    instruction_length = int(batch["instruction_length"])
+    visual_token_count = int(batch["visual_token_counts"][0])
+    grid_height = int(batch["grid_heights"][0])
+    grid_width = int(batch["grid_widths"][0])
+    visual_embeddings = model.encode_images(
+        batch["pixel_values"], batch["grid_sizes"], batch["merge_sizes"]
+    )
+    generated: list[torch.Tensor] = []
+    for _ in range(max_new_tokens):
+        response_length = input_ids.shape[1] - header_length - visual_token_count - instruction_length
+        mask_positions, rope_positions, visual_mask, instruction_mask = build_compact_positions(
+            torch.tensor([header_length], device=input_ids.device),
+            torch.tensor([instruction_length], device=input_ids.device),
+            torch.tensor([response_length], device=input_ids.device),
+            visual_token_counts=torch.tensor([visual_token_count], device=input_ids.device),
+            grid_heights=torch.tensor([grid_height], device=input_ids.device),
+            grid_widths=torch.tensor([grid_width], device=input_ids.device),
+            grid_h=model.config.vision.grid_height,
+            grid_w=model.config.vision.grid_width,
+            rope_mode=model.config.recurrent_visual.rope_mode,
+        )
+        text_embeddings = model.language_model.get_input_embeddings()(input_ids)
+        inputs_embeds = model._replace_visual_tokens(text_embeddings, visual_embeddings, visual_mask)
+        prefix_length = header_length + visual_token_count + instruction_length
+        token_type_ids = torch.zeros_like(input_ids)
+        token_type_ids[:, :prefix_length] = 1
+        outputs = model.recurrent_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=torch.ones_like(input_ids),
+            token_type_ids=token_type_ids,
+            position_ids=mask_positions,
+            rope_position_ids=rope_positions,
+            visual_mask=visual_mask,
+            instruction_mask=instruction_mask,
+            use_cache=False,
+        )
+        next_token = model.language_model.lm_head(outputs.last_hidden_state[:, -1]).argmax(dim=-1)
+        generated.append(next_token)
+        input_ids = torch.cat((input_ids, next_token[:, None]), dim=1)
+        if eos_token_id is not None and bool((next_token == eos_token_id).all()):
+            break
+    return torch.stack(generated, dim=1) if generated else input_ids[:, :0]

hrm_penguin/inference.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import torch
+from .prompting import PromptParts, tokenize_without_special_tokens
+
+
+def build_inference_batch(tokenizer, image_processor, image, instruction, config, device):
+    prompt = PromptParts()
+    header = tokenize_without_special_tokens(tokenizer, prompt.header)
+    instruction_ids = tokenize_without_special_tokens(
+        tokenizer, prompt.instruction_prefix + instruction + prompt.instruction_suffix
+    )
+    placeholder = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
+    processed = image_processor(image, return_tensors="pt")
+    grid = processed["grid_sizes"][0].long()
+    merge = int(processed["merge_sizes"][0])
+    if int(grid[0]) != 1 or int(grid[1]) % merge or int(grid[2]) % merge:
+        raise ValueError(f"Unsupported inference grid: {grid.tolist()}, merge={merge}")
+    grid_height = int(grid[1]) // merge
+    grid_width = int(grid[2]) // merge
+    visual_tokens = grid_height * grid_width
+    input_ids = torch.tensor(
+        [header + [placeholder] * visual_tokens + instruction_ids],
+        dtype=torch.long,
+        device=device,
+    )
+    return {
+        "input_ids": input_ids,
+        "header_length": len(header),
+        "instruction_length": len(instruction_ids),
+        "visual_token_counts": torch.tensor([visual_tokens], device=device),
+        "grid_heights": torch.tensor([grid_height], device=device),
+        "grid_widths": torch.tensor([grid_width], device=device),
+        "pixel_values": processed["pixel_values"].to(device),
+        "grid_sizes": processed["grid_sizes"].to(device),
+        "merge_sizes": processed["merge_sizes"].to(device),
+    }

hrm_penguin/model.py

@@ -0,0 +1,180 @@
+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,
+        )

hrm_penguin/prompting.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class PromptParts:
+    header: str = "<|im_start|>condition\nImage:\n"
+    text_header: str = "<|im_start|>condition\n"
+    instruction_prefix: str = "\nInstruction:\n"
+    instruction_suffix: str = "<|im_end|>\n"
+
+
+def tokenize_without_special_tokens(tokenizer, text: str) -> list[int]:
+    return tokenizer(text, add_special_tokens=False)["input_ids"]

hrm_penguin/recurrent/__init__.py

@@ -0,0 +1,11 @@
+from .hybrid_rope import HybridHrmRotaryEmbedding, apply_rotary_pos_emb
+from .visual_anchor import VisualAnchor
+from .visual_gate import LoopAwareVisualGate
+
+__all__ = [
+    "HybridHrmRotaryEmbedding",
+    "LoopAwareVisualGate",
+    "VisualAnchor",
+    "apply_rotary_pos_emb",
+]
+

hrm_penguin/recurrent/hrm_multimodal.py

@@ -0,0 +1,408 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.models.hrm_text.modeling_hrm_text import HrmTextModel
+
+from .hybrid_rope import HybridHrmRotaryEmbedding
+from .visual_anchor import VisualAnchor
+from .visual_gate import LoopAwareVisualGate, masked_mean
+
+
+@dataclass
+class RecurrentVisualDiagnostics:
+    gate_mean: list[torch.Tensor]
+    gate_std: list[torch.Tensor]
+    gate_min: list[torch.Tensor]
+    gate_max: list[torch.Tensor]
+    gate_entropy: list[torch.Tensor]
+    gate_saturation_rate: list[torch.Tensor]
+    alpha: list[torch.Tensor]
+    reinjection_norm: list[torch.Tensor]
+    current_visual_norm: list[torch.Tensor]
+    reinjection_ratio: list[torch.Tensor]
+    gate_maps: list[list[torch.Tensor]]
+    transition_names: list[str]
+    transition_delta_rms: list[torch.Tensor]
+    transition_relative_delta: list[torch.Tensor]
+    transition_cosine: list[torch.Tensor]
+    transition_visual_delta_rms: list[torch.Tensor]
+    transition_visual_cosine: list[torch.Tensor]
+    transition_response_delta_rms: list[torch.Tensor]
+    transition_response_cosine: list[torch.Tensor]
+    logit_lens_names: list[str]
+    logit_lens_response_hidden: list[torch.Tensor]
+
+    @classmethod
+    def empty(cls) -> "RecurrentVisualDiagnostics":
+        return cls(
+            [], [], [], [], [], [], [], [], [], [], [],
+            [], [], [], [], [], [], [], [], [], [],
+        )
+
+    @staticmethod
+    def _state_change(
+        before: torch.Tensor,
+        after: torch.Tensor,
+        mask: torch.BoolTensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        before_selected = before.detach()[mask].float()
+        after_selected = after.detach()[mask].float()
+        if not before_selected.numel():
+            nan = after.new_tensor(float("nan"), dtype=torch.float32)
+            return nan, nan, nan
+        delta = after_selected - before_selected
+        delta_rms = delta.square().mean().sqrt()
+        relative_delta = delta.norm() / before_selected.norm().clamp_min(1e-12)
+        cosine = F.cosine_similarity(before_selected, after_selected, dim=-1).mean()
+        return delta_rms, relative_delta, cosine
+
+    def record_transition(
+        self,
+        name: str,
+        before: torch.Tensor,
+        after: torch.Tensor,
+        valid_mask: torch.BoolTensor,
+        visual_mask: torch.BoolTensor | None,
+        response_mask: torch.BoolTensor,
+        *,
+        logit_lens_token_limit: int = 4,
+    ) -> None:
+        delta_rms, relative_delta, cosine = self._state_change(before, after, valid_mask)
+        self.transition_names.append(name)
+        self.transition_delta_rms.append(delta_rms)
+        self.transition_relative_delta.append(relative_delta)
+        self.transition_cosine.append(cosine)
+
+        if visual_mask is not None and bool(visual_mask.any()):
+            visual_delta, _, visual_cosine = self._state_change(before, after, visual_mask)
+        else:
+            visual_delta = after.new_tensor(float("nan"), dtype=torch.float32)
+            visual_cosine = after.new_tensor(float("nan"), dtype=torch.float32)
+        self.transition_visual_delta_rms.append(visual_delta)
+        self.transition_visual_cosine.append(visual_cosine)
+
+        response_delta, _, response_cosine = self._state_change(
+            before, after, response_mask
+        )
+        self.transition_response_delta_rms.append(response_delta)
+        self.transition_response_cosine.append(response_cosine)
+
+        response_hidden = after.detach()[response_mask]
+        if response_hidden.shape[0] > logit_lens_token_limit:
+            positions = torch.linspace(
+                0,
+                response_hidden.shape[0] - 1,
+                steps=logit_lens_token_limit,
+                device=response_hidden.device,
+            ).round().long()
+            response_hidden = response_hidden.index_select(0, positions)
+        self.logit_lens_names.append(name)
+        self.logit_lens_response_hidden.append(response_hidden)
+
+    def record(
+        self,
+        gate: torch.Tensor,
+        alpha: torch.Tensor,
+        reinjection: torch.Tensor,
+        current_visual: torch.Tensor,
+        valid_mask: torch.BoolTensor | None = None,
+        saturation_lower: float = 0.05,
+        saturation_upper: float = 1.95,
+    ) -> None:
+        detached = gate.detach().float()
+        if valid_mask is None:
+            valid_mask = torch.ones_like(detached, dtype=torch.bool)
+        selected = detached[valid_mask]
+        if not selected.numel():
+            return
+        probability = (selected / 2.0).clamp(1e-6, 1.0 - 1e-6)
+        entropy = -(probability * probability.log() + (1 - probability) * (1 - probability).log())
+        reinjection_norm = reinjection.detach().float()[valid_mask].norm()
+        current_visual_norm = current_visual.detach().float()[valid_mask].norm()
+        ratio = reinjection_norm / current_visual_norm.clamp_min(1e-12)
+        self.gate_mean.append(selected.mean())
+        self.gate_std.append(selected.std())
+        self.gate_min.append(selected.min())
+        self.gate_max.append(selected.max())
+        self.gate_entropy.append(entropy.mean())
+        self.gate_saturation_rate.append(
+            ((selected < saturation_lower) | (selected > saturation_upper)).float().mean()
+        )
+        self.alpha.append(alpha.detach().float())
+        self.reinjection_norm.append(reinjection_norm)
+        self.current_visual_norm.append(current_visual_norm)
+        self.reinjection_ratio.append(ratio)
+        self.gate_maps.append(
+            [detached[row, valid_mask[row]].cpu() for row in range(detached.shape[0])]
+        )
+
+
+class MultimodalHrmTextModel(HrmTextModel):
+    """Minimal HRM recurrent patch for spatial RoPE and gated visual re-grounding."""
+
+    def __init__(
+        self,
+        config,
+        rope_mode: str = "hybrid_2d",
+        use_visual_anchor: bool = True,
+        use_visual_gate: bool = True,
+        gate_dim: int = 256,
+        anchor_raw_init: float = -2.0,
+        anchor_max_scale: float = 0.5,
+        gate_saturation_metric_lower: float = 0.05,
+        gate_saturation_metric_upper: float = 1.95,
+    ):
+        super().__init__(config)
+        if rope_mode not in {"original_1d", "hybrid_2d"}:
+            raise ValueError(f"Unsupported rope_mode: {rope_mode}")
+        self.rope_mode = rope_mode
+        self.use_visual_anchor = use_visual_anchor
+        self.use_visual_gate = use_visual_gate
+        self.gate_saturation_metric_lower = float(gate_saturation_metric_lower)
+        self.gate_saturation_metric_upper = float(gate_saturation_metric_upper)
+        rope_parameters = getattr(config, "rope_parameters", None) or {}
+        rope_theta = getattr(config, "rope_theta", None) or rope_parameters.get("rope_theta", 10000.0)
+        self.hybrid_rotary_emb = HybridHrmRotaryEmbedding(
+            head_dim=config.head_dim,
+            theta=rope_theta,
+            config=config,
+        )
+        self.visual_gate = LoopAwareVisualGate(config.hidden_size, gate_dim)
+        self.visual_anchor = VisualAnchor(
+            config.H_cycles,
+            raw_init=anchor_raw_init,
+            max_scale=anchor_max_scale,
+        )
+        self.last_visual_diagnostics: RecurrentVisualDiagnostics | None = None
+
+    @classmethod
+    def from_text_model(cls, text_model: HrmTextModel, **kwargs) -> "MultimodalHrmTextModel":
+        model = cls(text_model.config, **kwargs)
+        incompatible = model.load_state_dict(text_model.state_dict(), strict=False)
+        allowed_missing = {
+            "visual_gate.visual_proj.weight",
+            "visual_gate.query_proj.weight",
+            "visual_gate.out_proj.weight",
+            "visual_gate.out_proj.bias",
+            "visual_anchor.raw_alpha",
+        }
+        unexpected_missing = set(incompatible.missing_keys).difference(allowed_missing)
+        if unexpected_missing or incompatible.unexpected_keys:
+            raise RuntimeError(
+                f"HRM conversion mismatch: missing={sorted(unexpected_missing)}, "
+                f"unexpected={incompatible.unexpected_keys}"
+            )
+        return model
+
+    @staticmethod
+    def _extract_tokens(
+        states: torch.Tensor, mask: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.BoolTensor]:
+        counts = mask.sum(dim=1)
+        max_count = int(counts.max()) if counts.numel() else 0
+        tokens = states.new_zeros((states.shape[0], max_count, states.shape[-1]))
+        valid = torch.zeros((states.shape[0], max_count), dtype=torch.bool, device=states.device)
+        for row, count in enumerate(counts.tolist()):
+            if count:
+                tokens[row, :count] = states[row, mask[row]]
+                valid[row, :count] = True
+        return tokens, valid
+
+    @staticmethod
+    def _scatter_tokens(
+        tokens: torch.Tensor,
+        valid: torch.BoolTensor,
+        mask: torch.Tensor,
+        reference: torch.Tensor,
+    ) -> torch.Tensor:
+        output = torch.zeros_like(reference)
+        for row in range(reference.shape[0]):
+            output[row, mask[row]] = tokens[row, valid[row]]
+        return output
+
+    def _run_stack(self, stack, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
+        return stack(hidden_states, **kwargs)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        token_type_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        rope_position_ids: torch.LongTensor | None = None,
+        visual_mask: torch.BoolTensor | None = None,
+        instruction_mask: torch.BoolTensor | None = None,
+        collect_visual_diagnostics: bool = False,
+        collect_depth_diagnostics: bool = False,
+        **kwargs,
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        inputs_embeds = inputs_embeds * self.embedding_scale
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+        if position_ids is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            position_ids = position_ids.unsqueeze(0)
+
+        mask_kwargs = {
+            "config": self.config,
+            "inputs_embeds": inputs_embeds,
+            "attention_mask": attention_mask,
+            "past_key_values": past_key_values,
+            "position_ids": position_ids,
+        }
+        is_first_iteration = past_key_values is None or not past_key_values.is_initialized
+        if token_type_ids is not None and is_first_iteration and self.config.prefix_lm:
+            mask_kwargs["block_sequence_ids"] = torch.where(token_type_ids == 1, 0, -1)
+        causal_attention_mask = create_causal_mask(**mask_kwargs)
+
+        if self.rope_mode == "hybrid_2d" and rope_position_ids is not None:
+            position_embeddings = self.hybrid_rotary_emb(inputs_embeds, rope_position_ids)
+        else:
+            position_embeddings = self.rotary_emb(inputs_embeds, position_ids)
+
+        hidden_states_high_cycle = inputs_embeds
+        hidden_states_low_cycle = (
+            self.z_L_init.to(dtype=hidden_states_high_cycle.dtype, device=hidden_states_high_cycle.device)
+            .expand_as(hidden_states_high_cycle)
+            .contiguous()
+        )
+        multimodal = visual_mask is not None and bool(visual_mask.any())
+        if multimodal:
+            if instruction_mask is None:
+                raise ValueError("instruction_mask is required when visual_mask is provided")
+            multimodal_rows = visual_mask.any(dim=1)
+            multimodal_indices = multimodal_rows.nonzero(as_tuple=False).flatten()
+            visual_anchor_tokens, visual_token_valid = self._extract_tokens(
+                inputs_embeds[multimodal_rows], visual_mask[multimodal_rows]
+            )
+            visual_anchor_tokens = visual_anchor_tokens.clone()
+        else:
+            multimodal_rows = None
+            multimodal_indices = None
+            visual_anchor_tokens = None
+            visual_token_valid = None
+
+        diagnostics = (
+            RecurrentVisualDiagnostics.empty()
+            if collect_visual_diagnostics or collect_depth_diagnostics
+            else None
+        )
+        valid_token_mask = (
+            attention_mask.bool()
+            if attention_mask is not None
+            else torch.ones(inputs_embeds.shape[:2], dtype=torch.bool, device=inputs_embeds.device)
+        )
+        response_mask = valid_token_mask & (
+            token_type_ids == 0 if token_type_ids is not None else valid_token_mask
+        )
+        num_layers_per_stack = self.config.num_layers_per_stack
+        for high_cycle_idx in range(self.config.H_cycles):
+            if multimodal and self.use_visual_anchor:
+                assert multimodal_rows is not None and multimodal_indices is not None
+                current_visual, current_visual_valid = self._extract_tokens(
+                    hidden_states_high_cycle[multimodal_rows], visual_mask[multimodal_rows]
+                )
+                assert visual_token_valid is not None
+                if not torch.equal(current_visual_valid, visual_token_valid):
+                    raise RuntimeError("Dynamic visual-token layout changed inside HRM recurrence")
+                instruction_state = masked_mean(
+                    hidden_states_high_cycle[multimodal_rows],
+                    instruction_mask[multimodal_rows],
+                )
+                if self.use_visual_gate:
+                    gates = self.visual_gate(current_visual, instruction_state)
+                else:
+                    gates = torch.ones(current_visual.shape[:2], dtype=current_visual.dtype, device=current_visual.device)
+                gates = gates.masked_fill(~visual_token_valid, 0.0)
+                reinjection_tokens = self.visual_anchor(visual_anchor_tokens, gates, high_cycle_idx)
+                reinjection_rows = self._scatter_tokens(
+                    reinjection_tokens,
+                    visual_token_valid,
+                    visual_mask[multimodal_rows],
+                    hidden_states_high_cycle[multimodal_rows],
+                )
+                reinjection_full = torch.zeros_like(hidden_states_high_cycle).index_copy(
+                    0, multimodal_indices, reinjection_rows
+                )
+                hidden_states_high_cycle = hidden_states_high_cycle + reinjection_full
+                if diagnostics is not None:
+                    diagnostics.record(
+                        gates,
+                        self.visual_anchor.alpha(high_cycle_idx),
+                        reinjection_tokens,
+                        current_visual,
+                        visual_token_valid,
+                        self.gate_saturation_metric_lower,
+                        self.gate_saturation_metric_upper,
+                    )
+
+            for low_cycle_idx in range(self.config.L_cycles):
+                cycle_offset = (
+                    high_cycle_idx * (self.config.L_cycles + 1) + low_cycle_idx
+                ) * num_layers_per_stack
+                previous_low = hidden_states_low_cycle
+                hidden_states_low_cycle = self._run_stack(
+                    self.L_module,
+                    hidden_states_low_cycle.to(hidden_states_high_cycle.device) + hidden_states_high_cycle,
+                    attention_mask=causal_attention_mask,
+                    past_key_values=past_key_values,
+                    position_embeddings=position_embeddings,
+                    position_ids=position_ids,
+                    cycle_offset=cycle_offset,
+                    **kwargs,
+                )
+                if diagnostics is not None and collect_depth_diagnostics:
+                    diagnostics.record_transition(
+                        f"h{high_cycle_idx}_l{low_cycle_idx}",
+                        previous_low,
+                        hidden_states_low_cycle,
+                        valid_token_mask,
+                        visual_mask,
+                        response_mask,
+                    )
+
+            cycle_offset = (
+                high_cycle_idx * (self.config.L_cycles + 1) + self.config.L_cycles
+            ) * num_layers_per_stack
+            previous_high = hidden_states_high_cycle
+            hidden_states_high_cycle = self._run_stack(
+                self.H_module,
+                hidden_states_high_cycle + hidden_states_low_cycle.to(hidden_states_high_cycle.device),
+                attention_mask=causal_attention_mask,
+                past_key_values=past_key_values,
+                position_embeddings=position_embeddings,
+                position_ids=position_ids,
+                cycle_offset=cycle_offset,
+                **kwargs,
+            )
+            if diagnostics is not None and collect_depth_diagnostics:
+                diagnostics.record_transition(
+                    f"h{high_cycle_idx}_h",
+                    previous_high,
+                    hidden_states_high_cycle,
+                    valid_token_mask,
+                    visual_mask,
+                    response_mask,
+                )
+
+        self.last_visual_diagnostics = diagnostics
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states_high_cycle,
+            past_key_values=past_key_values,
+        )

hrm_penguin/recurrent/hybrid_rope.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    first, second = x.chunk(2, dim=-1)
+    return torch.cat((-second, first), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    unsqueeze_dim: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    return (
+        query * cos + rotate_half(query) * sin,
+        key * cos + rotate_half(key) * sin,
+    )
+
+
+class HybridHrmRotaryEmbedding(nn.Module):
+    """2D visual RoPE that exactly reduces to HRM 1D RoPE for positions (p, p)."""
+
+    def __init__(
+        self,
+        head_dim: int = 128,
+        theta: float = 10000.0,
+        config=None,
+    ):
+        super().__init__()
+        if head_dim % 2:
+            raise ValueError("head_dim must be even")
+        attention_scaling = 1.0
+        rope_parameters = getattr(config, "rope_parameters", None) or {}
+        rope_type = rope_parameters.get("rope_type", "default")
+        if config is not None and rope_type != "default":
+            from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+
+            if rope_type not in ROPE_INIT_FUNCTIONS:
+                raise ValueError(f"Unsupported hybrid RoPE type: {rope_type}")
+            inv_freq, attention_scaling = ROPE_INIT_FUNCTIONS[rope_type](
+                config,
+                device=None,
+            )
+            inv_freq = inv_freq.float()
+        else:
+            inv_freq = 1.0 / (
+                theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
+            )
+        self.head_dim = head_dim
+        self.rope_type = rope_type
+        self.attention_scaling = float(attention_scaling)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope_position_ids: torch.LongTensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if rope_position_ids.ndim == 2:
+            rope_position_ids = rope_position_ids.unsqueeze(0).expand(2, -1, -1)
+        if rope_position_ids.ndim != 3 or rope_position_ids.shape[0] != 2:
+            raise ValueError("rope_position_ids must have shape [2, B, S]")
+        positions = rope_position_ids.float()
+        freqs = positions.unsqueeze(-1) * self.inv_freq.float().view(1, 1, 1, -1)
+        # Standard HRM cos is [freqs, freqs]. Selecting axis0's first half
+        # and axis1's second half preserves exact 1D behavior for (p, p).
+        angles = torch.cat((freqs[0], freqs[1]), dim=-1)
+        if angles.shape[-1] != self.head_dim:
+            raise RuntimeError(f"RoPE produced {angles.shape[-1]} dims, expected {self.head_dim}")
+        scale = self.attention_scaling
+        return (angles.cos() * scale).to(x.dtype), (angles.sin() * scale).to(x.dtype)

hrm_penguin/recurrent/positions.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import torch
+
+
+def build_compact_positions(
+    header_lengths: torch.LongTensor,
+    instruction_lengths: torch.LongTensor,
+    response_lengths: torch.LongTensor,
+    visual_token_counts: torch.LongTensor | None = None,
+    grid_heights: torch.LongTensor | None = None,
+    grid_widths: torch.LongTensor | None = None,
+    grid_h: int = 24,
+    grid_w: int = 24,
+    rope_mode: str = "hybrid_2d",
+) -> tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]:
+    """Build true sequence positions and independent compact spatial RoPE positions."""
+    batch = int(header_lengths.numel())
+    visual_tokens = grid_h * grid_w
+    if visual_token_counts is None:
+        visual_token_counts = torch.full_like(header_lengths, visual_tokens)
+    if grid_heights is None:
+        grid_heights = torch.where(
+            visual_token_counts > 0,
+            torch.full_like(visual_token_counts, grid_h),
+            torch.zeros_like(visual_token_counts),
+        )
+    if grid_widths is None:
+        grid_widths = torch.where(
+            visual_token_counts > 0,
+            torch.full_like(visual_token_counts, grid_w),
+            torch.zeros_like(visual_token_counts),
+        )
+    if not (
+        grid_heights.shape == visual_token_counts.shape
+        and grid_widths.shape == visual_token_counts.shape
+    ):
+        raise ValueError("Per-sample grid tensors must match visual_token_counts")
+    expected_counts = grid_heights * grid_widths
+    if bool((expected_counts != visual_token_counts).any()):
+        raise ValueError(
+            "visual token counts must exactly equal per-sample grid area: "
+            f"counts={visual_token_counts.tolist()}, "
+            f"heights={grid_heights.tolist()}, widths={grid_widths.tolist()}"
+        )
+    total_lengths = header_lengths + visual_token_counts + instruction_lengths + response_lengths
+    max_length = int(total_lengths.max().item())
+    device = header_lengths.device
+    mask_positions = torch.arange(max_length, device=device).expand(batch, -1).clone()
+    rope_positions = mask_positions.unsqueeze(0).expand(2, -1, -1).clone()
+    visual_mask = torch.zeros((batch, max_length), dtype=torch.bool, device=device)
+    instruction_mask = torch.zeros_like(visual_mask)
+
+    for index in range(batch):
+        header = int(header_lengths[index])
+        instruction = int(instruction_lengths[index])
+        response = int(response_lengths[index])
+        row_visual_tokens = int(visual_token_counts[index])
+        row_grid_h = int(grid_heights[index])
+        row_grid_w = int(grid_widths[index])
+        visual_start = header
+        visual_end = visual_start + row_visual_tokens
+        instruction_start = visual_end
+        instruction_end = instruction_start + instruction
+        response_end = instruction_end + response
+        visual_mask[index, visual_start:visual_end] = True
+        instruction_mask[index, instruction_start:instruction_end] = True
+        if rope_mode == "hybrid_2d" and row_visual_tokens:
+            rows = torch.arange(row_grid_h, device=device).repeat_interleave(row_grid_w) + header
+            cols = torch.arange(row_grid_w, device=device).repeat(row_grid_h) + header
+            rope_positions[0, index, visual_start:visual_end] = rows
+            rope_positions[1, index, visual_start:visual_end] = cols
+            text_start = header + max(row_grid_h, row_grid_w)
+            tail = torch.arange(response_end - instruction_start, device=device) + text_start
+            rope_positions[:, index, instruction_start:response_end] = tail.unsqueeze(0)
+    return mask_positions, rope_positions, visual_mask, instruction_mask

hrm_penguin/recurrent/visual_anchor.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class VisualAnchor(nn.Module):
+    def __init__(
+        self,
+        H_cycles: int = 2,
+        raw_init: float = -2.0,
+        max_scale: float = 0.5,
+    ):
+        super().__init__()
+        self.raw_alpha = nn.Parameter(torch.full((H_cycles,), float(raw_init)))
+        self.max_scale = float(max_scale)
+
+    def alpha(self, cycle: int) -> torch.Tensor:
+        return self.max_scale * torch.sigmoid(self.raw_alpha[cycle])
+
+    def forward(
+        self,
+        visual_anchor: torch.Tensor,
+        gates: torch.Tensor,
+        cycle: int,
+    ) -> torch.Tensor:
+        alpha = self.alpha(cycle).to(dtype=visual_anchor.dtype)
+        return alpha * gates.unsqueeze(-1) * visual_anchor

hrm_penguin/recurrent/visual_gate.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class LoopAwareVisualGate(nn.Module):
+    def __init__(self, hidden_size: int = 1536, gate_dim: int = 256):
+        super().__init__()
+        self.visual_proj = nn.Linear(hidden_size, gate_dim, bias=False)
+        self.query_proj = nn.Linear(hidden_size, gate_dim, bias=False)
+        self.out_proj = nn.Linear(gate_dim, 1, bias=True)
+        self.reset_output_head()
+
+    def reset_output_head(self) -> None:
+        """Make the gate neutral while retaining its learned feature projections."""
+        nn.init.zeros_(self.out_proj.weight)
+        nn.init.zeros_(self.out_proj.bias)
+
+    def reset_parameters(self) -> None:
+        """Initialize a neutral gate before checkpoint loading."""
+        self.visual_proj.reset_parameters()
+        self.query_proj.reset_parameters()
+        self.reset_output_head()
+
+    def forward(
+        self,
+        visual_states: torch.Tensor,
+        instruction_state: torch.Tensor,
+    ) -> torch.Tensor:
+        output_dtype = visual_states.dtype
+        parameter_dtype = self.visual_proj.weight.dtype
+        # Keep these linears and sigmoid in fp32 when the surrounding model
+        # uses bf16 autocast, preserving the checkpoint's gate precision.
+        with torch.autocast(device_type=visual_states.device.type, enabled=False):
+            visual = self.visual_proj(visual_states.to(parameter_dtype))
+            query = self.query_proj(instruction_state.to(parameter_dtype)).unsqueeze(1)
+            score = self.out_proj(F.silu(visual + query)).squeeze(-1)
+            # Preserve the evaluated checkpoint's floating-point operation
+            # order, including the bounded auxiliary gate calculation.
+            bounded_score = score.clamp(-8.0, 8.0)
+            surrogate_score = score + (bounded_score - score).detach()
+            raw_gate = 2.0 * torch.sigmoid(score)
+            surrogate_gate = 2.0 * torch.sigmoid(surrogate_score)
+            gate = raw_gate.detach() + surrogate_gate - surrogate_gate.detach()
+        return gate.to(output_dtype)
+
+
+def masked_mean(states: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    weights = mask.to(states.dtype).unsqueeze(-1)
+    denominator = weights.sum(dim=1).clamp_min(1.0)
+    return (states * weights).sum(dim=1) / denominator

hrm_penguin/vision/__init__.py

@@ -0,0 +1,17 @@
+from .encoder import PenguinVisionEncoder
+from .preprocessing import (
+    PenguinDynamicResolutionProcessor,
+    PenguinLetterboxProcessor,
+    build_penguin_image_processor,
+    letterbox_image,
+)
+from .projector import VisionProjector
+
+__all__ = [
+    "PenguinVisionEncoder",
+    "PenguinLetterboxProcessor",
+    "PenguinDynamicResolutionProcessor",
+    "build_penguin_image_processor",
+    "VisionProjector",
+    "letterbox_image",
+]

hrm_penguin/vision/encoder.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+from types import MethodType
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+def _rotate_half(value: torch.Tensor) -> torch.Tensor:
+    first, second = value.chunk(2, dim=-1)
+    return torch.cat((-second, first), dim=-1)
+
+
+def _apply_penguin_rotary(query, key, cos, sin):
+    section = [cos.shape[-1] // 2, cos.shape[-1] // 2]
+    cos = torch.cat([axis[index % 2] for index, axis in enumerate(cos.split(section, dim=-1))], dim=-1)
+    sin = torch.cat([axis[index % 2] for index, axis in enumerate(sin.split(section, dim=-1))], dim=-1)
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    return query * cos + _rotate_half(query) * sin, key * cos + _rotate_half(key) * sin
+
+
+def _penguin_sdpa_attention_forward(
+    module,
+    hidden_states: torch.Tensor,
+    position_embeddings,
+    attention_mask=None,
+    past_key_value=None,
+    cache_position=None,
+    cu_seqlens=None,
+    **kwargs,
+):
+    """SDPA equivalent of Penguin's packed, non-causal FlashAttention path."""
+    input_shape = hidden_states.shape[:-1]
+    hidden_shape = (*input_shape, -1, module.head_dim)
+    query = module.q_norm(module.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+    key = module.k_norm(module.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+    value = module.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+    cos, sin = position_embeddings
+    query, key = _apply_penguin_rotary(query, key, cos, sin)
+    if past_key_value is not None:
+        cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+        key, value = past_key_value.update(key, value, module.layer_idx, cache_kwargs)
+
+    if cu_seqlens is None:
+        attended = F.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            dropout_p=0.0,
+            is_causal=False,
+            enable_gqa=query.shape[1] != key.shape[1],
+        )
+    else:
+        # Penguin packs image sequences along S inside a singleton batch. Splitting
+        # here prevents patches from different images attending to one another.
+        boundaries = cu_seqlens.detach().cpu().tolist()
+        chunks = []
+        for start, end in zip(boundaries[:-1], boundaries[1:]):
+            chunks.append(
+                F.scaled_dot_product_attention(
+                    query[:, :, start:end],
+                    key[:, :, start:end],
+                    value[:, :, start:end],
+                    dropout_p=0.0,
+                    is_causal=False,
+                    enable_gqa=query.shape[1] != key.shape[1],
+                )
+            )
+        attended = torch.cat(chunks, dim=2)
+    attended = attended.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+    return module.o_proj(attended), None
+
+
+def install_penguin_sdpa(model: nn.Module) -> int:
+    patched = 0
+    for module in model.modules():
+        if module.__class__.__name__ == "PenguinVLAttention":
+            module.forward = MethodType(_penguin_sdpa_attention_forward, module)
+            patched += 1
+    if patched == 0:
+        raise RuntimeError("No PenguinVLAttention modules found for SDPA patch")
+    return patched
+
+
+def install_penguin_transformers_compatibility() -> None:
+    """Bridge Penguin remote code written for Transformers 4.51 to 5.x RoPE APIs."""
+    from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+
+    if "default" not in ROPE_INIT_FUNCTIONS:
+        ROPE_INIT_FUNCTIONS["default"] = compute_default_rope_parameters
+
+
+def compute_default_rope_parameters(config, device=None, **kwargs):
+    parameters = getattr(config, "rope_parameters", None) or getattr(config, "rope_scaling", None) or {}
+    theta = float(parameters.get("rope_theta", getattr(config, "rope_theta", 10000.0)))
+    dimension = int(getattr(config, "head_dim", config.hidden_size // config.num_attention_heads))
+    inverse_frequency = 1.0 / (
+        theta ** (torch.arange(0, dimension, 2, dtype=torch.float32, device=device) / dimension)
+    )
+    return inverse_frequency, 1.0
+
+
+class PenguinVisionEncoder(nn.Module):
+    """Thin checked wrapper around the official Penguin encoder."""
+
+    def __init__(self, model: nn.Module, hidden_size: int = 1024):
+        super().__init__()
+        self.model = model
+        self.hidden_size = hidden_size
+
+    def forward(self, pixel_values, grid_sizes, merge_sizes) -> torch.Tensor:
+        features = self.model(
+            pixel_values=pixel_values,
+            grid_sizes=grid_sizes,
+            merge_sizes=merge_sizes,
+        )
+        if hasattr(features, "last_hidden_state"):
+            features = features.last_hidden_state
+        if features.ndim != 2 or features.shape[-1] != self.hidden_size:
+            raise RuntimeError(
+                f"Expected packed Penguin output [sum(N_i), {self.hidden_size}], "
+                f"got {tuple(features.shape)}"
+            )
+        if grid_sizes.ndim != 2 or grid_sizes.shape[1] != 3:
+            raise RuntimeError(f"Invalid grid_sizes shape: {tuple(grid_sizes.shape)}")
+        if merge_sizes.ndim != 1 or merge_sizes.shape[0] != grid_sizes.shape[0]:
+            raise RuntimeError(f"Invalid merge_sizes shape: {tuple(merge_sizes.shape)}")
+        spatial = grid_sizes[:, 1:] // merge_sizes[:, None]
+        if bool((spatial * merge_sizes[:, None] != grid_sizes[:, 1:]).any()):
+            raise RuntimeError("Grid dimensions must be divisible by merge_size")
+        counts = grid_sizes[:, 0] * spatial[:, 0] * spatial[:, 1]
+        if features.shape[0] != int(counts.sum()):
+            raise RuntimeError(
+                "Packed Penguin features do not match dynamic grids: "
+                f"features={features.shape[0]}, counts={counts.tolist()}"
+            )
+        return features

hrm_penguin/vision/preprocessing.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+
+from PIL import Image, ImageOps
+
+
+def letterbox_image(
+    image: Image.Image,
+    size: int = 336,
+    fill: tuple[int, int, int] = (128, 128, 128),
+) -> Image.Image:
+    """EXIF-correct, convert to RGB, preserve aspect ratio, and center pad."""
+    image = ImageOps.exif_transpose(image).convert("RGB")
+    width, height = image.size
+    if width < 1 or height < 1:
+        raise ValueError(f"Invalid image size: {image.size}")
+    scale = min(size / width, size / height)
+    new_width = max(1, round(width * scale))
+    new_height = max(1, round(height * scale))
+    image = image.resize((new_width, new_height), Image.Resampling.BICUBIC)
+    canvas = Image.new("RGB", (size, size), fill)
+    canvas.paste(image, ((size - new_width) // 2, (size - new_height) // 2))
+    return canvas
+
+
+class PenguinLetterboxProcessor:
+    """Letterbox first, then delegate rescale/normalize/patchify to Penguin."""
+
+    def __init__(self, processor, image_size: int = 336):
+        self.processor = processor
+        self.image_size = image_size
+
+    @property
+    def neutral_fill(self) -> tuple[int, int, int]:
+        mean = getattr(self.processor, "image_mean", (0.5, 0.5, 0.5))
+        if isinstance(mean, (float, int)):
+            mean = (mean, mean, mean)
+        return tuple(max(0, min(255, round(float(value) * 255))) for value in mean)
+
+    def __call__(self, images: Image.Image | Sequence[Image.Image], return_tensors: str = "pt"):
+        if isinstance(images, Image.Image):
+            images = [images]
+        boxed = [letterbox_image(image, self.image_size, self.neutral_fill) for image in images]
+        result = self.processor(
+            images=boxed,
+            do_resize=False,
+            merge_size=1,
+            return_tensors=return_tensors,
+        )
+        required = {"pixel_values", "grid_sizes", "merge_sizes"}
+        missing = required.difference(result.keys())
+        if missing:
+            raise RuntimeError(f"Penguin processor missing fields: {sorted(missing)}")
+        return result
+
+
+class PenguinDynamicResolutionProcessor:
+    """Use Penguin's native aspect-ratio preserving resize without image padding.
+
+    Penguin already packs differently sized images into one flat patch tensor.
+    This wrapper fixes the supported token range, applies EXIF orientation, and
+    validates that every emitted token corresponds to a real resized image patch.
+    """
+
+    def __init__(
+        self,
+        processor,
+        min_visual_tokens: int = 256,
+        max_visual_tokens: int = 1024,
+    ):
+        if min_visual_tokens < 1 or max_visual_tokens < min_visual_tokens:
+            raise ValueError("Invalid dynamic visual-token range")
+        self.processor = processor
+        self.min_visual_tokens = int(min_visual_tokens)
+        self.max_visual_tokens = int(max_visual_tokens)
+        # The remote Penguin processor reads these attributes while selecting
+        # aligned target sizes. merge_size=1 keeps one encoder token per patch.
+        self.processor.min_tokens = self.min_visual_tokens
+        self.processor.max_tokens = self.max_visual_tokens
+
+    def _fit_single(
+        self,
+        image: Image.Image,
+        return_tensors: str,
+        min_visual_tokens: int | None = None,
+        max_visual_tokens: int | None = None,
+    ):
+        """Fit one real image to the configured patch-token contract.
+
+        Penguin's helper internally multiplies ``min_tokens`` by 1.5 and
+        ``max_tokens`` by 0.95 before choosing a grid.  Compensate for those
+        implementation factors, then tighten and retry if aspect-ratio rounding
+        crosses the public limits.  No canvas or synthetic patch is introduced.
+        """
+
+        minimum = int(min_visual_tokens or self.min_visual_tokens)
+        maximum = int(max_visual_tokens or self.max_visual_tokens)
+        if minimum < 1 or maximum < minimum or maximum > self.max_visual_tokens:
+            raise ValueError(f"Invalid per-image token range [{minimum}, {maximum}]")
+        original_min = self.processor.min_tokens
+        original_max = self.processor.max_tokens
+        # Floor is intentional: using ceil(256 / 1.5) makes Penguin's own
+        # ``ceil`` alignment jump a square 16x16 grid to 17x17.
+        delegate_min = max(1, math.floor(minimum / 1.5))
+        delegate_max = max(delegate_min, math.ceil(maximum / 0.95))
+        last_count = None
+        try:
+            for _ in range(12):
+                self.processor.min_tokens = delegate_min
+                self.processor.max_tokens = delegate_max
+                result = self.processor(
+                    images=[image],
+                    do_resize=True,
+                    merge_size=1,
+                    return_tensors=return_tensors,
+                )
+                grids = result["grid_sizes"]
+                count = int(grids.prod(dim=1)[0])
+                last_count = count
+                if minimum <= count <= maximum:
+                    return result
+                if count > maximum:
+                    delegate_max = max(
+                        delegate_min,
+                        math.floor(
+                            delegate_max
+                            * maximum
+                            / max(count, 1)
+                            * 0.98
+                        ),
+                    )
+                else:
+                    delegate_min = min(
+                        delegate_max,
+                        math.ceil(
+                            delegate_min
+                            * minimum
+                            / max(count, 1)
+                            * 1.02
+                        ),
+                    )
+            raise RuntimeError(
+                "Penguin token-budget tightening did not converge: "
+                f"last_count={last_count}, range="
+                f"[{minimum}, {maximum}]"
+            )
+        finally:
+            self.processor.min_tokens = original_min
+            self.processor.max_tokens = original_max
+
+    def __call__(
+        self,
+        images: Image.Image | Sequence[Image.Image],
+        return_tensors: str = "pt",
+        max_visual_tokens: int | Sequence[int] | None = None,
+    ):
+        if isinstance(images, Image.Image):
+            images = [images]
+        images = [ImageOps.exif_transpose(image).convert("RGB") for image in images]
+        if not images:
+            raise ValueError("At least one image is required")
+        if any(image.width < 1 or image.height < 1 for image in images):
+            raise ValueError("Images must have positive dimensions")
+        # Penguin's video-oriented processor shares max_tokens across every
+        # item passed in one call. LoopVL defines the budget per image, so
+        # select each grid independently and concatenate the already-packed
+        # patch tensors afterwards.
+        if max_visual_tokens is None:
+            caps = [self.max_visual_tokens] * len(images)
+        elif isinstance(max_visual_tokens, int):
+            caps = [int(max_visual_tokens)] * len(images)
+        else:
+            caps = [int(value) for value in max_visual_tokens]
+            if len(caps) != len(images):
+                raise ValueError("One visual-token cap is required per image")
+        per_image = [
+            self._fit_single(
+                image,
+                return_tensors,
+                min_visual_tokens=self.min_visual_tokens,
+                max_visual_tokens=cap,
+            )
+            for image, cap in zip(images, caps, strict=True)
+        ]
+        first = per_image[0]
+        pixel_values = first["pixel_values"]
+        grid_sizes = first["grid_sizes"]
+        merge_sizes = first["merge_sizes"]
+        if len(per_image) > 1:
+            import torch
+
+            pixel_values = torch.cat([item["pixel_values"] for item in per_image], dim=0)
+            grid_sizes = torch.cat([item["grid_sizes"] for item in per_image], dim=0)
+            merge_sizes = torch.cat([item["merge_sizes"] for item in per_image], dim=0)
+        result = {
+            "pixel_values": pixel_values,
+            "grid_sizes": grid_sizes,
+            "merge_sizes": merge_sizes,
+        }
+        required = {"pixel_values", "grid_sizes", "merge_sizes"}
+        missing = required.difference(result.keys())
+        if missing:
+            raise RuntimeError(f"Penguin processor missing fields: {sorted(missing)}")
+
+        grids = result["grid_sizes"]
+        merges = result["merge_sizes"]
+        if grids.ndim != 2 or grids.shape[1] != 3 or grids.shape[0] != len(images):
+            raise RuntimeError(f"Unexpected Penguin grid_sizes shape: {tuple(grids.shape)}")
+        if merges.ndim != 1 or merges.shape[0] != len(images) or bool((merges != 1).any()):
+            raise RuntimeError("Dynamic LoopVL input requires merge_size=1 for every image")
+        if bool((grids[:, 0] != 1).any()):
+            raise RuntimeError("LoopVL supports one image frame per sample")
+
+        counts = grids.prod(dim=1)
+        import torch
+
+        cap_tensor = torch.tensor(caps, device=counts.device, dtype=counts.dtype)
+        if bool(((counts < self.min_visual_tokens) | (counts > cap_tensor)).any()):
+            raise RuntimeError(
+                "Penguin emitted visual token counts outside per-image budgets: "
+                f"counts={counts.tolist()}, caps={caps}"
+            )
+        pixel_values = result["pixel_values"]
+        if pixel_values.ndim != 2 or pixel_values.shape[0] != int(counts.sum()):
+            raise RuntimeError(
+                "Packed pixel patches do not exactly match the real dynamic grids: "
+                f"patches={tuple(pixel_values.shape)}, grids={grids.tolist()}"
+            )
+        return result
+
+
+def build_penguin_image_processor(processor, vision_config):
+    if getattr(vision_config, "dynamic_resolution", False):
+        return PenguinDynamicResolutionProcessor(
+            processor,
+            min_visual_tokens=vision_config.min_visual_tokens,
+            max_visual_tokens=vision_config.max_visual_tokens,
+        )
+    return PenguinLetterboxProcessor(processor, vision_config.input_size)

hrm_penguin/vision/projector.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class VisionProjector(nn.Module):
+    def __init__(
+        self,
+        input_size: int = 1024,
+        hidden_size: int = 1536,
+        target_raw_rms: float = 1.0,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.fc1 = nn.Linear(input_size, hidden_size, bias=True)
+        self.act = nn.GELU()
+        self.fc2 = nn.Linear(hidden_size, hidden_size, bias=True)
+        self.eps = eps
+        self.register_buffer(
+            "visual_embedding_scale",
+            torch.tensor(float(target_raw_rms), dtype=torch.float32),
+        )
+        self.last_rms_diagnostics: dict[str, torch.Tensor] | None = None
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        features = features.to(dtype=self.fc1.weight.dtype)
+        projected = self.fc2(self.act(self.fc1(features)))
+        scaled = F.rms_norm(
+            projected,
+            normalized_shape=(projected.shape[-1],),
+            weight=None,
+            eps=self.eps,
+        )
+        scaled = scaled * self.visual_embedding_scale.to(projected.dtype)
+        self.last_rms_diagnostics = {
+            "vision_encoder_rms": features.detach().float().square().mean().sqrt(),
+            "projector_rms": projected.detach().float().square().mean().sqrt(),
+            "visual_scaled_rms": scaled.detach().float().square().mean().sqrt(),
+        }
+        return scaled

metadata.py

@@ -0,0 +1,82 @@
+"""Standard-library-only metadata reader for LoopVL weight snapshots.
+
+No code is imported from the snapshot. Version 2 stores all data files at its
+root; version 1 retains its original nested image-processor JSON for backwards
+compatibility. Both versions use the same code distributed in this runtime.
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+
+COMPONENT_DTYPES = {
+    "vision": "bfloat16",
+    "language": "bfloat16",
+    "projector": "float32",
+    "visual_gate": "float32",
+    "visual_anchor": "float32",
+}
+
+FLAT_DATA_FILES = (
+    "model.safetensors",
+    "config.json",
+    "generation_config.json",
+    "preprocessor_config.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+)
+
+PROCESSOR_FIELDS = (
+    "do_convert_rgb", "do_normalize", "do_rescale", "do_resize",
+    "image_mean", "image_std", "max_tokens", "min_tokens", "patch_size",
+    "resample", "rescale_factor",
+)
+
+
+def read_checkpoint_metadata(model_dir: str | Path) -> dict:
+    """Read and validate format metadata without torch, transformers or imports."""
+    root = Path(model_dir).expanduser().resolve()
+    raw = json.loads((root / "config.json").read_text(encoding="utf-8"))
+    if raw.get("model_type") != "loopvl" or raw.get("format_version") not in (1, 2):
+        raise ValueError("Expected LoopVL single-file checkpoint format_version 1 or 2")
+    if raw.get("component_dtypes") != COMPONENT_DTYPES:
+        raise ValueError("This loader requires the original BF16 cores / FP32 adapters")
+    for field in ("loopvl_config", "text_config", "vision_config"):
+        if not isinstance(raw.get(field), dict):
+            raise ValueError(f"Checkpoint lacks embedded {field}")
+    if raw["text_config"].get("model_type") != "hrm_text":
+        raise ValueError("Expected the embedded hrm_text language configuration")
+    if raw["vision_config"].get("model_type") != "penguinvl_vision_encoder":
+        raise ValueError("Expected the embedded penguinvl_vision_encoder configuration")
+    return raw
+
+
+def read_processor_config(model_dir: str | Path, raw: dict | None = None) -> dict:
+    """Return the original processor parameters, never a remote-code reference.
+
+    The old root preprocessor_config.json was only a bundle description and is
+    deliberately not used for format 1. Omitting the true image mean/std would
+    silently select CLIP defaults and change pixel values.
+    """
+    root = Path(model_dir).expanduser().resolve()
+    if raw is None:
+        raw = read_checkpoint_metadata(root)
+    version = raw.get("format_version")
+    if version == 1:
+        path = root / "penguin_encoder" / "preprocessor_config.json"
+    elif version == 2:
+        path = root / "preprocessor_config.json"
+    else:
+        raise ValueError("Unsupported processor metadata format")
+    values = json.loads(path.read_text(encoding="utf-8"))
+    missing = sorted(set(PROCESSOR_FIELDS) - set(values))
+    if missing:
+        raise ValueError(f"Missing actual Penguin image-processor parameters: {missing}")
+    kind = values.get("image_processor_type", "PenguinVLImageProcessor")
+    if kind != "PenguinVLImageProcessor":
+        raise ValueError(f"Unexpected image_processor_type: {kind!r}")
+    # Loading uses the pinned GitHub class directly. Old auto_map metadata is
+    # not executable and must never redirect imports into a checkpoint folder.
+    values.pop("auto_map", None)
+    return values

modeling_loopvl.py

@@ -0,0 +1,169 @@
+"""Load LoopVL's unchanged architecture from a data-only weight snapshot.
+
+This is a small explicit loader, not a claim of native Transformers AutoModel
+support. All model parameters and persistent buffers live in model.safetensors.
+The original hrm_penguin implementation and Penguin classes live in this GitHub
+runtime, not in the weight repository. Loading neither imports checkpoint code
+nor needs the old split checkpoint or the network.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+from transformers import AutoConfig, AutoTokenizer
+from transformers import HrmTextForCausalLM
+
+from metadata import COMPONENT_DTYPES, read_checkpoint_metadata, read_processor_config
+from penguin_encoder.configuration_penguinvl_encoder import PenguinVLVisionEncoderConfig
+from penguin_encoder.image_processing_penguinvl import PenguinVLImageProcessor
+from penguin_encoder.modeling_penguinvl_encoder import (
+    PenguinVLVisionEncoderModel,
+    VisualRotaryEmbedding,
+)
+
+from hrm_penguin.configuration import HrmPenguinConfig
+from hrm_penguin.model import HrmPenguinForConditionalGeneration
+from hrm_penguin.recurrent.hrm_multimodal import MultimodalHrmTextModel
+from hrm_penguin.vision.encoder import (
+    PenguinVisionEncoder,
+    compute_default_rope_parameters,
+    install_penguin_sdpa,
+    install_penguin_transformers_compatibility,
+)
+from hrm_penguin.vision.preprocessing import build_penguin_image_processor
+from hrm_penguin.vision.projector import VisionProjector
+
+
+def read_config(model_dir: str | Path) -> tuple[dict, HrmPenguinConfig]:
+    """Read embedded model configs without importing code from model_dir."""
+    root = Path(model_dir).expanduser().resolve()
+    raw = read_checkpoint_metadata(root)
+    source = raw["loopvl_config"]
+    config = HrmPenguinConfig.from_dict(source)
+    # Metadata paths only: both towers are constructed from embedded configs.
+    # No tower-specific directory or Python file is required in a v2 snapshot.
+    config.vision.model_path = str(root)
+    config.language.model_path = str(root)
+    config.validate()
+    return raw, config
+
+
+def load_tokenizer_and_processor(model_dir: str | Path, config=None):
+    """Use the original tokenizer, image preprocessor and dynamic-resolution policy."""
+    root = Path(model_dir).expanduser().resolve()
+    raw = read_checkpoint_metadata(root)
+    if config is None:
+        _, config = read_config(root)
+    tokenizer = AutoTokenizer.from_pretrained(str(root), local_files_only=True)
+    official = PenguinVLImageProcessor.from_dict(read_processor_config(root, raw))
+    processor = build_penguin_image_processor(official, config.vision)
+    return tokenizer, processor
+
+
+# Short public synonym used by the bundled preprocessor metadata.
+load_components = load_tokenizer_and_processor
+
+
+class LoopVLForConditionalGeneration(HrmPenguinForConditionalGeneration):
+    """The existing LoopVL computation graph with a single-weight-file loader.
+
+    Usage::
+
+        model = LoopVLForConditionalGeneration.from_pretrained("./model", device="cuda:0")
+
+    No dtype/quantization override is offered: the checkpoint intentionally mixes
+    BF16 language/vision parameters with FP32 projector, visual gate and anchor.
+    Moving the model to a device is safe; calling model.half()/bfloat16() would
+    change those FP32 adapters and is not a lossless operation.
+    """
+
+    @classmethod
+    def from_pretrained(cls, model_dir: str | Path, *, device="cpu"):
+        root = Path(model_dir).expanduser().resolve()
+        raw, config = read_config(root)
+        if torch.get_default_dtype() != torch.float32:
+            raise ValueError("Load with torch's default dtype set to torch.float32")
+
+        # Build under HF's BF16 construction context, not model.to(BF16).
+        # Penguin explicitly creates FP32 nonpersistent rotary buffers; converting
+        # the entire vision module afterward would round those frequencies.
+        install_penguin_transformers_compatibility()
+        # Patch the exact static class instantiated below, rather than a class
+        # copied from a weight repository into transformers' remote-code cache.
+        if not hasattr(VisualRotaryEmbedding, "compute_default_rope_parameters"):
+            VisualRotaryEmbedding.compute_default_rope_parameters = staticmethod(
+                compute_default_rope_parameters
+            )
+        vision_values = dict(raw["vision_config"])
+        vision_values.pop("auto_map", None)
+        vision_config = PenguinVLVisionEncoderConfig.from_dict(vision_values)
+        # AutoModel.from_config previously dispatched to this same _from_config
+        # method. Keep its dtype construction context and attention setup.
+        vision_model = PenguinVLVisionEncoderModel._from_config(
+            vision_config,
+            dtype=torch.bfloat16,
+            attn_implementation=config.vision.attn_implementation,
+        )
+        if config.vision.attn_implementation == "sdpa":
+            install_penguin_sdpa(vision_model)
+        vision = PenguinVisionEncoder(vision_model, hidden_size=config.vision.hidden_size)
+
+        text_values = dict(raw["text_config"])
+        model_type = text_values.pop("model_type")
+        if model_type != "hrm_text":
+            raise ValueError(f"Unexpected language model_type: {model_type!r}")
+        text_config = AutoConfig.for_model(model_type, **text_values)
+        text_config.max_position_embeddings = config.language.max_position_embeddings
+        text_config.rope_theta = config.language.rope_theta
+        if config.language.rope_parameters:
+            text_config.rope_parameters = dict(config.language.rope_parameters)
+        text_config.use_cache = bool(config.inference.use_cache)
+        language = HrmTextForCausalLM._from_config(
+            text_config,
+            dtype=torch.bfloat16,
+            attn_implementation=config.language.attn_implementation,
+        )
+
+        # Reproduce the historical construction order exactly. The multimodal
+        # model is first constructed in FP32 and then converted to BF16, including
+        # its nonpersistent hybrid RoPE buffers. Only adapters are restored to FP32.
+        recurrent = config.recurrent_visual
+        language.model = MultimodalHrmTextModel.from_text_model(
+            language.model,
+            rope_mode=recurrent.rope_mode,
+            use_visual_anchor=recurrent.use_visual_anchor,
+            use_visual_gate=recurrent.use_visual_gate,
+            gate_dim=recurrent.gate_dim,
+            anchor_raw_init=recurrent.anchor_raw_init,
+            anchor_max_scale=recurrent.anchor_max_scale,
+            gate_saturation_metric_lower=recurrent.gate_saturation_metric_lower,
+            gate_saturation_metric_upper=recurrent.gate_saturation_metric_upper,
+        ).to(dtype=torch.bfloat16)
+        language.model.visual_gate.float()
+        language.model.visual_anchor.float()
+        projector = VisionProjector(
+            config.vision.hidden_size, config.language.hidden_size
+        ).float()
+        # visual_embedding_scale is a persistent checkpoint buffer, so loading
+        # its stored value is restored alongside the model parameters.
+        model = cls(vision, language, projector, config)
+
+        weights = load_file(str(root / "model.safetensors"), device="cpu")
+        expected = model.state_dict()
+        missing = sorted(set(expected) - set(weights))
+        unexpected = sorted(set(weights) - set(expected))
+        if missing or unexpected:
+            raise RuntimeError(f"Checkpoint key mismatch: missing={missing}, unexpected={unexpected}")
+        for name, tensor in weights.items():
+            target = expected[name]
+            if tensor.shape != target.shape or tensor.dtype != target.dtype:
+                raise RuntimeError(
+                    f"Refusing implicit conversion for {name}: checkpoint "
+                    f"{tuple(tensor.shape)} {tensor.dtype}, model {tuple(target.shape)} {target.dtype}"
+                )
+        model.load_state_dict(weights, strict=True)
+        del expected, weights
+        model.checkpoint_directory = str(root)
+        return model.to(device=torch.device(device)).eval().requires_grad_(False)

penguin_encoder/LICENSE.txt

@@ -0,0 +1,289 @@
+Tencent is pleased to support the open source community by making  Penguin-VL 高效视觉语言理解模型 available.
+
+Copyright (C) 2026 Tencent.  All rights reserved. 
+
+The open-source software and/or models included in this distribution may have been modified by Tencent (“Tencent Modifications”). All Tencent Modifications are Copyright (C) Tencent.
+
+ Penguin-VL 高效视觉语言理解模型 is licensed under the License Terms of  Penguin-VL 高效视觉语言理解模型, except for the third-party components listed below, which remain licensed under their respective original terms.  Penguin-VL 高效视觉语言理解模型 does not impose any additional restrictions beyond those specified in the original licenses of these third-party components. Users are required to comply with all applicable terms and conditions of the original licenses and to ensure that the use of these third-party components conforms to all relevant laws and regulations.
+
+For the avoidance of doubt,  Penguin-VL 高效视觉语言理解模型 refers solely to training code, inference code, parameters, and weights made publicly available by Tencent in accordance with the License Terms of  Penguin-VL 高效视觉语言理解模型. 
+
+Terms of the License Terms of  Penguin-VL 高效视觉语言理解模型:
+--------------------------------------------------------------------
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+0. Additional Territorial Limitation ** Penguin-VL 高效视觉语言理解模型 IS NOT INTENDED FOR USE WITHIN THE EUROPEAN UNION.** IN THE EVENT OF ANY CONFLICT, THIS CLAUSE SHALL PREVAIL.
+
+1. Definitions.
+
+"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document.
+
+"Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License.
+
+"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%!)(MISSING) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
+
+"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
+
+"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
+
+"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
+
+"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below).
+
+"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof.
+
+"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
+
+"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
+
+You must give any other recipients of the Work or Derivative Works a copy of this License; and 
+
+You must cause any modified files to carry prominent notices stating that You changed the files; and 
+
+You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and 
+
+If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. 
+
+You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 
+
+5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
+
+
+END OF TERMS AND CONDITIONS
+
+
+
+Dependencies and Licenses:
+
+This open-source project,  Penguin-VL 高效视觉语言理解模型, builds upon the following open-source models and/or software components, each of which remains licensed under its original license. Certain models or software may include modifications made by Tencent (“Tencent Modifications”), which are Copyright (C) Tencent.
+
+In case you believe there have been errors in the attribution below, you may submit the concerns to us for review and correction.
+
+
+Open Source Model Licensed under the Apache License 2.0:
+--------------------------------------------------------------------
+1. Qwen3
+Copyright 2024 Alibaba Cloud
+Please note this model has been modified by Tencent in this distribution.
+
+ 
+
+Terms of the Apache License 2.0：
+--------------------------------------------------------------------
+Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+
+

penguin_encoder/__init__.py

@@ -0,0 +1 @@
+"""Pinned Penguin vision implementation, distributed as code in GitHub only."""

penguin_encoder/configuration_penguinvl_encoder.py

@@ -0,0 +1,33 @@
+"""PenguinVL vision encoder model configuration."""
+
+from transformers import Qwen3Config
+
+
+class PenguinVLVisionEncoderConfig(Qwen3Config):
+
+    model_type = "penguinvl_vision_encoder"
+
+    def __init__(
+        self,
+        hidden_size=1536,
+        intermediate_size=8960,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        num_channels=3,
+        patch_size=14,
+        layer_norm_eps=1e-6,
+        attention_dropout=0.0,
+        num_key_value_heads=2,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.attention_dropout = attention_dropout
+        self.num_key_value_heads = num_key_value_heads
+        self.layer_norm_eps = layer_norm_eps

penguin_encoder/image_processing_penguinvl.py

@@ -0,0 +1,548 @@
+# Adopted from https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py.
+# Below is the original copyright:
+# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Image processor class for PenguinVL."""
+
+import math
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+
+import torch
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
+from transformers.image_utils import ImageInput
+from transformers.image_transforms import (
+    convert_to_rgb,
+    resize,
+    to_channel_dimension_format,
+)
+from transformers.image_utils import (
+    OPENAI_CLIP_MEAN,
+    OPENAI_CLIP_STD,
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    get_image_size,
+    infer_channel_dimension_format,
+    is_scaled_image,
+    is_valid_image,
+    make_list_of_images,
+    to_numpy_array,
+)
+try:
+    from transformers.image_utils import VideoInput
+except:
+    from transformers.video_utils import VideoInput
+from transformers.utils import TensorType, is_vision_available, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+if is_vision_available():
+    from PIL import Image
+
+
+def is_valid_video(video) -> bool:
+    if isinstance(video, (list, tuple)):
+        return all(is_valid_image(frame) for frame in video)
+    elif isinstance(video, np.ndarray):
+        return video.ndim == 4
+    elif isinstance(video, torch.Tensor):
+        return video.ndim == 4
+    return False
+
+
+def make_batched_images(images) -> List[List[ImageInput]]:
+    """
+    Normalize visual inputs to ``List[List[ImageInput]]`` – a list of *clips*,
+    where each clip is a list of frames.
+
+    Supported input formats::
+
+        Nested clips : [[image], [f1, f2, ...], ...]   → returned as-is
+        Flat frames  : [f1, f2, ...]                    → [[f1, f2, ...]]
+        Single image : image                            → [[image]]
+
+    Returns:
+        List of clips, where each clip is a list of valid images / frames.
+    """
+    if isinstance(images, (list, tuple)) and len(images) > 0:
+        if isinstance(images[0], (list, tuple)):
+            return [list(clip) for clip in images]
+        if all(is_valid_image(f) for f in images):
+            return [list(images)]
+    if is_valid_image(images):
+        return [[images]]
+    raise ValueError(f"Could not make batched images from {images}")
+
+
+def simple_batched_resize(
+    images, 
+    factor: int = 28, 
+    min_tokens: int = 4 * 4, 
+    max_tokens: int = 16384, 
+    input_data_format: str = None, 
+    frame_types=None
+):
+    """
+    Compute per-frame target (h, w) for a video frame list under a token budget (key/intermediate may differ).
+
+    Uses the Temporal Redundancy-Aware (TRA) token compression strategy: key and intermediate frames
+    can have different target areas (e.g. 1:16 ratio when compressing) to stay within max_tokens.
+
+    Args:
+        images: List of video frames (each PIL Image or ndarray).
+        factor: Alignment granularity (height and width are multiples of factor), default 28.
+        min_tokens: Minimum tokens per frame (used to derive min_pixels), default 16.
+        max_tokens: Token cap for total pixel budget, default 16384.
+        input_data_format: Channel format when not PIL, e.g. "channels_first".
+        frame_types: Per-frame type list, 0=key, 1=intermediate; None means all key.
+
+    Returns:
+        image_sizes: List of (h, w) per frame, one-to-one with images.
+    """
+    min_pixels = min_tokens * factor * factor * 1.5
+    max_pixels = max_tokens * factor * factor * 0.95
+
+    # --- Base info ---
+    first_image = images[0]
+    if isinstance(first_image, Image.Image):
+        width, height = first_image.size
+    else:
+        height, width = get_image_size(first_image, channel_dim=input_data_format)
+
+    aspect_ratio = height / width
+    raw_area = height * width
+
+    num_frames = len(images)
+    if frame_types is not None:
+        ft_list = frame_types.tolist() if hasattr(frame_types, 'tolist') else frame_types
+        num_intermediate = ft_list.count(1)
+        num_key = ft_list.count(0)
+    else:
+        num_key = num_frames
+        num_intermediate = 0
+        ft_list = [0] * num_frames
+
+    def get_dims_from_area(target_area, ar, fac):
+        """Compute aligned (h, w) from target area and aspect ratio; area = w²·ar => w = sqrt(area/ar)."""
+        w_new = math.sqrt(target_area / ar)
+        h_new = w_new * ar
+
+        h_bar = round(h_new / fac) * fac
+        w_bar = round(w_new / fac) * fac
+        h_bar = max(h_bar, fac)
+        w_bar = max(w_bar, fac)
+
+        return h_bar, w_bar
+
+    # --- Stage 1: No-downscale check ---
+    # If total pixels within budget, keep original size for both key and intermediate frames.
+    total_raw_pixels = num_frames * raw_area
+    target_key_area = raw_area
+    target_intermediate_area = raw_area
+
+    if total_raw_pixels > max_pixels:
+        # --- Stage 2: Sync compression ---
+        # Over budget: compress with 1:16 area ratio, intermediate_area = key_area / 16.
+        # Constraint: N_key·A_key + N_intermediate·(A_key/16) = max_pixels => A_key = max_pixels / (N_key + N_intermediate/16).
+        effective_count = num_key + (num_intermediate / 16.0)
+        calc_key_area = max_pixels / effective_count
+        calc_intermediate_area = calc_key_area / 16.0
+
+        # --- Stage 3: Intermediate-frame floor ---
+        # If computed intermediate area is below min_pixels, pin intermediate to min_pixels and give remaining budget to key.
+        if calc_intermediate_area >= min_pixels:
+            target_key_area = calc_key_area
+            target_intermediate_area = calc_intermediate_area
+        else:
+            target_intermediate_area = min_pixels
+            pixels_taken_by_intermediate = num_intermediate * min_pixels
+            remaining_for_key = max_pixels - pixels_taken_by_intermediate
+            target_key_area = remaining_for_key / num_key
+
+        # --- Stage 4: Key-frame hard floor ---
+        if target_key_area < min_pixels:
+            target_key_area = min_pixels
+
+    # --- Area to aligned dimensions ---
+    k_h, k_w = get_dims_from_area(target_key_area, aspect_ratio, factor)
+    if num_intermediate > 0:
+        i_h, i_w = get_dims_from_area(target_intermediate_area, aspect_ratio, factor)
+    else:
+        i_h, i_w = 0, 0
+
+    def ensure_min_hw(h, w, min_p, raw_ar):
+        """If area still below min_pixels after alignment (rounding), recompute from min area and align upward."""
+        if h * w < min_p:
+            w = math.sqrt(min_p / raw_ar)
+            h = w * raw_ar
+            h = math.ceil(h / factor) * factor
+            w = math.ceil(w / factor) * factor
+        return h, w
+
+    k_h, k_w = ensure_min_hw(k_h, k_w, min_pixels, aspect_ratio)
+    if num_intermediate > 0:
+        i_h, i_w = ensure_min_hw(i_h, i_w, min_pixels, aspect_ratio)
+
+    image_sizes = [
+        (i_h, i_w) if ft_list[i] == 1 else (k_h, k_w)
+        for i in range(num_frames)
+    ]
+    return image_sizes
+
+
+class PenguinVLImageProcessor(BaseImageProcessor):
+    r"""
+    Constructs a PenguinVL image processor that dynamically resizes images based on the original images.
+
+    Args:
+        do_resize (`bool`, *optional*, defaults to `True`):
+            Whether to resize the image's (height, width) dimensions.
+        resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):
+            Resampling filter to use when resizing the image.
+        do_rescale (`bool`, *optional*, defaults to `True`):
+            Whether to rescale the image by the specified scale `rescale_factor`.
+        rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
+            Scale factor to use if rescaling the image.
+        do_normalize (`bool`, *optional*, defaults to `True`):
+            Whether to normalize the image.
+        image_mean (`float` or `List[float]`, *optional*, defaults to `[0.48145466, 0.4578275, 0.40821073]`):
+            Mean to use if normalizing the image. This is a float or list of floats for each channel in the image.
+        image_std (`float` or `List[float]`, *optional*, defaults to `[0.26862954, 0.26130258, 0.27577711]`):
+            Standard deviation to use if normalizing the image. This is a float or list of floats for each channel in the image.
+        do_convert_rgb (`bool`, *optional*, defaults to `True`):
+            Whether to convert the image to RGB.
+        min_pixels (`int`, *optional*, defaults to `56 * 56`):
+            The min pixels of the image to resize the image.
+        max_pixels (`int`, *optional*, defaults to `28 * 28 * 1280`):
+            The max pixels of the image to resize the image.
+        patch_size (`int`, *optional*, defaults to 14):
+            The spacial patch size of the vision encoder.
+    """
+
+    model_input_names = ["pixel_values", "grid_sizes", "merge_sizes"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        resample: PILImageResampling = PILImageResampling.BICUBIC,
+        do_rescale: bool = True,
+        rescale_factor: Union[int, float] = 1 / 255,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: bool = True,
+        min_tokens: int = 4 * 4,
+        max_tokens: int = 16384,
+        patch_size: int = 14,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.do_resize = do_resize
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN
+        self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD
+        self.min_tokens = min_tokens
+        self.max_tokens = max_tokens
+        self.patch_size = patch_size
+        self.do_convert_rgb = do_convert_rgb
+
+    def _allocate_token_budget(self, clips, clip_merge_sizes, input_data_format):
+        """Distribute self.max_tokens across clips proportionally to their raw token counts."""
+        clip_raw_tokens = []
+        for clip, ms in zip(clips, clip_merge_sizes):
+            first_frame = clip[0]
+            if isinstance(first_frame, Image.Image):
+                w, h = first_frame.size
+            else:
+                h, w = get_image_size(first_frame, channel_dim=input_data_format)
+            factor = self.patch_size * ms
+            clip_raw_tokens.append(len(clip) * h * w / (factor * factor))
+
+        total_raw_tokens = sum(clip_raw_tokens)
+        if total_raw_tokens <= self.max_tokens:
+            return [self.max_tokens] * len(clips)
+
+        return [
+            max(self.min_tokens * len(clip), raw * self.max_tokens / total_raw_tokens)
+            for clip, raw in zip(clips, clip_raw_tokens)
+        ]
+
+    def _preprocess(
+        self,
+        images: Union[ImageInput, VideoInput],
+        target_size: List[int],
+        merge_size: int = 1,
+        do_resize: bool = None,
+        resample: PILImageResampling = None,
+        do_rescale: bool = None,
+        rescale_factor: float = None,
+        do_normalize: bool = None,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: bool = None,
+        data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+    ):
+        """
+        Preprocess an image or batch of images. Copy of the `preprocess` method from `CLIPImageProcessor`.
+
+        Args:
+            images (`ImageInput`):
+                Image or batch of images to preprocess. Expects pixel values ranging from 0 to 255. If pixel values range from 0 to 1, set `do_rescale=False`.
+            target_size (`List[int]`):
+                The target size to resize the image to. Should be a list of two integers: [target_height, target_width].
+            merge_size (`int`, *optional*, defaults to `1`):
+                The merge size after the vision encoder.
+            do_resize (`bool`, *optional*, defaults to `self.do_resize`):
+                Whether to resize the image.
+            resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
+                Resampling filter to use if resizing the image. This can be one of the `PILImageResampling` enums.
+            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+                Whether to rescale the image.
+            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+                Scale factor to use if rescaling the image.
+            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
+                Whether to normalize the image.
+            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
+                Mean to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image.
+            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
+                Standard deviation to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image.
+            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+                Whether to convert the image to RGB.
+            data_format (`ChannelDimension`, *optional*, defaults to `ChannelDimension.FIRST`):
+                The channel dimension format for the output image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - Unset: Use the channel dimension format of the input image.
+            input_data_format (`ChannelDimension` or `str`, *optional*):
+                The channel dimension format for the input image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.   - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+        """
+        images = make_list_of_images(images)
+
+        if do_convert_rgb:
+            images = [convert_to_rgb(image) for image in images]
+
+        # All transformations expect numpy arrays.
+        images = [to_numpy_array(image) for image in images]
+
+        if is_scaled_image(images[0]) and do_rescale:
+            logger.warning_once(
+                "It looks like you are trying to rescale already rescaled images. If the input"
+                " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
+            )
+        if input_data_format is None:
+            # We assume that all images have the same channel dimension format.
+            input_data_format = infer_channel_dimension_format(images[0])
+
+        height, width = get_image_size(images[0], channel_dim=input_data_format)
+        resized_height, resized_width = height, width
+        processed_images = []
+        for image in images:
+            if do_resize:
+                resized_height, resized_width = target_size
+                image = resize(
+                    image, size=(resized_height, resized_width), resample=resample, input_data_format=input_data_format
+                )
+
+            if do_rescale:
+                image = self.rescale(image, scale=rescale_factor, input_data_format=input_data_format)
+
+            if do_normalize:
+                image = self.normalize(
+                    image=image, mean=image_mean, std=image_std, input_data_format=input_data_format
+                )
+
+            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        patches = np.array(processed_images)
+        if data_format == ChannelDimension.LAST:
+            patches = patches.transpose(0, 3, 1, 2)
+        t = patches.shape[0]
+        channel = patches.shape[1]
+        grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size
+        patches = patches.reshape(
+            t,
+            channel,
+            grid_h // merge_size,
+            merge_size,
+            self.patch_size,
+            grid_w // merge_size,
+            merge_size,
+            self.patch_size,
+        )
+        patches = patches.transpose(0, 2, 5, 3, 6, 1, 4, 7)
+        flatten_patches = patches.reshape(
+            t * grid_h * grid_w, channel * self.patch_size * self.patch_size
+        )
+
+        return flatten_patches, (t, grid_h, grid_w)
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: bool = None,
+        resample: PILImageResampling = None,
+        do_rescale: bool = None,
+        rescale_factor: float = None,
+        do_normalize: bool = None,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: bool = None,
+        merge_size: Optional[Union[int, List[int]]] = None,
+        frame_types: Optional[Union[int, List[int]]] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+    ):
+        """
+        Args:
+            images (`ImageInput`):
+                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
+                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
+            do_resize (`bool`, *optional*, defaults to `self.do_resize`):
+                Whether to resize the image.
+            resample (`int`, *optional*, defaults to `self.resample`):
+                Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only
+                has an effect if `do_resize` is set to `True`.
+            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+                Whether to rescale the image.
+            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+                Rescale factor to rescale the image by if `do_rescale` is set to `True`.
+            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
+                Whether to normalize the image.
+            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
+                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
+                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to
+                `True`.
+            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+                Whether to convert the image to RGB.
+            return_tensors (`str` or `TensorType`, *optional*):
+                The type of tensors to return. Can be one of:
+                - Unset: Return a list of `np.ndarray`.
+                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
+                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
+                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
+                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
+            data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
+                The channel dimension format for the output image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - Unset: Use the channel dimension format of the input image.
+            input_data_format (`ChannelDimension` or `str`, *optional*):
+                The channel dimension format for the input image. If unset, the channel dimension format is inferred
+                from the input image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+
+        """
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        merge_size = merge_size if merge_size is not None else self.merge_size
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+
+        clips = make_batched_images(images)
+        num_clips = len(clips)
+
+        if isinstance(merge_size, (list, tuple)):
+            assert len(merge_size) == num_clips, (
+                f"merge_size length ({len(merge_size)}) must match number of clips ({num_clips})"
+            )
+            clip_merge_sizes = list(merge_size)
+        else:
+            clip_merge_sizes = [merge_size] * num_clips
+
+        if frame_types is None:
+            clip_frame_types = [None] * num_clips
+        elif isinstance(frame_types, (list, tuple)) and len(frame_types) > 0:
+            if isinstance(frame_types[0], (list, tuple)) or frame_types[0] is None:
+                assert len(frame_types) == num_clips, (
+                    f"frame_types length ({len(frame_types)}) must match number of clips ({num_clips})"
+                )
+                clip_frame_types = list(frame_types)
+            else:
+                assert num_clips == 1, "Flat frame_types is only supported for a single clip"
+                clip_frame_types = [frame_types]
+        else:
+            clip_frame_types = [None] * num_clips
+
+        pixel_values, grid_sizes, per_frame_merge_sizes = [], [], []
+
+        clip_max_tokens_list = self._allocate_token_budget(
+            clips, clip_merge_sizes, input_data_format,
+        )
+
+        for clip, ms, ft, clip_max_tokens in zip(clips, clip_merge_sizes, clip_frame_types, clip_max_tokens_list):
+            target_sizes = simple_batched_resize(
+                clip,
+                factor=self.patch_size * ms,
+                min_tokens=self.min_tokens,
+                max_tokens=clip_max_tokens,
+                input_data_format=input_data_format,
+                frame_types=ft,
+            )
+
+            for frame, target_size in zip(clip, target_sizes):
+                patches, grid_size = self._preprocess(
+                    frame,
+                    target_size=target_size,
+                    merge_size=ms,
+                    do_resize=do_resize,
+                    resample=resample,
+                    do_rescale=do_rescale,
+                    rescale_factor=rescale_factor,
+                    do_normalize=do_normalize,
+                    image_mean=image_mean,
+                    image_std=image_std,
+                    data_format=data_format,
+                    do_convert_rgb=do_convert_rgb,
+                    input_data_format=input_data_format,
+                )
+                pixel_values.append(patches)
+                grid_sizes.append(grid_size)
+                per_frame_merge_sizes.append(ms)
+
+        pixel_values = np.concatenate(pixel_values, axis=0)
+        grid_sizes = np.array(grid_sizes)
+        merge_sizes = np.array(per_frame_merge_sizes)
+
+        data = {
+            "pixel_values": pixel_values,
+            "grid_sizes": grid_sizes,
+            "merge_sizes": merge_sizes,
+        }
+
+        return BatchFeature(data=data, tensor_type=return_tensors)

penguin_encoder/modeling_penguinvl_encoder.py

@@ -0,0 +1,522 @@
+from torch import nn
+import torch
+import math
+import warnings
+from .configuration_penguinvl_encoder import PenguinVLVisionEncoderConfig
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.qwen3.modeling_qwen3 import Qwen3Model, Qwen3Attention, rotate_half, Qwen3DecoderLayer
+from typing import List, Optional, Tuple, Union
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.processing_utils import Unpack
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.utils import logging, is_flash_attn_greater_or_equal_2_10, is_flash_attn_2_available
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from torch.nn.init import _calculate_fan_in_and_fan_out
+import torch.nn.functional as F
+if is_flash_attn_2_available():
+    from transformers.modeling_flash_attention_utils import _flash_attention_forward
+    from flash_attn import flash_attn_varlen_func
+
+logger = logging.get_logger(__name__)
+
+class PenguinVLVisionEncoderEmbeddings(nn.Module):
+
+    def __init__(self, config: PenguinVLVisionEncoderConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.patch_size = config.patch_size
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            padding="valid",
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.view(
+            -1, self.config.num_channels, self.patch_size, self.patch_size
+        )
+        patch_embeds = self.patch_embedding(hidden_states)
+        embeddings = patch_embeds.view(-1, self.embed_dim)
+
+        return embeddings
+    
+
+# Adapted from Qwen2VLRotaryEmbedding in transformers/models/qwen2/modeling_qwen2.py
+class VisualRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim=None,
+        max_position_embeddings=2048,
+        base=10000,
+        device=None,
+        scaling_factor=1.0,
+        rope_type="default",
+        config = None,
+    ):
+        super().__init__()
+        # TODO (joao): remove the `if` below, only used for BC
+        self.rope_kwargs = {}
+        if config is None:
+            logger.warning_once(
+                "`Qwen2VLRotaryEmbedding` can now be fully parameterized by passing the model config through the "
+                "`config` argument. All other arguments will be removed in v4.46"
+            )
+            self.rope_kwargs = {
+                "rope_type": rope_type,
+                "factor": scaling_factor,
+                "dim": dim,
+                "base": base,
+                "max_position_embeddings": max_position_embeddings,
+            }
+            self.rope_type = rope_type
+            self.max_seq_len_cached = max_position_embeddings
+            self.original_max_seq_len = max_position_embeddings
+        else:
+            # BC: "rope_type" was originally "type"
+            if config.rope_scaling is not None:
+                self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+            else:
+                self.rope_type = "default"
+            self.max_seq_len_cached = config.max_position_embeddings
+            self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, **self.rope_kwargs)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    def _dynamic_frequency_update(self, position_ids, device):
+        """
+        dynamic RoPE layers should recompute `inv_freq` in the following situations:
+        1 - growing beyond the cached sequence length (allow scaling)
+        2 - the current sequence length is in the original scale (avoid losing precision with small sequences)
+        """
+        seq_len = torch.max(position_ids) + 1
+        if seq_len > self.max_seq_len_cached:  # growth
+            inv_freq, self.attention_scaling = self.rope_init_fn(
+                self.config, device, seq_len=seq_len, **self.rope_kwargs
+            )
+            self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO joao: may break with compilation
+            self.max_seq_len_cached = seq_len
+
+        if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:  # reset
+            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
+            self.max_seq_len_cached = self.original_max_seq_len
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        if "dynamic" in self.rope_type:
+            self._dynamic_frequency_update(position_ids, device=x.device)
+
+        inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(2, position_ids.shape[1], -1, 1)
+        position_ids_expanded = position_ids[:, :, None, :].float()  # shape (2, bs, 1, positions)
+        # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
+        device_type = x.device.type
+        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+
+        # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
+        cos = cos * self.attention_scaling
+        sin = sin * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def apply_multimodal_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    rope_section = [cos.shape[-1] // 2, cos.shape[-1] // 2]
+    cos = torch.cat([m[i % 2] for i, m in enumerate(cos.split(rope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim)
+    sin = torch.cat([m[i % 2] for i, m in enumerate(sin.split(rope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim)
+
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class PenguinVLAttention(Qwen3Attention):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.is_causal = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_value: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_multimodal_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        # This is before the transpose
+        seq_len = query_states.shape[2]
+
+        # Restore the attention dtype if normalization promoted hidden states to FP32.
+        target_dtype = None
+        if query_states.dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = next(layer for layer in self.modules() if isinstance(layer, torch.nn.Linear)).weight.dtype
+
+        # FA2 always relies on the value set in the module, so remove it if present in kwargs to avoid passing it twice
+        kwargs.pop("is_causal", None)
+
+        # Reashape to the expected shape for Flash Attention
+        query_states = query_states.transpose(1, 2).squeeze(0)
+        key_states = key_states.transpose(1, 2).squeeze(0)
+        value_states = value_states.transpose(1, 2).squeeze(0)
+
+        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+        attn_output = flash_attn_varlen_func(
+            query_states,
+            key_states,
+            value_states,
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            dropout_p=0.0 if not self.training else self.attention_dropout,
+            causal=self.is_causal
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, None
+    
+
+class PenguinVLDecoderLayer(Qwen3DecoderLayer):
+    def __init__(self, config: PenguinVLVisionEncoderConfig, layer_idx: int):
+        super(PenguinVLDecoderLayer, self).__init__(config, layer_idx)
+        self.self_attn = PenguinVLAttention(config, layer_idx)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
+        cu_seqlens: Optional[torch.Tensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            cu_seqlens=cu_seqlens,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        return outputs
+    
+    
+class PenguinVLVisionEncoderFromQwen3Model(Qwen3Model):
+    def __init__(self, config: PenguinVLVisionEncoderConfig):
+        super().__init__(config)
+        self.layers = nn.ModuleList(
+            [PenguinVLDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.rotary_emb = VisualRotaryEmbedding(config=config)
+        del self.embed_tokens
+
+    @staticmethod
+    def _prepare_4d_causal_attention_mask_with_cache_position(
+        attention_mask: torch.Tensor,
+        sequence_length: int,
+        target_length: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        cache_position: torch.Tensor,
+        batch_size: int,
+        config: PenguinVLVisionEncoderConfig,
+        past_key_values: Cache,
+    ):
+        """
+        Override the original causal mask method to create full attention mask instead.
+        Creates a full attention 4D mask of shape `(batch_size, 1, query_length, key_value_length)` 
+        from a 2D mask of shape `(batch_size, key_value_length)`.
+        
+        For vision encoding, we want full attention between all patches, not causal attention.
+        """
+        if attention_mask is not None and attention_mask.dim() == 4:
+            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+            full_attention_mask = attention_mask
+        else:
+            # Create full attention mask (all zeros, meaning attend to all positions)
+            # We only mask based on the provided attention_mask for padding
+            if attention_mask is not None:
+                # Use the provided attention_mask to handle padding
+                min_dtype = torch.finfo(dtype).min
+                full_attention_mask = torch.zeros(
+                    (sequence_length, target_length), dtype=dtype, device=device
+                )
+                # Expand to 4D
+                full_attention_mask = full_attention_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
+                
+                # Apply padding mask if provided
+                full_attention_mask = full_attention_mask.clone()  # copy to contiguous memory for in-place edit
+                if attention_mask.shape[-1] > target_length:
+                    attention_mask = attention_mask[:, :target_length]
+                mask_length = attention_mask.shape[-1]
+                padding_mask = attention_mask[:, None, None, :] == 0
+                full_attention_mask[:, :, :, :mask_length] = full_attention_mask[:, :, :, :mask_length].masked_fill(
+                    padding_mask, min_dtype
+                )
+            else:
+                # No attention mask provided, create all-zeros mask (full attention)
+                full_attention_mask = torch.zeros(
+                    (batch_size, 1, sequence_length, target_length), dtype=dtype, device=device
+                )
+        return full_attention_mask
+
+    def get_rope_index(self, grid_sizes, merge_sizes, position_ids):
+        position_ids = position_ids.contiguous()
+        batch_size = grid_sizes.shape[0]
+
+        # Vision Part: Generate 2D position indices for vision tokens
+        vision_pos_ids = []
+        for (t, h, w), merge_size in zip(grid_sizes, merge_sizes):
+            # Generate height position indices
+            hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w).to(position_ids.device)
+            hpos_ids = hpos_ids.reshape(
+                h // merge_size,
+                merge_size,
+                w // merge_size,
+                merge_size,
+            )
+            hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+            hpos_ids = hpos_ids.flatten()
+
+            # Generate width position indices  
+            wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1).to(position_ids.device)
+            wpos_ids = wpos_ids.reshape(
+                h // merge_size,
+                merge_size,
+                w // merge_size,
+                merge_size,
+            )
+            wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+            wpos_ids = wpos_ids.flatten()
+            
+            # Stack height and width to create 2D positions
+            vision_pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
+
+        num_start_idx = 0
+        for batch_idx in range(batch_size):
+            pos_len = vision_pos_ids[batch_idx].shape[0]
+            position_ids[:, 0, num_start_idx: num_start_idx+pos_len] = vision_pos_ids[batch_idx].permute(1, 0)
+            num_start_idx += pos_len
+        
+        return position_ids
+
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        grid_sizes: Optional[torch.Tensor] = None,
+        merge_sizes: Optional[torch.Tensor] = None,
+        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+    ) -> BaseModelOutputWithPast:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache
+        if not isinstance(past_key_values, (type(None), Cache)):
+            raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        # the hard coded `2` is for temporal, height and width.
+        if position_ids is None:
+            position_ids = cache_position.view(1, 1, -1).expand(2, inputs_embeds.shape[0], -1)
+        elif position_ids.dim() == 2:
+            position_ids = position_ids[None, ...].expand(2, position_ids.shape[0], -1)
+        position_ids = self.get_rope_index(grid_sizes, merge_sizes, position_ids)
+
+        causal_mask = None
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
+        # Calculate cumulative sequence lengths for the grid sizes
+        cu_seqlens = torch.repeat_interleave(grid_sizes[:, 1] * grid_sizes[:, 2], grid_sizes[:, 0]).cumsum(dim=0, dtype=torch.int32)
+        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                cu_seqlens=cu_seqlens,
+                **flash_attn_kwargs,
+            )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class PenguinVLVisionEncoderModel(PreTrainedModel):
+
+    config_class = PenguinVLVisionEncoderConfig
+    base_model_prefix = "penguinvl_vision_encoder"
+    main_input_name = "pixel_values"
+    _no_split_modules = [
+        "PenguinVLVisionEncoderEmbeddings",
+    ]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+
+    def __init__(self, config: PenguinVLVisionEncoderConfig):
+        super().__init__(config=config)
+        self.embeddings = PenguinVLVisionEncoderEmbeddings(config)
+        self.encoder = PenguinVLVisionEncoderFromQwen3Model(config)
+
+        self.post_init()
+
+
+    def forward(self, pixel_values, grid_sizes, merge_sizes=None) -> torch.Tensor:
+        hidden_states = self.embeddings(pixel_values)
+        encoder_output = self.encoder(
+            inputs_embeds=hidden_states[None, ...],
+            grid_sizes=grid_sizes,
+            merge_sizes=merge_sizes,
+            output_hidden_states=True,
+        )
+        hidden_states = encoder_output.hidden_states
+        hidden_states = hidden_states[-1].squeeze(0)
+
+        hidden_states_chunks = hidden_states.split(grid_sizes.prod(dim=1).tolist(), dim=0)
+        outputs = []
+
+        for hidden_states, grid_size, merge_size in zip(hidden_states_chunks, grid_sizes, merge_sizes):
+            # NOTE: previous implementation, which supports downsampling with any factor
+            c = hidden_states.shape[-1]
+            hidden_states = hidden_states.view(
+                grid_size[0], grid_size[1] // merge_size, grid_size[2] // merge_size, merge_size, merge_size,  c
+            ).permute(0, 1, 3, 2, 4, 5)
+            hidden_states = hidden_states.reshape(
+                grid_size[0], grid_size[1], grid_size[2], c
+            ).permute(0, 3, 1, 2)
+            hidden_states = torch.nn.functional.interpolate(
+                hidden_states,
+                size=(grid_size[1] // merge_size, grid_size[2] // merge_size),
+                mode='bilinear'
+            )
+            hidden_states = hidden_states.permute(0, 2, 3, 1).view(-1, c)
+
+            # NOTE: simplified implementation, which only supports downsampling with integer factor
+            # NOTE: this implementation is mathematically equivalent to the previous one when merge_size is 1 or 2 but may cause slightly different results
+            # hidden_states = hidden_states.view(-1, merge_size * merge_size, hidden_states.size(-1))
+            # hidden_states = hidden_states.mean(dim=1)
+
+            outputs.append(hidden_states)
+        return torch.cat(outputs, dim=0)

requirements.txt

@@ -0,0 +1,5 @@
+transformers==5.15.1
+safetensors==0.8.0
+pillow==12.2.0
+pyyaml==6.0.3
+accelerate==1.15.0