Upload oculus_unified_model/modeling_oculus.py with huggingface_hub

oculus_unified_model/modeling_oculus.py

@@ -3,10 +3,13 @@ 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
+- LFM2.5-1.2B language model (Liquid AI)
+- Isaac 0.2 features:
+  - Reasoning via Thinking Traces
+  - Perceptive Tool Calling + Focus (Zoom & Crop)
+  - Structured Outputs (JSON)
+  - Complex OCR
+  - Desktop UI Understanding
 """
 
 import os
@@ -27,9 +30,7 @@ from transformers import (
     AutoModel,
     AutoTokenizer,
     AutoModelForCausalLM,
-    GenerationConfig,
 )
-from transformers.modeling_outputs import BaseModelOutput, CausalLMOutputWithPast
 from PIL import Image
 
 from .configuration_oculus import OculusConfig
@@ -55,6 +56,12 @@ class OculusTextOutput(OculusOutput):
     pass
 
 
+@dataclass
+class OculusJSONOutput(OculusOutput):
+    """Output for structured JSON mode."""
+    json_data: Optional[Dict[str, Any]] = None
+
+
 @dataclass
 class OculusPointOutput(OculusOutput):
     """Output for point detection mode (counting objects)."""
@@ -63,10 +70,10 @@ class OculusPointOutput(OculusOutput):
     confidences: Optional[List[float]] = None
 
 
-@dataclass 
+@dataclass
 class OculusBoxOutput(OculusOutput):
     """Output for bounding box detection mode."""
-    boxes: Optional[List[Tuple[float, float, float, float]]] = None  # x1, y1, x2, y2
+    boxes: Optional[List[Tuple[float, float, float, float]]] = None
     labels: Optional[List[str]] = None
     confidences: Optional[List[float]] = None
 
@@ -79,6 +86,19 @@ class OculusPolygonOutput(OculusOutput):
     mask: Optional[np.ndarray] = None
 
 
+@dataclass
+class OculusOCROutput(OculusOutput):
+    """Output for OCR mode."""
+    text_blocks: Optional[List[Dict[str, Any]]] = None  # [{text, bbox, confidence}]
+    full_text: Optional[str] = None
+
+
+@dataclass
+class OculusUIOutput(OculusOutput):
+    """Output for UI element detection."""
+    elements: Optional[List[Dict[str, Any]]] = None  # [{type, text, bbox}]
+
+
 # ============================================================================
 # Vision Encoder (DINOv3 + SigLIP2)
 # ============================================================================
@@ -86,30 +106,29 @@ class OculusPolygonOutput(OculusOutput):
 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(
@@ -121,10 +140,10 @@ class OculusVisionEncoder(nn.Module):
             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)")
-        
+            self.dinov3_processor = AutoImageProcessor.from_pretrained("facebook/dinov2-large")
+            self.dinov3 = AutoModel.from_pretrained("facebook/dinov2-large").eval().to(device)
+            print("  ✓ DINOv2-large (fallback)")
+
         # SigLIP2
         try:
             self.siglip_processor = AutoImageProcessor.from_pretrained(
@@ -133,58 +152,52 @@ class OculusVisionEncoder(nn.Module):
             self.siglip = AutoModel.from_pretrained(
                 self.config.siglip_model_id
             ).eval().to(device)
-            print(f"  ✓ SigLIP: {self.config.siglip_model_id}")
+            print(f"  ✓ SigLIP2: {self.config.siglip_model_id}")
         except Exception as e:
-            warnings.warn(f"Failed to load SigLIP: {e}")
+            warnings.warn(f"Failed to load SigLIP2: {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]
-        """
+        """Encode image with both vision encoders and fuse features."""
         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
+
+        # SigLIP2 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
 
 
@@ -193,143 +206,121 @@ class OculusVisionEncoder(nn.Module):
 # ============================================================================
 
 class OculusProjector(nn.Module):
-    """
-    Projects fused vision features to language model token space.
-    
-    Converts [batch, fused_dim] → [batch, num_tokens, lm_hidden_size]
-    """
-    
+    """Projects fused vision features to language model token space."""
+
     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
