Update app.py

app.py

@@ -1,268 +1,732 @@
+#!/usr/bin/env python3
 import os
+import sys
+import traceback
+from typing import Optional, Tuple, Dict, Any, List
+import warnings
+
+import importlib.util
+import time
+import cv2
 import torch
+import numpy as np
+import gradio as gr
+from PIL import Image, ImageOps
+from torchvision import transforms
+import torch.nn as nn
+import torch.nn.functional as F
+import traceback
+from torchvision.models import vit_b_16
+from transformers import AutoModel, CLIPImageProcessor, AutoConfig
+import joblib
+import zipfile
 import json
-from huggingface_hub import hf_hub_download
-import safetensors.torch
+from datetime import datetime
+import requests
+import base64
+import io
 
-print("=" * 60)
-print("C-RADIOv3-B Model Deep Inspection")
-print("=" * 60)
+# --------------------------------------------------------------------------------------
+# PATCHED MODEL LOADING
+# --------------------------------------------------------------------------------------
 
-# Step 1: Download and inspect the model file directly
-def inspect_model_weights():
-    """Directly inspect the safetensors file to see what keys exist"""
-    
-    print("\n📥 Downloading model weights for inspection...")
-    
-    # Download the model file
-    model_path = hf_hub_download(
-        repo_id="nvidia/C-RADIOv3-B",
-        filename="model.safetensors"
-    )
-    
-    print(f"Downloaded to: {model_path}")
-    
-    # Load the safetensors file
-    print("\n🔍 Inspecting model weights...")
-    state_dict = safetensors.torch.load_file(model_path)
-    
-    # Analyze the keys
-    all_keys = list(state_dict.keys())
-    print(f"Total keys in model: {len(all_keys)}")
-    
-    # Look for ls1 related keys
-    ls1_keys = [k for k in all_keys if 'ls1' in k.lower()]
-    ls_keys = [k for k in all_keys if 'ls' in k.lower()]
-    gamma_keys = [k for k in all_keys if 'gamma' in k.lower()]
-    block_keys = [k for k in all_keys if k.startswith('blocks.')]
-    
-    print(f"\n📊 Key Analysis:")
-    print(f"  Keys with 'ls1': {len(ls1_keys)}")
-    print(f"  Keys with 'ls': {len(ls_keys)}")
-    print(f"  Keys with 'gamma': {len(gamma_keys)}")
-    print(f"  Keys starting with 'blocks.': {len(block_keys)}")
-    
-    # Show first few block keys
-    print(f"\n📝 First 20 block keys:")
-    for i, key in enumerate(sorted([k for k in all_keys if k.startswith('blocks.0.')])[:20]):
-        print(f"  {key}")
-    
-    # Check what's actually in blocks.0
-    blocks_0_keys = [k for k in all_keys if k.startswith('blocks.0.')]
-    print(f"\n🔎 All blocks.0 submodules:")
-    submodules = set()
-    for key in blocks_0_keys:
-        parts = key.split('.')
-        if len(parts) > 2:
-            submodules.add(parts[2])
-    for submodule in sorted(submodules):
-        count = len([k for k in blocks_0_keys if f'blocks.0.{submodule}.' in k])
-        print(f"  blocks.0.{submodule}.*: {count} parameters")
-    
-    return state_dict, all_keys
+def patch_transformers_for_radio():
+    """Patch transformers to handle missing ls1 parameters in C-RADIOv3-B"""
+    try:
+        import transformers.modeling_utils
+        
+        # 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
+        
+        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
+        
+        # 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")
+        return True
+        
+    except Exception as e:
+        print(f"⚠️ Could not apply patch: {e}")
+        return False
 
