Update app.py

app.py

@@ -27,57 +27,60 @@ import base64
 import io
 
 # --------------------------------------------------------------------------------------
-# PATCHED MODEL LOADING
+# FIXED PATCHING FOR C-RADIOv3-B
 # --------------------------------------------------------------------------------------
 
-def patch_transformers_for_radio():
-    """Patch transformers to handle missing ls1 parameters in C-RADIOv3-B"""
+def patch_dinov2_architecture():
+    """Patch the DINOv2 architecture directly to handle missing ls1 parameters"""
     try:
-        import transformers.modeling_utils
+        # Try to import and patch the dinov2_arch module if it exists
+        import sys
+        from huggingface_hub import hf_hub_download
         
-        # Store original function
-        if not hasattr(transformers.modeling_utils, '_original_load_state_dict'):
-            transformers.modeling_utils._original_load_state_dict = transformers.modeling_utils._load_state_dict_into_meta_model
+        # Download the dinov2_arch.py file
+        dinov2_path = hf_hub_download(
+            repo_id="nvidia/C-RADIOv3-B",
+            filename="dinov2_arch.py",
+            cache_dir=".cache"
+        )
         
-        def patched_load_state_dict_into_meta_model(model, state_dict, device_map=None, 
-                                                   offload_folder=None, dtype=None, 
-                                                   offload_state_dict=None, 
-                                                   offload_buffers=None,
-                                                   keep_in_fp32_modules=None,
-                                                   tied_params=None,
-                                                   **kwargs):
-            """Patched loader that ignores missing ls1 keys"""
-            
-            # Filter out any existing ls1 fake keys if they exist
-            filtered_state = {k: v for k, v in state_dict.items() 
-                            if not ('ls1.gamma' in k or 'ls1.grandma' in k)}
-            
-            # Try loading with the original function
-            try:
-                return transformers.modeling_utils._original_load_state_dict(
-                    model, filtered_state, device_map, offload_folder, dtype, 
-                    offload_state_dict, offload_buffers, keep_in_fp32_modules,
-                    tied_params, **kwargs
-                )
-            except KeyError as e:
-                if "ls1.gamma" in str(e) or "ls1.grandma" in str(e):
-                    print(f"⚠️ Ignoring missing layer scaling parameters: {e}")
-                    # Return empty dicts to indicate successful loading
-                    return {}, {}
-                raise
+        # Read the file
+        with open(dinov2_path, 'r') as f:
+            dinov2_code = f.read()
         
-        # Apply the patch
-        transformers.modeling_utils._load_state_dict_into_meta_model = patched_load_state_dict_into_meta_model
-        print("✅ Applied compatibility patch for C-RADIOv3-B")
+        # Replace the error-raising code with a warning
+        dinov2_code = dinov2_code.replace(
+            'raise KeyError(f"Couldn\'t find the key {key_a} nor {key_b} in the state dict!")',
+            '''
+# Patched: Use default values instead of raising error
+import torch.nn as nn
+if not hasattr(self, 'ls1'):
+    self.ls1 = nn.Identity()  # Use identity as fallback
+print(f"Warning: Missing keys {key_a} and {key_b}, using Identity layer as fallback")
+return
+'''
+        )
+        
+        # Save patched version
+        patched_path = ".cache/dinov2_arch_patched.py"
+        os.makedirs(".cache", exist_ok=True)
+        with open(patched_path, 'w') as f:
+            f.write(dinov2_code)
+        
+        # Import the patched version
+        import importlib.util
+        spec = importlib.util.spec_from_file_location("dinov2_arch_patched", patched_path)
+        patched_module = importlib.util.module_from_spec(spec)
+        sys.modules['dinov2_arch'] = patched_module
+        spec.loader.exec_module(patched_module)
+        
+        print("✅ Applied architecture patch for DINOv2")
         return True
         
     except Exception as e:
-        print(f"⚠️ Could not apply patch: {e}")
+        print(f"⚠️ Could not patch DINOv2 architecture: {e}")
         return False
 
-# Apply the patch at module load time
-patch_transformers_for_radio()
-
 # --------------------------------------------------------------------------------------
 # Check Detectron2
 # --------------------------------------------------------------------------------------
@@ -128,92 +131,92 @@ else:
 
 print(f"🖥️ Using device: {DEVICE}")
 
