Upload oculus_unified_model/modeling_oculus.py with huggingface_hub

oculus_unified_model/modeling_oculus.py

@@ -0,0 +1,842 @@
+"""
+Oculus Unified Model
+
+HuggingFace-compatible vision-language model with:
+- Multi-encoder vision (DINOv3 + SigLIP2)
+- Trained projector for vision-to-language
+- Optional reasoning with thinking traces
+- Multiple output modes (Text, Point, Box, Polygon)
+- Focus/Zoom tool calling for fine-grained perception
+"""
+
+import os
+import json
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Tuple, List, Dict, Any, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import (
+    PreTrainedModel,
+    PretrainedConfig,
+    AutoImageProcessor,
+    AutoModel,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    GenerationConfig,
+)
+from transformers.modeling_outputs import BaseModelOutput, CausalLMOutputWithPast
+from PIL import Image
+
+from .configuration_oculus import OculusConfig
+
+
+# ============================================================================
+# Output Data Classes
+# ============================================================================
+
+@dataclass
+class OculusOutput:
+    """Base output class for Oculus model."""
+    text: Optional[str] = None
+    thinking_trace: Optional[str] = None
+    logits: Optional[torch.Tensor] = None
+    hidden_states: Optional[torch.Tensor] = None
+    vision_tokens: Optional[torch.Tensor] = None
+
+
+@dataclass
+class OculusTextOutput(OculusOutput):
+    """Output for text/caption mode."""
+    pass
+
+
+@dataclass
+class OculusPointOutput(OculusOutput):
+    """Output for point detection mode (counting objects)."""
+    points: Optional[List[Tuple[float, float]]] = None
+    labels: Optional[List[str]] = None
+    confidences: Optional[List[float]] = None
+
+
+@dataclass 
+class OculusBoxOutput(OculusOutput):
+    """Output for bounding box detection mode."""
+    boxes: Optional[List[Tuple[float, float, float, float]]] = None  # x1, y1, x2, y2
+    labels: Optional[List[str]] = None
+    confidences: Optional[List[float]] = None
+
+
+@dataclass
+class OculusPolygonOutput(OculusOutput):
+    """Output for polygon/segmentation mode."""
+    polygons: Optional[List[List[Tuple[float, float]]]] = None
+    labels: Optional[List[str]] = None
+    mask: Optional[np.ndarray] = None
+
+
+# ============================================================================
+# Vision Encoder (DINOv3 + SigLIP2)
+# ============================================================================
+
+class OculusVisionEncoder(nn.Module):
+    """
+    Dual vision encoder combining DINOv3 and SigLIP2.
+    
+    DINOv3: Excellent at semantic understanding, object boundaries
+    SigLIP2: Strong at text/language alignment
+    """
+    
+    def __init__(self, config: OculusConfig):
+        super().__init__()
+        self.config = config
+        
+        # Will be loaded lazily
+        self.dinov3 = None
+        self.dinov3_processor = None
+        self.siglip = None
+        self.siglip_processor = None
+        
+        self._loaded = False
+    
+    def load_encoders(self, device: str = "cpu"):
+        """Load vision encoders from HuggingFace."""
+        if self._loaded:
+            return
+            
+        print("[Oculus] Loading vision encoders...")
+        
+        # DINOv3
+        try:
+            self.dinov3_processor = AutoImageProcessor.from_pretrained(
+                self.config.dinov3_model_id
+            )
+            self.dinov3 = AutoModel.from_pretrained(
+                self.config.dinov3_model_id
+            ).eval().to(device)
+            print(f"  ✓ DINOv3: {self.config.dinov3_model_id}")
+        except Exception as e:
+            warnings.warn(f"Failed to load DINOv3: {e}")
+            self.dinov3_processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
+            self.dinov3 = AutoModel.from_pretrained("facebook/dinov2-base").eval().to(device)
+            print("  ✓ DINOv2-base (fallback)")
+        
+        # SigLIP2
+        try:
+            self.siglip_processor = AutoImageProcessor.from_pretrained(
+                self.config.siglip_model_id
+            )
+            self.siglip = AutoModel.from_pretrained(
+                self.config.siglip_model_id
+            ).eval().to(device)
+            print(f"  ✓ SigLIP: {self.config.siglip_model_id}")
+        except Exception as e:
+            warnings.warn(f"Failed to load SigLIP: {e}")
+            from transformers import SiglipVisionModel
+            self.siglip_processor = AutoImageProcessor.from_pretrained("google/siglip-base-patch16-224")
+            self.siglip = SiglipVisionModel.from_pretrained("google/siglip-base-patch16-224").eval().to(device)
+            print("  ✓ SigLIP-base (fallback)")
+        
+        self._loaded = True
+    
+    @torch.no_grad()
+    def forward(self, image: Union[Image.Image, torch.Tensor, np.ndarray]) -> torch.Tensor:
+        """
+        Encode image with both vision encoders and fuse features.
+        
+        Returns:
+            Fused vision features [batch, fused_dim]
+        """
+        if not self._loaded:
+            self.load_encoders()
+        
+        # Handle different input types
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        elif isinstance(image, torch.Tensor):
+            image = Image.fromarray(image.cpu().numpy().astype(np.uint8))
+        
+        if isinstance(image, Image.Image):
+            image = image.convert('RGB')
+        
+        device = next(self.dinov3.parameters()).device
+        
+        # DINOv3 encoding
+        d_inputs = self.dinov3_processor(images=image, return_tensors="pt")
+        d_inputs = {k: v.to(device) for k, v in d_inputs.items()}
+        d_out = self.dinov3(**d_inputs)
+        d_pooled = d_out.pooler_output if hasattr(d_out, 'pooler_output') and d_out.pooler_output is not None else d_out.last_hidden_state[:, 0]
+        
+        # SigLIP encoding
+        s_inputs = self.siglip_processor(images=image, return_tensors="pt")
+        s_inputs = {k: v.to(device) for k, v in s_inputs.items()}
+        
+        if hasattr(self.siglip, 'vision_model'):
+            s_hidden = self.siglip.vision_model.embeddings(s_inputs['pixel_values'])
+            s_pooled = s_hidden.mean(dim=1)
+        else:
+            s_out = self.siglip(**s_inputs)
+            s_pooled = s_out.pooler_output if hasattr(s_out, 'pooler_output') else s_out.last_hidden_state[:, 0]
+        
+        # Fuse features
+        fused = torch.cat([d_pooled, s_pooled], dim=-1)
+        
+        return fused
+
+
+# ============================================================================
+# Vision Projector
+# ============================================================================
+
+class OculusProjector(nn.Module):
+    """
+    Projects fused vision features to language model token space.
+    
+    Converts [batch, fused_dim] → [batch, num_tokens, lm_hidden_size]
+    """
+    
+    def __init__(self, config: OculusConfig):
+        super().__init__()
+        self.config = config
+        
+        fused_dim = config.fused_vision_dim
+        hidden_dim = config.projector_hidden_dim
+        num_tokens = config.num_vision_tokens
+        embed_dim = config.lm_hidden_size
+        
+        self.fc1 = nn.Linear(fused_dim, hidden_dim)
+        self.act1 = nn.GELU()
+        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
+        self.act2 = nn.GELU()
+        self.fc3 = nn.Linear(hidden_dim, num_tokens * embed_dim)
+        self.norm = nn.LayerNorm(embed_dim)
+        
+        self.num_tokens = num_tokens
+        self.embed_dim = embed_dim
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Project vision features to token embeddings.
+        
+        Args:
+            x: Vision features [batch, fused_dim]
+        
+        Returns:
+            Vision tokens [batch, num_tokens, embed_dim]
+        """
+        batch_size = x.shape[0]
+        
+        h = self.fc1(x)
+        h = self.act1(h)
+        h = self.fc2(h)
+        h = self.act2(h)
+        h = self.fc3(h)
+        
+        h = h.reshape(batch_size, self.num_tokens, self.embed_dim)
+        h = self.norm(h)
+        
+        return h
+    
+    @classmethod
+    def from_pretrained(cls, path: str, config: OculusConfig):
+        """Load projector from saved weights."""
+        projector = cls(config)
+        
+        weights_path = Path(path) / "projector.npz"
+        if weights_path.exists():
+            import numpy as np
+            weights = np.load(weights_path, allow_pickle=True)
+            
+            state_dict = {}
+            for key in weights.files:
+                layer_dict = weights[key].item()
+                for param_name, param_val in layer_dict.items():
+                    full_key = f"{key}.{param_name}"
+                    # Convert from MLX array if needed
+                    if hasattr(param_val, 'tolist'):
+                        param_val = np.array(param_val.tolist())
+                    state_dict[full_key] = torch.from_numpy(np.array(param_val))
+            
+            projector.load_state_dict(state_dict, strict=False)
+            print(f"  ✓ Loaded projector from {path}")
+        
+        return projector
+
+
+# ============================================================================
+# Detection/Segmentation Heads
+# ============================================================================
+
+class OculusDetectionHead(nn.Module):
+    """Head for bounding box detection."""
+    
+    def __init__(self, config: OculusConfig):
+        super().__init__()
+        hidden_dim = config.lm_hidden_size
+        num_classes = config.num_detection_classes
+        
+        self.cls_head = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.GELU(),
+            nn.Linear(hidden_dim // 2, num_classes)
+        )
+        
+        self.box_head = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.GELU(),
+            nn.Linear(hidden_dim // 2, 4)  # x1, y1, x2, y2
+        )
+    
+    def forward(self, vision_tokens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Predict boxes and classes from vision tokens.
+        
+        Returns:
+            cls_logits: [batch, num_tokens, num_classes]
+            box_coords: [batch, num_tokens, 4]
+        """
+        cls_logits = self.cls_head(vision_tokens)
+        box_coords = self.box_head(vision_tokens).sigmoid()  # Normalize to [0, 1]
+        return cls_logits, box_coords
+
+
+class OculusPointHead(nn.Module):
+    """Head for point detection (object counting)."""
+    
+    def __init__(self, config: OculusConfig):
+        super().__init__()
+        hidden_dim = config.lm_hidden_size
+        num_classes = config.num_detection_classes
+        
+        self.point_head = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.GELU(),
+            nn.Linear(hidden_dim // 2, 2)  # x, y
+        )
+        
+        self.cls_head = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.GELU(),
+            nn.Linear(hidden_dim // 2, num_classes)
+        )
+        
+        self.conf_head = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 4),
+            nn.GELU(),
+            nn.Linear(hidden_dim // 4, 1)
+        )
+    
+    def forward(self, vision_tokens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        points = self.point_head(vision_tokens).sigmoid()
+        cls_logits = self.cls_head(vision_tokens)
+        confidence = self.conf_head(vision_tokens).sigmoid()
+        return points, cls_logits, confidence
+
+
+class OculusSegmentationHead(nn.Module):
+    """Head for polygon/mask segmentation."""
+    
+    def __init__(self, config: OculusConfig):
+        super().__init__()
+        hidden_dim = config.lm_hidden_size
+        num_classes = config.num_segmentation_classes
+        
+        # Predict mask logits
+        self.mask_head = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, 14 * 14 * num_classes)  # Output spatial mask
+        )
+        
+        self.num_classes = num_classes
+    
+    def forward(self, vision_tokens: torch.Tensor) -> torch.Tensor:
+        batch_size = vision_tokens.shape[0]
+        pooled = vision_tokens.mean(dim=1)
+        mask_logits = self.mask_head(pooled)
+        mask_logits = mask_logits.reshape(batch_size, self.num_classes, 14, 14)
+        return mask_logits
+
+
+# ============================================================================
+# Main Model
+# ============================================================================
+
+class OculusForConditionalGeneration(PreTrainedModel):
+    """
+    Oculus: Unified Vision-Language Model
+    
+    Features:
+    - Multi-encoder vision (DINOv3 + SigLIP2)
+    - Optional reasoning with thinking traces
+    - Multiple output modes: Text, Point, Box, Polygon
+    - Focus/Zoom tool calling for fine-grained perception
+    
+    Usage:
+        ```python
+        from oculus_unified_model import OculusForConditionalGeneration
+        
+        model = OculusForConditionalGeneration.from_pretrained("OceanirAI/oculus-0.2")
+        
+        # Caption mode
+        output = model.generate(image, mode="text", prompt="Describe this image")
+        
+        # VQA mode
+        output = model.generate(image, mode="text", prompt="What color is the cat?")
+        
+        # With reasoning
+        output = model.generate(image, mode="text", prompt="Count the people", think=True)
+        
+        # Detection mode
+        output = model.generate(image, mode="box", prompt="Find all cars")
+        
+        # Point mode (counting)
+        output = model.generate(image, mode="point", prompt="Count the birds")
+        
+        # Segmentation mode
+        output = model.generate(image, mode="polygon", prompt="Segment the road")
+        ```
+    """
+    
+    config_class = OculusConfig
+    base_model_prefix = "oculus"
+    
+    def __init__(self, config: OculusConfig):
+        super().__init__(config)
+        self.config = config
+        
+        # Vision encoder
+        self.vision_encoder = OculusVisionEncoder(config)
+        
+        # Vision adapter (handles dimension mismatch if needed)
+        self.vision_adapter = None
+        self._actual_vision_dim = None
+        
+        # Projector
+        self.projector = OculusProjector(config)
+        
+        # Task-specific heads
+        self.detection_head = OculusDetectionHead(config)
+        self.point_head = OculusPointHead(config)
+        self.segmentation_head = OculusSegmentationHead(config)
+        
+        # Language model (loaded lazily)
+        self.lm_tokenizer = None
+        self.lm_model = None
+        self._lm_loaded = False
+        
+        # Special tokens for reasoning
+        self.thinking_token = config.thinking_token
+        self.thinking_end_token = config.thinking_end_token
+        self.focus_token = config.focus_token
+        self.focus_end_token = config.focus_end_token
+    
+    def load_language_model(self, device: str = "cpu"):
+        """Load language model for text generation."""
+        if self._lm_loaded:
+            return
+        
+        print("[Oculus] Loading language model...")
+        
+        try:
+            # Try BLIP for now (works well for captioning/VQA)
+            from transformers import BlipProcessor, BlipForConditionalGeneration, BlipForQuestionAnswering
+            
+            self.lm_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+            self.lm_caption_model = BlipForConditionalGeneration.from_pretrained(
+                "Salesforce/blip-image-captioning-base"
+            ).to(device)
+            
+            self.lm_vqa_processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
+            self.lm_vqa_model = BlipForQuestionAnswering.from_pretrained(
+                "Salesforce/blip-vqa-base"
+            ).to(device)
+            
+            print("  ✓ BLIP (captioning + VQA)")
+            self._lm_loaded = True
+            
+        except Exception as e:
+            warnings.warn(f"Failed to load language model: {e}")
+    
+    def encode_image(self, image: Union[Image.Image, str, np.ndarray]) -> torch.Tensor:
+        """
+        Encode image to vision tokens.
+        
+        Args:
+            image: PIL Image, file path, or numpy array
+        
+        Returns:
+            Vision tokens [1, num_tokens, embed_dim]
+        """
+        # Load image if path
+        if isinstance(image, str):
+            image = Image.open(image)
+        
+        # Encode with vision encoders
+        vision_features = self.vision_encoder(image)
+        
+        # Check if we need an adapter for dimension mismatch
+        actual_dim = vision_features.shape[-1]
+        expected_dim = self.config.fused_vision_dim
+        
+        if actual_dim != expected_dim:
+            if self.vision_adapter is None or self._actual_vision_dim != actual_dim:
+                # Create adapter layer
+                print(f"  [Adapter] Creating vision adapter: {actual_dim} -> {expected_dim}")
+                self.vision_adapter = nn.Linear(actual_dim, expected_dim)
+                self._actual_vision_dim = actual_dim
+                # Initialize with small weights
+                nn.init.xavier_uniform_(self.vision_adapter.weight)
+                nn.init.zeros_(self.vision_adapter.bias)
+            
+            vision_features = self.vision_adapter(vision_features)
+        
+        # Project to language space
+        vision_tokens = self.projector(vision_features)
+        
+        return vision_tokens
+    
+    def _generate_thinking_trace(
+        self,
+        image: Image.Image,
+        prompt: str,
+        max_tokens: int = 256
+    ) -> str:
+        """
+        Generate a thinking/reasoning trace before answering.
+        
+        This enables multi-step reasoning for complex tasks.
+        """
+        thinking_prompt = f"""Let me think about this step by step:
+1. First, I'll analyze what I see in the image.
+2. Then, I'll consider the question: "{prompt}"
+3. Finally, I'll formulate my answer.
+
+Observation: """
+        
+        # Generate reasoning (simplified for now)
+        if self._lm_loaded and hasattr(self, 'lm_caption_model'):
+            inputs = self.lm_processor(image, thinking_prompt, return_tensors="pt")
+            inputs = {k: v.to(self.lm_caption_model.device) for k, v in inputs.items()}
+            
+            with torch.no_grad():
+                out = self.lm_caption_model.generate(
+                    **inputs,
+                    max_new_tokens=max_tokens,
+                    do_sample=True,
+                    temperature=0.7
+                )
+            thinking = self.lm_processor.decode(out[0], skip_special_tokens=True)
+        else:
+            thinking = "I observe the image and analyze its contents."
+        
+        return thinking
+    
+    def _detect_focus_regions(
+        self,
+        image: Image.Image,
+        prompt: str
+    ) -> List[Tuple[int, int, int, int]]:
+        """
+        Detect regions that need closer inspection (Focus/Zoom system).
+        
+        Returns list of (x1, y1, x2, y2) crop regions.
+        """
+        # Simplified: return full image as single region
+        # In full implementation, would use attention maps to find regions of interest
+        w, h = image.size
+        return [(0, 0, w, h)]
+    
+    def generate(
+        self,
+        image: Union[Image.Image, str, np.ndarray],
+        prompt: str = "Describe this image",
+        mode: str = "text",
+        think: bool = False,
+        focus: bool = False,
+        max_new_tokens: Optional[int] = None,
+        temperature: float = 0.7,
+        return_thinking: bool = True,
+        **kwargs
+    ) -> Union[OculusTextOutput, OculusPointOutput, OculusBoxOutput, OculusPolygonOutput]:
+        """
+        Generate output from image.
+        
+        Args:
+            image: Input image (PIL, path, or array)
+            prompt: Text prompt/question
+            mode: Output mode ("text", "point", "box", "polygon")
+            think: Enable reasoning traces
+            focus: Enable zoom/crop for fine-grained perception
+            max_new_tokens: Maximum tokens to generate
+            temperature: Sampling temperature
+            return_thinking: Include thinking trace in output
+            
+        Returns:
+            Mode-specific output dataclass
+        """
+        # Load models if needed
+        self.vision_encoder.load_encoders()
+        if mode == "text":
+            self.load_language_model()
+        
+        # Load image
+        if isinstance(image, str):
+            image = Image.open(image).convert('RGB')
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image).convert('RGB')
+        
+        # Encode image
+        vision_tokens = self.encode_image(image)
+        
+        # Generate thinking trace if enabled
+        thinking_trace = None
+        if think and self.config.reasoning_enabled:
+            thinking_trace = self._generate_thinking_trace(image, prompt)
+        
+        # Focus system: zoom/crop if needed
+        if focus and self.config.enable_focus:
+            focus_regions = self._detect_focus_regions(image, prompt)
+            # Could re-encode cropped regions here
+        
+        # Mode-specific generation
+        if mode == "text":
+            return self._generate_text(image, prompt, vision_tokens, thinking_trace, max_new_tokens, **kwargs)
+        elif mode == "point":
+            return self._generate_points(vision_tokens, thinking_trace, **kwargs)
+        elif mode == "box":
+            return self._generate_boxes(vision_tokens, thinking_trace, **kwargs)
+        elif mode == "polygon":
+            return self._generate_polygons(vision_tokens, thinking_trace, **kwargs)
+        else:
+            raise ValueError(f"Unknown mode: {mode}")
+    
+    def _generate_text(
+        self,
+        image: Image.Image,
+        prompt: str,
+        vision_tokens: torch.Tensor,
+        thinking_trace: Optional[str],
+        max_new_tokens: Optional[int],
+        **kwargs
+    ) -> OculusTextOutput:
+        """Generate text output (caption or VQA)."""
+        
+        device = vision_tokens.device if vision_tokens.is_cuda else "cpu"
+        max_tokens = max_new_tokens or self.config.max_new_tokens
+        
+        # Determine if this is a question
+        is_question = any(q in prompt.lower() for q in ["what", "where", "who", "how", "why", "is", "are", "does", "do", "can", "?"])
+        
+        if is_question and hasattr(self, 'lm_vqa_model'):
+            # VQA mode
+            inputs = self.lm_vqa_processor(image, prompt, return_tensors="pt")
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+            
+            with torch.no_grad():
+                out = self.lm_vqa_model.generate(**inputs, max_new_tokens=50)
+            text = self.lm_vqa_processor.decode(out[0], skip_special_tokens=True)
+        else:
+            # Caption mode
+            inputs = self.lm_processor(image, prompt, return_tensors="pt")
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+            
+            with torch.no_grad():
+                out = self.lm_caption_model.generate(**inputs, max_new_tokens=max_tokens)
+            text = self.lm_processor.decode(out[0], skip_special_tokens=True)
+        
+        return OculusTextOutput(
+            text=text,
+            thinking_trace=thinking_trace,
+            vision_tokens=vision_tokens
+        )
+    
+    def _generate_points(
+        self,
+        vision_tokens: torch.Tensor,
+        thinking_trace: Optional[str],
+        threshold: float = 0.5,
+        **kwargs
+    ) -> OculusPointOutput:
+        """Generate point detections."""
+        
+        points, cls_logits, confidence = self.point_head(vision_tokens)
+        
+        # Filter by confidence
+        mask = confidence.squeeze(-1) > threshold
+        
+        filtered_points = []
+        filtered_labels = []
+        filtered_conf = []
+        
+        for i in range(vision_tokens.shape[0]):
+            token_mask = mask[i]
+            pts = points[i][token_mask].detach().cpu().numpy().tolist()
+            confs = confidence[i][token_mask].squeeze(-1).detach().cpu().numpy().tolist()
+            cls_ids = cls_logits[i][token_mask].argmax(dim=-1).detach().cpu().numpy().tolist()
+            
+            filtered_points.extend([tuple(p) for p in pts])
+            filtered_conf.extend(confs)
+            filtered_labels.extend([str(c) for c in cls_ids])
+        
+        return OculusPointOutput(
+            points=filtered_points,
+            labels=filtered_labels,
+            confidences=filtered_conf,
+            thinking_trace=thinking_trace,
+            vision_tokens=vision_tokens
+        )
+    
+    def _generate_boxes(
+        self,
+        vision_tokens: torch.Tensor,
+        thinking_trace: Optional[str],
+        threshold: float = 0.3,
+        **kwargs
+    ) -> OculusBoxOutput:
+        """Generate bounding box detections."""
+        
+        cls_logits, box_coords = self.detection_head(vision_tokens)
+        
+        # Get confidence from class logits
+        confidence = F.softmax(cls_logits, dim=-1).max(dim=-1).values
+        
+        filtered_boxes = []
+        filtered_labels = []
+        filtered_conf = []
+        
+        for i in range(vision_tokens.shape[0]):
+            mask = confidence[i] > threshold
+            boxes = box_coords[i][mask].detach().cpu().numpy()
+            confs = confidence[i][mask].detach().cpu().numpy().tolist()
+            cls_ids = cls_logits[i][mask].argmax(dim=-1).detach().cpu().numpy().tolist()
+            
+            filtered_boxes.extend([tuple(b) for b in boxes])
+            filtered_conf.extend(confs)
+            filtered_labels.extend([str(c) for c in cls_ids])
+        
+        return OculusBoxOutput(
+            boxes=filtered_boxes,
+            labels=filtered_labels,
+            confidences=filtered_conf,
+            thinking_trace=thinking_trace,
+            vision_tokens=vision_tokens
+        )
+    
+    def _generate_polygons(
+        self,
+        vision_tokens: torch.Tensor,
+        thinking_trace: Optional[str],
+        **kwargs
+    ) -> OculusPolygonOutput:
+        """Generate polygon/mask segmentation."""
+        
+        mask_logits = self.segmentation_head(vision_tokens)
+        
+        # Get predicted mask
+        mask = mask_logits.argmax(dim=1).detach().cpu().numpy()
+        
+        # Convert to polygons (simplified)
+        # In full implementation, would use cv2.findContours
+        polygons = []
+        labels = []
+        
+        unique_classes = np.unique(mask[0])
+        for cls_id in unique_classes:
+            if cls_id == 0:  # Skip background
+                continue
+            labels.append(str(cls_id))
+            # Placeholder polygon
+            polygons.append([(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 1.0)])
+        
+        return OculusPolygonOutput(
+            polygons=polygons,
+            labels=labels,
+            mask=mask[0],
+            thinking_trace=thinking_trace,
+            vision_tokens=vision_tokens
+        )
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        """
+        Load model from pretrained weights.
+        
+        Args:
+            pretrained_model_name_or_path: HuggingFace repo ID or local path
+        """
+        path = Path(pretrained_model_name_or_path)
+        
+        # Load config
+        config_path = path / "config.json"
+        if config_path.exists():
+            import json
+            with open(config_path) as f:
+                proj_config = json.load(f)
+            
+            # Create config with correct dimensions from projector
+            config = OculusConfig(
+                dinov3_hidden_size=proj_config.get("fused_dim", 2048) - 768,  # Infer from fused
+                siglip_hidden_size=768,
+                projector_hidden_dim=proj_config.get("hidden_dim", 2048),
+                num_vision_tokens=proj_config.get("num_tokens", 64),
+                lm_hidden_size=proj_config.get("embed_dim", 1536),
+            )
+        else:
+            config = OculusConfig()
+        
+        # Create model
+        model = cls(config)
+        
+        # Load projector weights
+        projector_path = path / "projector.npz"
+        if projector_path.exists():
+            model.projector = OculusProjector.from_pretrained(path, config)
+        
+        # Load detection/segmentation heads if available
+        heads_path = path / "heads.pth"
+        if heads_path.exists():
+            heads_state = torch.load(heads_path, map_location="cpu")
+            model.detection_head.load_state_dict(heads_state.get("detection", {}), strict=False)
+            model.point_head.load_state_dict(heads_state.get("point", {}), strict=False)
+            model.segmentation_head.load_state_dict(heads_state.get("segmentation", {}), strict=False)
+        
+        return model
+    
+    def save_pretrained(self, save_directory: str):
+        """Save model to directory."""
+        path = Path(save_directory)
+        path.mkdir(parents=True, exist_ok=True)
+        
+        # Save config
+        self.config.save_pretrained(path)
+        
+        # Save projector
+        projector_state = self.projector.state_dict()
+        # Convert to numpy for MLX compatibility
+        np_weights = {}
+        for k, v in projector_state.items():
+            parts = k.split(".")
+            layer = parts[0]
+            param = ".".join(parts[1:])
+            if layer not in np_weights:
+                np_weights[layer] = {}
+            np_weights[layer][param] = v.cpu().numpy()
+        np.savez(path / "projector.npz", **{k: v for k, v in np_weights.items()})
+        
+        # Save heads
+        torch.save({
+            "detection": self.detection_head.state_dict(),
+            "point": self.point_head.state_dict(),
+            "segmentation": self.segmentation_head.state_dict(),
+        }, path / "heads.pth")
+        
+        print(f"✓ Saved model to {path}")
+
+
+# Register for auto-loading
+OculusForConditionalGeneration.register_for_auto_class("AutoModelForVision2Seq")