-# Step 2: Check the model architecture expectations
-def inspect_model_code():
-    """Download and inspect the model code to understand what it expects"""
-    
-    print("\n📜 Downloading model code...")
-    
-    # Download the dinov2_arch.py file
-    dinov2_path = hf_hub_download(
-        repo_id="nvidia/C-RADIOv3-B",
-        filename="dinov2_arch.py"
-    )
-    
-    print(f"Downloaded dinov2_arch.py to: {dinov2_path}")
-    
-    # Read the problematic part of the code
-    with open(dinov2_path, 'r') as f:
-        lines = f.readlines()
-    
-    # Find the error location (around line 309)
-    print("\n🔍 Code around line 309 (error location):")
-    for i in range(max(0, 308-10), min(len(lines), 308+10)):
-        if i == 308:  # Line 309 (0-indexed)
-            print(f">>> {i+1}: {lines[i].rstrip()}")
-        else:
-            print(f"    {i+1}: {lines[i].rstrip()}")
+# Apply the patch at module load time
+patch_transformers_for_radio()
+
+# --------------------------------------------------------------------------------------
+# Check Detectron2
+# --------------------------------------------------------------------------------------
+
+DETECTRON2_AVAILABLE = False
+try:
+    from detectron2.engine import DefaultPredictor
+    from detectron2.config import get_cfg
+    from detectron2.utils.visualizer import Visualizer, ColorMode
+    from detectron2 import model_zoo
+    DETECTRON2_AVAILABLE = True
+    print("✅ Detectron2 imported successfully")
+except ImportError as e:
+    print(f"⚠️ Detectron2 not available: {e}")
+    DETECTRON2_AVAILABLE = False
+
+# Try to download model from Hugging Face
+huggingface_model_path = None
+try:
+    from huggingface_hub import hf_hub_download
+    
+    repo = os.getenv('PRIVATE_REPO', 'fallback')
+    token = os.getenv('key')
+    
+    if repo != 'fallback' and token:
+        huggingface_model_path = hf_hub_download(
+            repo_id=repo,
+            filename="V1.pkl",
+            token=token
+        )
+        print(f"✅ Model downloaded from Hugging Face: {huggingface_model_path}")
+except Exception as e:
+    print(f"⚠️ Could not download model from Hugging Face: {e}")
+    print("🔄 Will use demo mode with simulated results")
+    huggingface_model_path = None
+
+# --------------------------------------------------------------------------------------
+# Basics
+# --------------------------------------------------------------------------------------
+
+# Initialize device for model
+if torch.backends.mps.is_available():
+    DEVICE = torch.device("mps")
+elif torch.cuda.is_available():
+    DEVICE = torch.device("cuda")
+else:
+    DEVICE = torch.device("cpu")
+
+print(f"🖥️ Using device: {DEVICE}")
+
+# Global variables for C model
+image_processor = None
+model = None
+ai_detection_classifier = None
+_preloaded = False
+
+# --------------------------------------------------------------------------------------
+# FIXED Model Loading
+# --------------------------------------------------------------------------------------
+
+def preload_models():
+    """Preload models with compatibility fixes"""
+    global image_processor, model, _preloaded
     
-    # Look for _load_from_state_dict method
-    print("\n📖 Looking for _load_from_state_dict method...")
-    for i, line in enumerate(lines):
-        if '_load_from_state_dict' in line:
-            print(f"Found at line {i+1}: {line.rstrip()}")
-            # Show context
-            for j in range(max(0, i-2), min(len(lines), i+15)):
-                print(f"  {j+1}: {lines[j].rstrip()}")
-            break
-
-# Step 3: Create a working loader
-def create_fixed_loader():
-    """Create a fixed loading function that handles the missing keys"""
+    if _preloaded:
+        print("✅ Models already loaded")
+        return True
     
-    print("\n🔧 Creating Fixed Model Loader...")
+    print("🔄 Preloading C-RADIOv3-B model...")
     