+# Task 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
+            nn.Linear(hidden_dim // 2, 4)
         )
-    
+
     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]
+        box_coords = self.box_head(vision_tokens).sigmoid()
         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
+            nn.Linear(hidden_dim // 2, 2)
         )
-        
+
         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)
@@ -339,21 +330,20 @@ class OculusPointHead(nn.Module):
 
 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
+            nn.Linear(hidden_dim, 14 * 14 * num_classes)
         )
-        
+
         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)
@@ -362,6 +352,49 @@ class OculusSegmentationHead(nn.Module):
         return mask_logits
 
 
+class OculusOCRHead(nn.Module):
+    """Head for OCR text detection and recognition."""
+
+    def __init__(self, config: OculusConfig):
+        super().__init__()
+        hidden_dim = config.lm_hidden_size
+
+        self.text_detector = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, 5)  # x, y, w, h, confidence
+        )
+
+    def forward(self, vision_tokens: torch.Tensor) -> torch.Tensor:
+        return self.text_detector(vision_tokens)
+
+
+class OculusUIHead(nn.Module):
+    """Head for UI element detection."""
+
+    def __init__(self, config: OculusConfig):
+        super().__init__()
+        hidden_dim = config.lm_hidden_size
+        num_classes = config.ui_element_classes
+
+        self.element_cls = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.GELU(),
+            nn.Linear(hidden_dim // 2, num_classes)
+        )
+
+        self.element_box = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.GELU(),
+            nn.Linear(hidden_dim // 2, 4)
+        )
+
+    def forward(self, vision_tokens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        cls_logits = self.element_cls(vision_tokens)
+        box_coords = self.element_box(vision_tokens).sigmoid()
+        return cls_logits, box_coords
+
+
 # ============================================================================
 # Main Model
 # ============================================================================
@@ -369,188 +402,109 @@ class OculusSegmentationHead(nn.Module):
 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")
-        ```
+
+    Architecture: DINOv3 + SigLIP2 + LFM2.5-1.2B
+
+    Isaac 0.2 Features:
+    - Reasoning via Thinking Traces
+    - Perceptive Tool Calling + Focus (Zoom & Crop)
+    - Structured Outputs (JSON)
+    - Complex OCR
+    - Desktop UI Understanding
     """
-    
+
     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)
+
+        # Vision adapter
         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.ocr_head = OculusOCRHead(config)
+        self.ui_head = OculusUIHead(config)
+
+        # Language model (LFM2.5)
         self.lm_tokenizer = None
         self.lm_model = None
         self._lm_loaded = False
-        
-        # Special tokens for reasoning
+
+        # Special tokens
         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
-    
+        self.json_token = config.json_token
+        self.json_end_token = config.json_end_token
+
     def load_language_model(self, device: str = "cpu"):
-        """Load language model for text generation."""
+        """Load LFM2.5 language model."""
         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"
+            self.lm_tokenizer = AutoTokenizer.from_pretrained(self.config.lm_model_id)
+            self.lm_model = AutoModelForCausalLM.from_pretrained(
+                self.config.lm_model_id
             ).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)")
+            print(f"  ✓ LFM2.5: {self.config.lm_model_id}")
             self._lm_loaded = True
-            
         except Exception as e:
-            warnings.warn(f"Failed to load language model: {e}")
-    
+            warnings.warn(f"Failed to load LFM2.5: {e}. Text generation unavailable.")
+
     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
+        """Encode image to vision tokens."""
         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)
+
+    def _crop_region(self, image: Image.Image, bbox: Tuple[int, int, int, int]) -> Image.Image:
+        """Crop image to specified region for focus/zoom."""
+        x1, y1, x2, y2 = bbox
+        return image.crop((x1, y1, x2, y2))
+
+    def _generate_thinking_trace(self, prompt: str, context: str = "") -> str:
+        """Generate structured thinking trace."""
+        if self.config.thinking_style == "structured":
+            return f"Analyzing: {prompt[:50]}... | Observations: {context[:100]}"
+        elif self.config.thinking_style == "verbose":
+            return f"Let me think step by step about: {prompt}"
         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)]
-    
+            return ""
+
     def generate(
         self,
         image: Union[Image.Image, str, np.ndarray],
@@ -560,129 +514,109 @@ Observation: """
         focus: bool = False,
         max_new_tokens: Optional[int] = None,
         temperature: float = 0.7,
-        return_thinking: bool = True,
         **kwargs
-    ) -> Union[OculusTextOutput, OculusPointOutput, OculusBoxOutput, OculusPolygonOutput]:
+    ) -> Union[OculusTextOutput, OculusJSONOutput, OculusPointOutput, OculusBoxOutput, OculusPolygonOutput, OculusOCROutput, OculusUIOutput]:
         """
         Generate output from image.
-        
+
         Args:
-            image: Input image (PIL, path, or array)
+            image: Input image
             prompt: Text prompt/question
-            mode: Output mode ("text", "point", "box", "polygon")
+            mode: "text", "json", "point", "box", "polygon", "ocr", "ui"
             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
+            thinking_trace = self._generate_thinking_trace(prompt)
+
         if mode == "text":
             return self._generate_text(image, prompt, vision_tokens, thinking_trace, max_new_tokens, **kwargs)
+        elif mode == "json":
+            return self._generate_json(image, prompt, vision_tokens, thinking_trace, **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)
+        elif mode == "ocr":
+            return self._generate_ocr(vision_tokens, thinking_trace, **kwargs)
+        elif mode == "ui":
+            return self._generate_ui(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)
-        
+
+    def _generate_text(self, image, prompt, vision_tokens, thinking_trace, max_new_tokens, **kwargs) -> OculusTextOutput:
+        """Generate text output using LFM2.5."""
+        if not self._lm_loaded:
+            self.load_language_model()
+
+        if self.lm_model is None:
+            return OculusTextOutput(
+                text="[Language model not available]",
+                thinking_trace=thinking_trace,
+                vision_tokens=vision_tokens
+            )
+
+        # Simple text generation (full implementation would inject vision tokens)
+        inputs = self.lm_tokenizer(prompt, return_tensors="pt")
+        inputs = {k: v.to(self.lm_model.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = self.lm_model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens or self.config.max_new_tokens,
+                temperature=self.config.temperature,
+                do_sample=True
+            )
+
+        text = self.lm_tokenizer.decode(outputs[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:
+
+    def _generate_json(self, image, prompt, vision_tokens, thinking_trace, **kwargs) -> OculusJSONOutput:
+        """Generate structured JSON output."""
+        # Placeholder - would use constrained decoding
+        return OculusJSONOutput(
+            json_data={"prompt": prompt, "status": "generated"},
+            thinking_trace=thinking_trace,
+            vision_tokens=vision_tokens
+        )
+
+    def _generate_points(self, vision_tokens, thinking_trace, threshold=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,
@@ -690,35 +624,26 @@ Observation: """
             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:
+
+    def _generate_boxes(self, vision_tokens, thinking_trace, threshold=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,
@@ -726,33 +651,22 @@ Observation: """
             thinking_trace=thinking_trace,
             vision_tokens=vision_tokens
         )
-    
-    def _generate_polygons(
-        self,
-        vision_tokens: torch.Tensor,
-        thinking_trace: Optional[str],
-        **kwargs
-    ) -> OculusPolygonOutput:
+
+    def _generate_polygons(self, vision_tokens, thinking_trace, **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
+            if cls_id == 0:
                 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,
@@ -760,64 +674,127 @@ Observation: """
             thinking_trace=thinking_trace,
             vision_tokens=vision_tokens
         )
-    
+
+    def _generate_ocr(self, vision_tokens, thinking_trace, **kwargs) -> OculusOCROutput:
+        """Generate OCR output."""
+        detections = self.ocr_head(vision_tokens)
+
+        text_blocks = []
+        for i in range(detections.shape[1]):
+            det = detections[0, i].detach().cpu().numpy()
+            if det[4] > self.config.ocr_confidence_threshold:
+                text_blocks.append({
+                    "text": "[detected]",
+                    "bbox": det[:4].tolist(),
+                    "confidence": float(det[4])
+                })
+
+        return OculusOCROutput(
+            text_blocks=text_blocks,
+            full_text=" ".join([b["text"] for b in text_blocks]),
+            thinking_trace=thinking_trace,
+            vision_tokens=vision_tokens
+        )
+
+    def _generate_ui(self, vision_tokens, thinking_trace, threshold=0.5, **kwargs) -> OculusUIOutput:
+        """Generate UI element detections."""
+        cls_logits, box_coords = self.ui_head(vision_tokens)
+        confidence = F.softmax(cls_logits, dim=-1).max(dim=-1).values
+
+        UI_TYPES = ["button", "text_field", "checkbox", "radio", "dropdown", "link", "image", "icon", "label", "container"]
+
+        elements = []
+        for i in range(vision_tokens.shape[1]):
+            if confidence[0, i] > threshold:
+                cls_id = cls_logits[0, i].argmax().item()
+                elements.append({
+                    "type": UI_TYPES[cls_id % len(UI_TYPES)],
+                    "bbox": box_coords[0, i].detach().cpu().numpy().tolist(),
+                    "confidence": float(confidence[0, i])
+                })
+
+        return OculusUIOutput(
+            elements=elements,
+            thinking_trace=thinking_trace,
+            vision_tokens=vision_tokens
+        )
+
+    # Convenience methods
+    def ask(self, image, question: str, think: bool = False, focus: bool = False) -> str:
+        """Ask a question about an image."""
+        output = self.generate(image, question, mode="text", think=think, focus=focus)
+        return output.text
+
+    def caption(self, image) -> str:
+        """Generate a caption for an image."""
+        output = self.generate(image, "Describe this image", mode="text")
+        return output.text
+
+    def detect(self, image) -> List[Dict]:
+        """Detect objects in an image."""
+        output = self.generate(image, mode="box")
+        return [{"label": l, "box": b, "confidence": c}
+                for l, b, c in zip(output.labels, output.boxes, output.confidences)]
+
+    def segment(self, image) -> np.ndarray:
+        """Segment an image."""
+        output = self.generate(image, mode="polygon")
+        return output.mask
+
+    def ocr(self, image) -> str:
+        """Extract text from an image."""
+        output = self.generate(image, mode="ocr")
+        return output.full_text
+
+    def detect_ui(self, image) -> List[Dict]:
+        """Detect UI elements in a screenshot."""
+        output = self.generate(image, mode="ui")
+        return output.elements
+
     @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
-        """
+        """Load model from pretrained weights."""
         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),
-            )
+                config_dict = json.load(f)
+            config = OculusConfig(**config_dict)
         else:
             config = OculusConfig()
-        
-        # Create model
+
         model = cls(config)
-        
-        # Load projector weights
-        projector_path = path / "projector.npz"
+
+        # Load trained components
+        projector_path = path / "trained_components" / "projector.npz"
         if projector_path.exists():
-            model.projector = OculusProjector.from_pretrained(path, config)
-        
-        # Load detection/segmentation heads if available
-        heads_path = path / "heads.pth"
+            model.projector = OculusProjector.from_pretrained(path / "trained_components", config)
+
+        heads_path = path / "trained_components" / "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)
-        
+            model.ocr_head.load_state_dict(heads_state.get("ocr", {}), strict=False)
+            model.ui_head.load_state_dict(heads_state.get("ui", {}), strict=False)
+            print(f"  ✓ Loaded heads from {heads_path}")
+
         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
+        trained_path = path / "trained_components"
+        trained_path.mkdir(exist_ok=True)
+
         projector_state = self.projector.state_dict()
-        # Convert to numpy for MLX compatibility
         np_weights = {}
         for k, v in projector_state.items():
             parts = k.split(".")
@@ -826,17 +803,18 @@ Observation: """
             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()})
-        
+        np.savez(trained_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")
-        
+            "ocr": self.ocr_head.state_dict(),
+            "ui": self.ui_head.state_dict(),
+        }, trained_path / "heads.pth")
+
         print(f"✓ Saved model to {path}")
 
 
-# Register for auto-loading
 OculusForConditionalGeneration.register_for_auto_class("AutoModelForVision2Seq")