-# Global variables for C model
+# Global variables for model
 image_processor = None
 model = None
 ai_detection_classifier = None
 _preloaded = False
+_use_clip_fallback = False
 
 # --------------------------------------------------------------------------------------
-# FIXED Model Loading
+# SIMPLIFIED Model Loading - Direct CLIP fallback
 # --------------------------------------------------------------------------------------
 
 def preload_models():
-    """Preload models with compatibility fixes"""
-    global image_processor, model, _preloaded
+    """Preload models - try RADIO first, fall back to CLIP"""
+    global image_processor, model, _preloaded, _use_clip_fallback
     
     if _preloaded:
         print("✅ Models already loaded")
         return True
     
-    print("🔄 Preloading C-RADIOv3-B model...")
+    print("🔄 Loading visual encoder model...")
     
-    try:
-        hf_repo = os.getenv('MODEL_REPO', 'nvidia/C-RADIOv3-B')
-        
-        if hf_repo == 'fallback':
-            hf_repo = 'nvidia/C-RADIOv3-B'
-        
-        print(f"📦 Loading from: {hf_repo}")
-        
-        # Method 1: Try with patched loader
+    # Try to load C-RADIOv3-B first
+    hf_repo = os.getenv('MODEL_REPO', 'nvidia/C-RADIOv3-B')
+    
+    if hf_repo != 'fallback':
         try:
-            # Ensure patch is applied
-            patch_transformers_for_radio()
+            print(f"📦 Attempting to load: {hf_repo}")
             
-            # Load image processor
-            from transformers import CLIPImageProcessor, AutoImageProcessor
-            try:
-                image_processor = CLIPImageProcessor.from_pretrained(hf_repo)
-            except:
-                image_processor = AutoImageProcessor.from_pretrained(hf_repo)
+            # Try patching first
+            patch_dinov2_architecture()
             
-            # Suppress the specific warning we know about
+            # Try loading with various workarounds
             with warnings.catch_warnings():
-                warnings.filterwarnings("ignore", message="Couldn't find the key")
-                
-                # Load model with low_cpu_mem_usage=False to avoid meta model issues
-                model = AutoModel.from_pretrained(
-                    hf_repo,
-                    trust_remote_code=True,
-                    low_cpu_mem_usage=False,  # Important: disable meta model loading
-                    ignore_mismatched_sizes=True
-                )
-            
-            model = model.to(DEVICE)
-            model.eval()
-            
-            print("✅ C-RADIOv3-B model loaded successfully with compatibility fixes!")
-            _preloaded = True
-            return True
-            
-        except Exception as e1:
-            print(f"⚠️ Method 1 failed: {e1}")
-            
-            # Method 2: Try loading without trust_remote_code
-            try:
-                print("Trying alternative loading method...")
-                
-                # Use a simpler CLIP model as fallback
-                from transformers import CLIPModel, CLIPProcessor
-                
-                fallback_model = "openai/clip-vit-base-patch32"
-                print(f"Loading fallback model: {fallback_model}")
-                
-                image_processor = CLIPProcessor.from_pretrained(fallback_model)
-                model = CLIPModel.from_pretrained(fallback_model)
-                model = model.to(DEVICE)
-                model.eval()
-                
-                print("✅ Loaded fallback CLIP model successfully!")
-                _preloaded = True
-                return True
-                
-            except Exception as e2:
-                print(f"⚠️ Method 2 failed: {e2}")
+                warnings.filterwarnings("ignore")
                 