-    # Create a custom model loading function
-    code = '''
-import torch
-from transformers import AutoModel, AutoConfig
-import warnings
-
-class RADIOModelFixed:
-    @staticmethod
-    def from_pretrained(repo_id="nvidia/C-RADIOv3-B"):
-        """Load RADIO model with compatibility fixes"""
-        
-        print("Loading with compatibility fixes...")
+    try:
+        hf_repo = os.getenv('MODEL_REPO', 'nvidia/C-RADIOv3-B')
         
-        # First, modify the environment to skip the problematic check
-        import sys
-        import transformers.modeling_utils as mu
+        if hf_repo == 'fallback':
+            hf_repo = 'nvidia/C-RADIOv3-B'
         
-        # Store original function
-        original_load = mu._load_state_dict_into_meta_model
+        print(f"📦 Loading from: {hf_repo}")
         
-        def patched_load(model, state_dict, device_map=None, offload_folder=None, 
-                        dtype=None, offload_state_dict=None, tie_weights=True,
-                        **kwargs):
-            """Patched loader that handles missing ls1 keys"""
+        # Method 1: Try with patched loader
+        try:
+            # Ensure patch is applied
+            patch_transformers_for_radio()
             
-            # Create a modified state dict with dummy ls1 keys if needed
-            modified_state = state_dict.copy()
+            # Load image processor
+            from transformers import CLIPImageProcessor, AutoImageProcessor
+            try:
+                image_processor = CLIPImageProcessor.from_pretrained(hf_repo)
+            except:
+                image_processor = AutoImageProcessor.from_pretrained(hf_repo)
             
-            # Check if we need to add dummy ls1 keys
-            block_keys = [k for k in state_dict.keys() if k.startswith('blocks.')]
-            if block_keys and not any('ls1' in k for k in block_keys):
-                print("  Adding compatibility keys for ls1 layers...")
+            # Suppress the specific warning we know about
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", message="Couldn't find the key")
                 
-                # Find all blocks
-                block_indices = set()
-                for key in block_keys:
-                    parts = key.split('.')
-                    if len(parts) > 1 and parts[1].isdigit():
-                        block_indices.add(int(parts[1]))
-                
-                # Add dummy ls1 parameters for each block
-                for idx in block_indices:
-                    # These will be ignored but prevent the error
-                    if f'blocks.{idx}.norm1.weight' in state_dict:
-                        # Use norm1 as a template for shape
-                        template = state_dict[f'blocks.{idx}.norm1.weight']
-                        modified_state[f'blocks.{idx}.ls1.gamma'] = torch.ones_like(template)
-                    else:
-                        # Default to scalar
-                        modified_state[f'blocks.{idx}.ls1.gamma'] = torch.tensor(1.0)
+                # 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}")
             
-            # Call original with modified state
-            return original_load(model, modified_state, device_map, offload_folder,
-                               dtype, offload_state_dict, tie_weights, **kwargs)
+            # 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}")
+                
+    except Exception as e:
+        print(f"❌ Could not preload model: {e}")
+        traceback.print_exc()
+    
+    return False
+
+# --------------------------------------------------------------------------------------
+# Paths
+# --------------------------------------------------------------------------------------
+DEFAULT_AI_DETECTION_MODEL_PATH = "./output/V1.pkl"
+DEFAULT_DAMAGE_MODEL_PATH = "./output/model_final.pth"
+
+# --------------------------------------------------------------------------------------
+# Damage Detection (Stage 1)
+# --------------------------------------------------------------------------------------
+
+_damage_predictor = None
+
+def load_damage_model(model_path: str, device_str: str = None):
+    """Load fine-tuned Detectron2 model once (Stage 1)."""
+    global _damage_predictor
+    if _damage_predictor is not None:
+        return _damage_predictor
+
+    if (not DETECTRON2_AVAILABLE) or (not model_path) or (not os.path.exists(model_path)):
+        print("ℹ️ Stage 1 damage model not available; using simulator")
+        return None
+
+    try:
+        cfg = get_cfg()
+        cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
+        cfg.MODEL.WEIGHTS = model_path
+        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
+        cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
+
+        if device_str is None:
+            device_str = "cuda" if torch.cuda.is_available() else "cpu"
+        cfg.MODEL.DEVICE = device_str
+
+        _damage_predictor = DefaultPredictor(cfg)
+        print(f"✅ Damage model loaded on {device_str}")
+        return _damage_predictor
+    except Exception as e:
+        print(f"❌ Could not load Detectron2 model: {e}")
+        return None
+
+def simulate_damage_detection(rgb_image: np.ndarray, seed_from: np.ndarray = None) -> List[Dict[str, Any]]:
+    """Deterministic fake detections for demo mode."""
+    import hashlib, random
+    h, w = rgb_image.shape[:2]
+    if seed_from is None:
+        seed_from = rgb_image
+    img_hash = hashlib.md5(seed_from.tobytes()).hexdigest()
+    seed = int(img_hash[:8], 16) % 10_000
+    random.seed(seed)
+    n = random.randint(0, 3)
+    boxes = []
+    for _ in range(n):
+        x1 = random.randint(0, max(0, w - w//3))
+        y1 = random.randint(0, max(0, h - h//3))
+        x2 = min(w-1, x1 + random.randint(w//8, w//3))
+        y2 = min(h-1, y1 + random.randint(h//8, h//3))
+        conf = round(random.uniform(0.6, 0.95), 3)
+        boxes.append({"bbox":[x1,y1,x2,y2], "score":conf, "label":"damage"})
+    return boxes
+
+def run_damage_detection(pil_image: Image.Image, score_thresh: float = 0.5):
+    """Run damage detection with fallback."""
+    try:
+        rgb = np.array(pil_image.convert("RGB"))
+        predictor = load_damage_model(DEFAULT_DAMAGE_MODEL_PATH)
         
-        # Temporarily replace the function
-        mu._load_state_dict_into_meta_model = patched_load
+        if predictor is None:
+            boxes = simulate_damage_detection(rgb, seed_from=rgb)
+            annotated = rgb.copy()
+            for i, b in enumerate(boxes):
+                x1,y1,x2,y2 = b["bbox"]
+                cv2.rectangle(annotated, (x1,y1), (x2,y2), (255,255,0), 2)
+                cv2.putText(annotated, f"Damage {i+1} {b['score']*100:.1f}%",
+                            (x1, max(0,y1-8)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,0), 2)
+            return boxes, annotated, True, "predictor not available"
         
-        try:
-            # Load the model
-            model = AutoModel.from_pretrained(
-                repo_id,
-                trust_remote_code=True,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-            )
-            print("  Model loaded successfully with compatibility fixes!")
+        # Real inference
+        outputs = predictor(rgb)
+        instances = outputs["instances"].to("cpu")
+        boxes = []
+        if len(instances) > 0:
+            pred_boxes = instances.pred_boxes.tensor.numpy()
+            scores = instances.scores.numpy()
+            for i, (box, sc) in enumerate(zip(pred_boxes, scores)):
+                if sc >= score_thresh:
+                    x1,y1,x2,y2 = [int(v) for v in box]
+                    boxes.append({"bbox":[x1,y1,x2,y2], "score":float(sc), "label":"damage"})
+        
+        annotated = rgb.copy()
+        for i, b in enumerate(boxes):
+            x1,y1,x2,y2 = b["bbox"]
+            cv2.rectangle(annotated, (x1,y1), (x2,y2), (255,255,0), 2)
+            cv2.putText(annotated, f"Damage {i+1} {b['score']*100:.1f}%",
+                        (x1, max(0,y1-8)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,0), 2)
+        
+        return boxes, annotated, False, None
+        
+    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)
+        annotated = rgb.copy()
+        for i, b in enumerate(boxes):
+            x1,y1,x2,y2 = b["bbox"]
+            cv2.rectangle(annotated, (x1,y1), (x2,y2), (255,255,0), 2)
+            cv2.putText(annotated, f"Damage {i+1} {b['score']*100:.1f}%",
+                        (x1, max(0,y1-8)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,0), 2)
+        return boxes, annotated, True, "stage1 error"
+
+# --------------------------------------------------------------------------------------
+# Stage 2: Feature extraction + classifier
+# --------------------------------------------------------------------------------------
+
+def load_ai_detection_classifier(model_path):
+    """Load the AI detection classifier (joblib)."""
+    global ai_detection_classifier
+    if ai_detection_classifier is not None:
+        print("✅ Classifier already loaded, reusing...")
+        return ai_detection_classifier
+
+    if model_path is None or not os.path.exists(model_path):
+        print(f"❌ AI detection model not found at: {model_path}")
+        return None
+
+    try:
+        ai_detection_classifier = joblib.load(model_path)
+        print(f"✅ AI detection classifier loaded from {model_path}")
+        print(f"   Classifier type: {type(ai_detection_classifier).__name__}")
+        return ai_detection_classifier
+    except Exception as e:
+        print(f"❌ Error loading classifier: {e}")
+        return None
+
+def preprocess_image(image) -> Optional[Image.Image]:
+    """Robust image preprocessing."""
+    try:
+        if image is None:
+            return None
+
+        if isinstance(image, Image.Image):
+            pil = image
+        elif isinstance(image, str):
+            pil = Image.open(image)
+        elif isinstance(image, np.ndarray):
+            if image.ndim == 2:
+                image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+            elif image.ndim == 3 and image.shape[2] == 4:
+                image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
+            
+            if image.dtype != np.uint8:
+                if image.max() <= 1.0:
+                    image = (image * 255).astype(np.uint8)
+                else:
+                    image = np.clip(image, 0, 255).astype(np.uint8)
             
-        finally:
-            # Restore original function
-            mu._load_state_dict_into_meta_model = original_load
+            pil = Image.fromarray(image, 'RGB')
+        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')
+
+        # Handle EXIF orientation
+        pil = ImageOps.exif_transpose(pil)
+        return pil.convert("RGB")
         
-        return model
+    except Exception as e:
+        print(f"❌ Preprocess error: {e}")
+        traceback.print_exc()
+        return None
 
-# Usage:
-model = RADIOModelFixed.from_pretrained()
-'''
-    
-    print(code)
-    
-    # Save to file
-    with open('radio_loader_fixed.py', 'w') as f:
-        f.write(code)
+def extract_features(image, return_stats=False):
+    """Extract features with proper handling for different model types."""
+    global image_processor, model
     