+                try:
+                    # Method 1: Load without meta model
+                    from transformers import AutoModel, CLIPImageProcessor
+                    
+                    image_processor = CLIPImageProcessor.from_pretrained(hf_repo)
+                    
+                    # Load with specific settings to avoid issues
+                    model = AutoModel.from_pretrained(
+                        hf_repo,
+                        trust_remote_code=True,
+                        low_cpu_mem_usage=False,
+                        torch_dtype=torch.float32
+                    )
+                    
+                    model = model.to(DEVICE)
+                    model.eval()
+                    
+                    print(f"✅ Successfully loaded {hf_repo}")
+                    _preloaded = True
+                    _use_clip_fallback = False
+                    return True
+                    
+                except KeyError as ke:
+                    if "ls1.gamma" in str(ke) or "ls1.grandma" in str(ke):
+                        print(f"⚠️ Known C-RADIOv3-B issue: {ke}")
+                    else:
+                        print(f"⚠️ Unexpected error: {ke}")
+                except Exception as e:
+                    print(f"⚠️ Could not load {hf_repo}: {e}")
+                    
+        except Exception as e:
+            print(f"⚠️ Error during RADIO loading: {e}")
+    
+    # Fall back to CLIP model which we know works
+    try:
+        print("📦 Loading fallback CLIP model...")
+        from transformers import CLIPModel, CLIPProcessor
+        
+        clip_model = "openai/clip-vit-base-patch32"
+        image_processor = CLIPProcessor.from_pretrained(clip_model)
+        model = CLIPModel.from_pretrained(clip_model)
+        model = model.to(DEVICE)
+        model.eval()
+        
+        print(f"✅ Successfully loaded fallback {clip_model}")
+        _preloaded = True
+        _use_clip_fallback = True
+        return True
+        
     except Exception as e:
-        print(f"❌ Could not preload model: {e}")
+        print(f"❌ Could not load any model: {e}")
         traceback.print_exc()
     
     return False
@@ -317,7 +320,6 @@ def run_damage_detection(pil_image: Image.Image, score_thresh: float = 0.5):
         
     except Exception as e:
         print(f"⚠️ Stage 1 error: {e}")
-        traceback.print_exc()
         # Fallback to simulator
         rgb = np.array(pil_image.convert("RGB"))
         boxes = simulate_damage_detection(rgb, seed_from=rgb)
@@ -375,13 +377,13 @@ def preprocess_image(image) -> Optional[Image.Image]:
                 else:
                     image = np.clip(image, 0, 255).astype(np.uint8)
             
-            pil = Image.fromarray(image, 'RGB')
+            pil = Image.fromarray(image)
         else:
             # Try to convert whatever it is
             arr = np.array(image)
             if arr.dtype != np.uint8:
                 arr = np.clip(arr, 0, 255).astype(np.uint8)
-            pil = Image.fromarray(arr, 'RGB')
+            pil = Image.fromarray(arr)
 
         # Handle EXIF orientation
         pil = ImageOps.exif_transpose(pil)
@@ -393,8 +395,8 @@ def preprocess_image(image) -> Optional[Image.Image]:
         return None
 
 def extract_features(image, return_stats=False):
-    """Extract features with proper handling for different model types."""
-    global image_processor, model
+    """Extract features - handles both CLIP and RADIO models."""
+    global image_processor, model, _use_clip_fallback
     
     if image_processor is None or model is None:
         raise Exception("Model not initialized")
@@ -410,42 +412,40 @@ def extract_features(image, return_stats=False):
     # Process image
     inputs = image_processor(images=image, return_tensors='pt', do_resize=True)
     
-    # Handle different processor outputs
+    # Get the right input tensor
     if hasattr(inputs, 'pixel_values'):
         pixel_values = inputs.pixel_values.to(DEVICE)
     else:
-        pixel_values = inputs['input_ids'].to(DEVICE) if 'input_ids' in inputs else inputs.to(DEVICE)
+        pixel_values = inputs['pixel_values'].to(DEVICE)
     
-    # Get features
+    # Extract features based on model type
     with torch.no_grad():
-        outputs = model(pixel_values)
-    
-    # Handle different model outputs
-    if hasattr(model, 'get_image_features'):
-        # CLIP model
-        features = model.get_image_features(pixel_values)
-    elif isinstance(outputs, dict):
-        # Dictionary output
-        if 'features' in outputs:
-            features = outputs['features']
-        elif 'last_hidden_state' in outputs:
-            features = outputs['last_hidden_state']
-        elif 'pooler_output' in outputs:
-            features = outputs['pooler_output']
+        if _use_clip_fallback and hasattr(model, 'get_image_features'):
+            # CLIP model
+            features = model.get_image_features(pixel_values)
         else:
-            # Take the first tensor value
-            features = next(iter(outputs.values()))
-    elif isinstance(outputs, (list, tuple)):
-        # Tuple/list output - take last element
-        features = outputs[-1] if len(outputs) > 1 else outputs[0]
-    else:
-        # Direct tensor output
-        features = outputs
+            # RADIO or other model
+            outputs = model(pixel_values)
+            
+            # Handle different output formats
+            if isinstance(outputs, dict):
+                if 'features' in outputs:
+                    features = outputs['features']
+                elif 'last_hidden_state' in outputs:
+                    features = outputs['last_hidden_state']
+                elif 'pooler_output' in outputs:
+                    features = outputs['pooler_output']
+                else:
+                    features = list(outputs.values())[0]
+            elif isinstance(outputs, (tuple, list)):
+                features = outputs[-1] if len(outputs) > 1 else outputs[0]
+            else:
+                features = outputs
     
     # Pool if needed
     if features.ndim == 3:  # (B, T, C)
         features = features.mean(dim=1)
-    elif features.ndim == 4:  # (B, C, H, W)
+    elif features.ndim == 4:  # (B, C, H, W)  
         features = features.mean(dim=(2, 3))
     
     # Normalize and flatten
@@ -458,7 +458,8 @@ def extract_features(image, return_stats=False):
             "std": float(features.std()),
             "min": float(features.min()),
             "max": float(features.max()),
-            "shape": features.shape
+            "shape": features.shape,
+            "model_type": "CLIP" if _use_clip_fallback else "RADIO"
         }
         return features, stats
     
@@ -648,10 +649,14 @@ def create_gradio_interface():
                     print(f"⚠️ Stage 1 error: {e}")
 
             # Status display
-            if isinstance(detailed_result, dict) and detailed_result.get("is_demo"):
-                status_html = '<div style="padding: 10px; background: #fef3c7; border-radius: 8px;"><p style="margin: 0; color: #f59e0b;">⚠️ Running in Demo Mode</p></div>'
+            if isinstance(detailed_result, dict):
+                if detailed_result.get("is_demo"):
+                    status_html = '<div style="padding: 10px; background: #fef3c7; border-radius: 8px;"><p style="margin: 0; color: #f59e0b;">⚠️ Running in Demo Mode (using fallback model)</p></div>'
+                else:
+                    model_info = detailed_result.get('feature_stats', {}).get('model_type', 'Unknown')
+                    status_html = f'<div style="padding: 10px; background: #d1fae5; border-radius: 8px;"><p style="margin: 0; color: #10b981;">✅ Analysis Complete (using {model_info} model)</p></div>'
             else:
-                status_html = '<div style="padding: 10px; background: #d1fae5; border-radius: 8px;"><p style="margin: 0; color: #10b981;">✅ Analysis Complete</p></div>'
+                status_html = '<div style="padding: 10px; background: #fee2e2; border-radius: 8px;"><p style="margin: 0; color: #dc2626;">❌ Analysis Failed</p></div>'
 
             return simple_result, detailed_result, status_html, dmg_results, annotated
 
@@ -688,12 +693,17 @@ def create_gradio_interface():
         with gr.Accordion("ℹ️ About", open=False):
             gr.Markdown("""
             ### Pipeline
-            - **Stage 1**: Detectron2 damage detection (optional)
-            - **Stage 2**: Visual features + AI detection classifier
+            - **Stage 1**: Detectron2 damage detection (simulated if not available)
+            - **Stage 2**: Visual feature extraction + AI detection classifier
+            
+            ### Models
+            - **Primary**: C-RADIOv3-B visual encoder (if available)
+            - **Fallback**: CLIP-ViT-B-32 (reliable alternative)
+            - **Classifier**: Scikit-learn model for AI detection
             
-            ### Notes
-            - Falls back to demo mode if models are unavailable
-            - C-RADIOv3-B model includes compatibility fixes for layer scaling issues
+            ### Status
+            - The app will show which model is being used in the status display
+            - Falls back gracefully if primary models are unavailable
             """)
 
     return app
@@ -712,14 +722,15 @@ if __name__ == "__main__":
     
     # Preload models with fixes
     if preload_models():
-        print("✅ Models preloaded successfully")
+        model_type = "CLIP" if _use_clip_fallback else "RADIO"
+        print(f"✅ Visual encoder loaded ({model_type})")
     else:
-        print("⚠️ Running in demo mode")
+        print("⚠️ Running in full demo mode")
     
     # Load classifier
     model_path = huggingface_model_path or DEFAULT_AI_DETECTION_MODEL_PATH
     if load_ai_detection_classifier(model_path):
-        print("✅ Classifier loaded")
+        print("✅ AI detection classifier loaded")
     
     print("=" * 60)