-    print("\n✅ Fixed loader saved to 'radio_loader_fixed.py'")
+    if image_processor is None or model is None:
+        raise Exception("Model not initialized")
 
-# Step 4: Alternative loading approach
-def try_alternative_loading():
-    """Try loading the model with different strategies"""
-    
-    print("\n🔄 Trying Alternative Loading Methods...")
-    
-    from transformers import AutoModel, AutoConfig
-    import transformers.modeling_utils
-    
-    repo_id = "nvidia/C-RADIOv3-B"
-    
-    # Method 1: Load config first and check architecture
-    print("\n1️⃣ Checking model config...")
-    config = AutoConfig.from_pretrained(repo_id, trust_remote_code=True)
-    print(f"  Architecture: {config.architectures}")
-    print(f"  Model type: {config.model_type}")
-    
-    # Method 2: Try loading without state dict verification
-    print("\n2️⃣ Attempting to load with strict=False...")
-    
-    # Monkey-patch the DINOv2 architecture
-    import importlib.util
-    import sys
-    
-    # Download the dinov2_arch.py
-    dinov2_path = hf_hub_download(repo_id=repo_id, filename="dinov2_arch.py")
-    
-    # Load it as a module
-    spec = importlib.util.spec_from_file_location("dinov2_arch_patched", dinov2_path)
-    dinov2_module = importlib.util.module_from_spec(spec)
-    
-    # Patch the _load_from_state_dict method before loading
-    original_code = open(dinov2_path, 'r').read()
-    
-    # Replace the error-raising code
-    patched_code = original_code.replace(
-        'raise KeyError(f"Couldn\'t find the key {key_a} nor {key_b} in the state dict!")',
-        '''
-print(f"  Warning: Missing keys {key_a} and {key_b}, using defaults")
-# Use identity/ones as default
-if "gamma" in key_a:
-    setattr(self, key_a.split(".")[-1], torch.nn.Parameter(torch.ones(self.dim)))
-elif "beta" in key_a:
-    setattr(self, key_a.split(".")[-1], torch.nn.Parameter(torch.zeros(self.dim)))
-return
-'''
-    )
+    if not isinstance(image, Image.Image):
+        image = preprocess_image(image)
+        if image is None:
+            raise Exception("Failed to preprocess image")
+
+    # Resize to 224x224
+    image = image.resize((224, 224), Image.Resampling.LANCZOS)
+    
+    # Process image
+    inputs = image_processor(images=image, return_tensors='pt', do_resize=True)
+    
+    # Handle different processor outputs
+    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)
+    
+    # Get features
+    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']
+        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
+    
+    # Pool if needed
+    if features.ndim == 3:  # (B, T, C)
+        features = features.mean(dim=1)
+    elif features.ndim == 4:  # (B, C, H, W)
+        features = features.mean(dim=(2, 3))
+    
+    # Normalize and flatten
+    features = features.detach().flatten()
+    features = F.normalize(features, p=2, dim=-1).cpu().numpy()
+    
+    if return_stats:
+        stats = {
+            "mean": float(features.mean()),
+            "std": float(features.std()),
+            "min": float(features.min()),
+            "max": float(features.max()),
+            "shape": features.shape
+        }
+        return features, stats
+    
+    return features
+
+def simulate_prediction(image) -> Dict[str, Any]:
+    """Fallback simulation when models/classifier aren't available."""
+    import hashlib, random
+    
+    if isinstance(image, Image.Image):
+        arr = np.array(image.convert("RGB"))
+    elif isinstance(image, np.ndarray):
+        arr = image
+    else:
+        arr = np.array(preprocess_image(image) or Image.new("RGB",(16,16),(0,0,0)))
+
+    img_hash = hashlib.md5(arr.tobytes()).hexdigest()
+    seed = int(img_hash[:8], 16) % 1000
+    random.seed(seed)
+    ai_prob = random.uniform(0.1, 0.9)
+    is_ai = ai_prob > 0.5
+    confidence_level = "HIGH" if abs(ai_prob - 0.5) > 0.3 else "MEDIUM" if abs(ai_prob - 0.5) > 0.15 else "LOW"
+
+    return {
+        "prediction": "AI-Generated" if is_ai else "Real",
+        "ai_probability": ai_prob,
+        "real_probability": 1 - ai_prob,
+        "confidence": confidence_level,
+        "is_demo": True
+    }
+
+def _predict_with_classifier(classifier, features: np.ndarray) -> Tuple[int, float, float]:
+    """Predict with classifier."""
+    features = features.reshape(1, -1)
+    pred = int(classifier.predict(features)[0])
     
-    # Save patched version
-    patched_path = "dinov2_arch_patched.py"
-    with open(patched_path, 'w') as f:
-        f.write(patched_code)
+    ai_prob = real_prob = 0.5
     
-    print(f"  Created patched architecture file: {patched_path}")
+    if hasattr(classifier, "predict_proba"):
+        try:
+            probs = classifier.predict_proba(features)[0]
+            if len(probs) >= 2:
+                real_prob = float(probs[0])
+                ai_prob = float(probs[1])
+            else:
+                ai_prob = float(probs[0]) if pred == 1 else 1 - float(probs[0])
+                real_prob = 1 - ai_prob
+        except:
+            pass
+    elif hasattr(classifier, "decision_function"):
+        try:
+            df = float(classifier.decision_function(features)[0])
+            ai_prob = 1.0 / (1.0 + np.exp(-df))
+            real_prob = 1.0 - ai_prob
+        except:
+            pass
     
-    print("\n✅ Alternative loading methods prepared")
+    return pred, ai_prob, real_prob
+
+# --------------------------------------------------------------------------------------
+# Gradio Interface
+# --------------------------------------------------------------------------------------
+
+def create_gradio_interface():
+    """Create the Gradio interface."""
+    
+    with gr.Blocks(title="AI Image Detection", css=".gradio-container { font-family: Inter, system-ui; }") as app:
+        gr.HTML("""
+        <div style="text-align: center; padding: 20px; background: linear-gradient(90deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 10px;">
+            <h1 style="margin: 0;">🤖 AI Image Detection</h1>
+            <p style="margin: 10px 0 0 0;">Stage 1 (Damage) + Stage 2 (AI Detection)</p>
+        </div>
+        """)
+
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(
+                    type="numpy",
+                    label="Upload Image",
+                    height=400
+                )
+                
+                with gr.Row():
+                    predict_btn = gr.Button("🔍 Analyze", variant="primary", size="lg")
+                    clear_btn = gr.Button("🗑️ Clear", variant="secondary", size="lg")
+                
+                enable_damage = gr.Checkbox(value=True, label="Enable Stage 1 (Damage Detection)")
+                damage_thresh = gr.Slider(0.1, 0.95, value=0.5, step=0.05, label="Damage Score Threshold")
+
+            with gr.Column():
+                output_text = gr.Textbox(
+                    label="Prediction Result",
+                    placeholder="Upload an image and click Analyze",
+                    interactive=False,
+                    lines=2
+                )
+                output_json = gr.JSON(label="Detailed Analysis (Stage 2)")
+                damage_json = gr.JSON(label="Stage 1: Damage Detections")
+                annotated_image = gr.Image(label="Annotated Output")
+                status_display = gr.HTML("""
+                <div style="padding: 10px; background: #f0f4f8; border-radius: 8px; margin-top: 10px;">
+                    <p style="margin: 0; color: #64748b;">Ready for analysis...</p>
+                </div>
+                """)
+
+        def analyze_with_status(image, enable_damage, damage_thresh):
+            """Analyze image."""
+            if image is None:
+                return (
+                    "❌ No image provided",
+                    {"error": "No image provided"},
+                    '<div style="padding: 10px; background: #fee2e2; border-radius: 8px;"><p style="margin: 0; color: #dc2626;">❌ No image provided</p></div>',
+                    [],
+                    None
+                )
+
+            # Initialize models
+            model_initialized = (image_processor is not None and model is not None) or preload_models()
+            model_path = huggingface_model_path or DEFAULT_AI_DETECTION_MODEL_PATH
+            classifier = ai_detection_classifier or load_ai_detection_classifier(model_path)
+
+            demo_reasons = []
+            if not model_initialized:
+                demo_reasons.append("feature extractor missing")
+            if classifier is None:
+                demo_reasons.append("classifier missing")
+
+            # Stage 2: AI Detection
+            try:
+                if demo_reasons:
+                    result = simulate_prediction(preprocess_image(image))
+                    result["demo_reasons"] = demo_reasons
+                    simple_result = f"{result['prediction']} (AI: {result['ai_probability']:.2%}) [Demo]"
+                    detailed_result = result
+                else:
+                    feats, stats = extract_features(preprocess_image(image), return_stats=True)
+                    pred, ai_prob, real_prob = _predict_with_classifier(classifier, feats)
+                    is_ai = pred == 1
+                    result_text = "AI-Generated" if is_ai else "Real"
+                    conf_score = max(ai_prob, real_prob)
+                    confidence = "HIGH" if conf_score > 0.80 else "MEDIUM" if conf_score > 0.60 else "LOW"
+                    
+                    detailed_result = {
+                        "prediction": result_text,
+                        "ai_probability": ai_prob,
+                        "real_probability": real_prob,
+                        "confidence": confidence,
+                        "confidence_score": conf_score,
+                        "is_demo": False,
+                        "feature_stats": stats
+                    }
+                    simple_result = f"{result_text} (Confidence: {conf_score:.2%})"
+                    
+            except Exception as e:
+                print(f"❌ Stage 2 error: {e}")
+                traceback.print_exc()
+                result = simulate_prediction(preprocess_image(image))
+                result["demo_reasons"] = ["stage2 error"]
+                simple_result = f"{result['prediction']} (AI: {result['ai_probability']:.2%}) [Demo]"
+                detailed_result = result
+
+            # Stage 1: Damage Detection
+            dmg_results = []
+            annotated = None
+            
+            if enable_damage:
+                try:
+                    pil = preprocess_image(image)
+                    if pil:
+                        boxes, annotated_rgb, demo, reason = run_damage_detection(pil, float(damage_thresh))
+                        dmg_results = boxes
+                        annotated = annotated_rgb
+                        
+                        # Add verdict overlay
+                        if annotated is not None and isinstance(detailed_result, dict):
+                            is_ai = (detailed_result.get("prediction") == "AI-Generated")
+                            ai_prob = float(detailed_result.get("ai_probability", 0.5))
+                            color = (0,0,255) if is_ai else (0,255,0)
+                            verdict = detailed_result.get("prediction", "Unknown")
+                            
+                            cv2.putText(annotated, verdict, (30, 50), 
+                                      cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 3)
+                            cv2.putText(annotated, f"Confidence: {ai_prob*100:.1f}%", 
+                                      (30, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
+                            
+                except Exception as e:
+                    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>'
+            else:
+                status_html = '<div style="padding: 10px; background: #d1fae5; border-radius: 8px;"><p style="margin: 0; color: #10b981;">✅ Analysis Complete</p></div>'
+
+            return simple_result, detailed_result, status_html, dmg_results, annotated
+
+        def clear_all():
+            """Clear all fields."""
+            return (
+                None,
+                "",
+                {},
+                '<div style="padding: 10px; background: #f0f4f8; border-radius: 8px;"><p style="margin: 0; color: #64748b;">Ready for analysis...</p></div>',
+                [],
+                None
+            )
+
+        # Wire up events
+        predict_btn.click(
+            fn=analyze_with_status,
+            inputs=[input_image, enable_damage, damage_thresh],
+            outputs=[output_text, output_json, status_display, damage_json, annotated_image]
+        )
+
+        clear_btn.click(
+            fn=clear_all,
+            outputs=[input_image, output_text, output_json, status_display, damage_json, annotated_image]
+        )
+
+        # Auto-analyze on image change
+        input_image.change(
+            fn=analyze_with_status,
+            inputs=[input_image, enable_damage, damage_thresh],
+            outputs=[output_text, output_json, status_display, damage_json, annotated_image]
+        )
+
+        with gr.Accordion("ℹ️ About", open=False):
+            gr.Markdown("""
+            ### Pipeline
+            - **Stage 1**: Detectron2 damage detection (optional)
+            - **Stage 2**: Visual features + AI detection classifier
+            
+            ### Notes
+            - Falls back to demo mode if models are unavailable
+            - C-RADIOv3-B model includes compatibility fixes for layer scaling issues
+            """)
+
+    return app
+
+# --------------------------------------------------------------------------------------
+# Main
+# --------------------------------------------------------------------------------------
 
-# Run all inspections
 if __name__ == "__main__":
-    try:
-        # Step 1: Inspect weights
-        state_dict, keys = inspect_model_weights()
-        
-        # Step 2: Inspect code
-        inspect_model_code()
-        
-        # Step 3: Create fixed loader
-        create_fixed_loader()
-        
-        # Step 4: Try alternatives
-        try_alternative_loading()
-        
-        print("\n" + "=" * 60)
-        print("DIAGNOSIS COMPLETE")
-        print("=" * 60)
-        
-    
-    
+    print("=" * 60)
+    print("🚀 Starting AI Image Detection App")
+    print("=" * 60)
+    print(f"📍 Device: {DEVICE}")
+    print(f"📦 Classifier: {huggingface_model_path or DEFAULT_AI_DETECTION_MODEL_PATH}")
+    print(f"🛠️ Damage Model: {DEFAULT_DAMAGE_MODEL_PATH}")
+    
+    # Preload models with fixes
+    if preload_models():
+        print("✅ Models preloaded successfully")
+    else:
+        print("⚠️ Running in 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("=" * 60)
+    
+    app = create_gradio_interface()
+    app.launch(
+        share=False,
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True
+    )