Update app.py

app.py

@@ -1,394 +1,268 @@
 import os
-import sys
-import warnings
-import traceback
-import json
-import pickle
 import torch
-import numpy as np
-from PIL import Image
-import gradio as gr
-from typing import Optional, Tuple, Dict, Any
-
-# ============================================================
-# CONFIGURATION AND SETUP
-# ============================================================
+import json
+from huggingface_hub import hf_hub_download
+import safetensors.torch
 
 print("=" * 60)
-print("DEBUGGING APP.PY - Model Loading Diagnostics")
+print("C-RADIOv3-B Model Deep Inspection")
 print("=" * 60)
 
-# Environment info
-print("\n📋 Environment Information:")
-print(f"Python version: {sys.version}")
-print(f"PyTorch version: {torch.__version__}")
-print(f"CUDA available: {torch.cuda.is_available()}")
-if torch.cuda.is_available():
-    print(f"CUDA version: {torch.version.cuda}")
-    print(f"GPU: {torch.cuda.get_device_name(0)}")
-
-# Check environment variables
-print("\n📦 Environment Variables:")
-env_vars = ['MODEL_REPO', 'HF_TOKEN', 'CUDA_VISIBLE_DEVICES']
-for var in env_vars:
-    value = os.getenv(var, 'NOT SET')
-    if var == 'HF_TOKEN' and value != 'NOT SET':
-        value = value[:10] + '...' if len(value) > 10 else value
-    print(f"  {var}: {value}")
-
-# Device configuration
-RADIO_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-AI_DETECT_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"\n🖥️ Using device: {RADIO_DEVICE}")
-
-# Global variables
-radio_l_image_processor = None
-radio_l_model = None
-ai_detection_classifier = None
-
-# ============================================================
-# DEBUGGING UTILITIES
-# ============================================================
-
-def inspect_model_architecture(model, model_name="Model"):
-    """Inspect and print model architecture details"""
-    print(f"\n🔍 Inspecting {model_name}:")
-    print(f"  Model type: {type(model).__name__}")
+# Step 1: Download and inspect the model file directly
+def inspect_model_weights():
+    """Directly inspect the safetensors file to see what keys exist"""
     
-    # Get state dict
-    try:
-        state_dict = model.state_dict()
-        print(f"  Total parameters: {len(state_dict)}")
-        
-        # Look for specific patterns
-        patterns = ['ls1', 'blocks', 'gamma', 'grandma', 'layer_norm', 'ln']
-        for pattern in patterns:
-            matching_keys = [k for k in state_dict.keys() if pattern in k.lower()]
-            if matching_keys:
-                print(f"  Keys containing '{pattern}': {len(matching_keys)}")
-                print(f"    First 3: {matching_keys[:3]}")
-        
-        # Show first 10 keys
-        print(f"  First 10 state dict keys:")
-        for i, key in enumerate(list(state_dict.keys())[:10]):
-            print(f"    {i+1}. {key}")
-            
-    except Exception as e:
-        print(f"  Could not inspect state dict: {e}")
-    
-    # Check for specific attributes
-    attrs_to_check = ['blocks', 'layers', 'encoder', 'decoder', 'visual', 'text']
-    print(f"  Model attributes:")
-    for attr in attrs_to_check:
-        if hasattr(model, attr):
-            print(f"    ✓ Has '{attr}'")
-            
-def test_pickle_file(filepath):
-    """Test loading a pickle file and inspect its contents"""
-    print(f"\n🥒 Testing pickle file: {filepath}")
-    try:
-        with open(filepath, 'rb') as f:
-            obj = pickle.load(f)
-        print(f"  ✓ Successfully loaded pickle")
-        print(f"  Object type: {type(obj).__name__}")
-        
-        if hasattr(obj, '__dict__'):
-            print(f"  Object attributes: {list(obj.__dict__.keys())[:5]}")
-        
-        if hasattr(obj, 'get_params'):
-            params = obj.get_params()
-            print(f"  Model parameters: {list(params.keys())[:5]}")
-            
-        return obj
-    except Exception as e:
-        print(f"  ✗ Failed to load pickle: {e}")
-        return None
-
-# ============================================================
-# MODEL LOADING FUNCTIONS
-# ============================================================
+    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 preload_c_model_debug():
-    """Debug version of C model preloading"""
-    global radio_l_image_processor, radio_l_model
+# Step 2: Check the model architecture expectations
+def inspect_model_code():
+    """Download and inspect the model code to understand what it expects"""
     
-    print("\n" + "=" * 60)
-    print("LOADING C MODEL")
-    print("=" * 60)
+    print("\n📜 Downloading model code...")
     
-    hf_repo = os.getenv('MODEL_REPO', 'fallback')
-    print(f"Repository: {hf_repo}")
+    # Download the dinov2_arch.py file
+    dinov2_path = hf_hub_download(
+        repo_id="nvidia/C-RADIOv3-B",
+        filename="dinov2_arch.py"
+    )
     
-    if not hf_repo or hf_repo == 'fallback':
-        print("⚠️ No MODEL_REPO environment variable set")
-        return False
+    print(f"Downloaded dinov2_arch.py to: {dinov2_path}")
     
-    try:
-        print("\n1️⃣ Importing transformers...")
-        from transformers import AutoModel, CLIPImageProcessor, AutoConfig
-        print("   ✓ Imports successful")
+    # 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()}")
+    
+    # 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"""
+    
+    print("\n🔧 Creating Fixed Model Loader...")
+    
+    # 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"""
         
-        # Try to load config first
-        print("\n2️⃣ Loading model config...")
-        try:
-            config = AutoConfig.from_pretrained(hf_repo, trust_remote_code=True)
-            print(f"   ✓ Config loaded: {type(config).__name__}")
-            if hasattr(config, 'architectures'):
-                print(f"   Architectures: {config.architectures}")
-        except Exception as e:
-            print(f"   ⚠️ Could not load config: {e}")
+        print("Loading with compatibility fixes...")
         
-        # Load image processor
-        print("\n3️⃣ Loading image processor...")
-        try:
-            radio_l_image_processor = CLIPImageProcessor.from_pretrained(hf_repo)
-            print(f"   ✓ Image processor loaded: {type(radio_l_image_processor).__name__}")
-        except Exception as e:
-            print(f"   ✗ Failed to load image processor: {e}")
-            # Try alternative processors
-            print("   Trying alternative processors...")
-            from transformers import AutoImageProcessor
-            try:
-                radio_l_image_processor = AutoImageProcessor.from_pretrained(hf_repo)
-                print(f"   ✓ Alternative processor loaded: {type(radio_l_image_processor).__name__}")
-            except:
-                print("   ✗ All processor attempts failed")
+        # First, modify the environment to skip the problematic check
+        import sys
+        import transformers.modeling_utils as mu
         
-        # Load model with detailed error catching
-        print("\n4️⃣ Loading model...")
-        try:
-            # First attempt - standard loading
-            radio_l_model = AutoModel.from_pretrained(
-                hf_repo, 
-                trust_remote_code=True,
-                torch_dtype=torch.float32,
-                device_map='auto' if torch.cuda.is_available() else None
-            )
-            print(f"   ✓ Model loaded: {type(radio_l_model).__name__}")
+        # Store original function
+        original_load = mu._load_state_dict_into_meta_model
+        
+        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"""
             
-        except Exception as e:
-            print(f"   ✗ Standard loading failed: {e}")
-            print("   Trying with force_download...")
+            # Create a modified state dict with dummy ls1 keys if needed
+            modified_state = state_dict.copy()
             
-            # Second attempt - force download
-            try:
-                radio_l_model = AutoModel.from_pretrained(
-                    hf_repo, 
-                    trust_remote_code=True,
-                    force_download=True,
-                    torch_dtype=torch.float32
-                )
-                print(f"   ✓ Model loaded with force_download")
-            except Exception as e2:
-                print(f"   ✗ Force download failed: {e2}")
-                raise
-        
-        # Inspect the loaded model
-        inspect_model_architecture(radio_l_model, "C Model")
+            # 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...")
+                
+                # 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)
+            
+            # Call original with modified state
+            return original_load(model, modified_state, device_map, offload_folder,
+                               dtype, offload_state_dict, tie_weights, **kwargs)
         
-        # Move to device
-        print(f"\n5️⃣ Moving model to {RADIO_DEVICE}...")
-        if RADIO_DEVICE.type != 'cpu':
-            radio_l_model = radio_l_model.to(RADIO_DEVICE)
-        radio_l_model.eval()
-        print("   ✓ Model moved and set to eval mode")
+        # Temporarily replace the function
+        mu._load_state_dict_into_meta_model = patched_load
         
-        # Test forward pass
-        print("\n6️⃣ Testing forward pass...")
         try:
-            with torch.no_grad():
-                # Create a dummy image
-                dummy_image = Image.new('RGB', (224, 224), color='white')
-                if radio_l_image_processor:
-                    inputs = radio_l_image_processor(dummy_image, return_tensors="pt")
-                    if RADIO_DEVICE.type != 'cpu':
-                        inputs = {k: v.to(RADIO_DEVICE) for k, v in inputs.items()}
-                    
-                    # Try forward pass
-                    outputs = radio_l_model(**inputs)
-                    print(f"   ✓ Forward pass successful!")
-                    print(f"   Output type: {type(outputs)}")
-                    if hasattr(outputs, 'keys'):
-                        print(f"   Output keys: {outputs.keys()}")
-                else:
-                    print("   ⚠️ No image processor available for test")
-        except Exception as e:
-            print(f"   ✗ Forward pass failed: {e}")
-            traceback.print_exc()
-        
-        print("\n✅ C model loading completed (with warnings)")
-        return True
+            # 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!")
+            
+        finally:
+            # Restore original function
+            mu._load_state_dict_into_meta_model = original_load
         
-    except Exception as e:
-        print(f"\n❌ C model loading failed completely: {e}")
-        traceback.print_exc()
-        return False
+        return model
 
-def preload_ai_detector_debug():
-    """Debug version of AI detector preloading"""
-    global ai_detection_classifier
+# Usage:
+model = RADIOModelFixed.from_pretrained()
+'''
     
-    print("\n" + "=" * 60)
-    print("LOADING AI DETECTION CLASSIFIER")
-    print("=" * 60)
+    print(code)
     
-    try:
-        print("\n1️⃣ Checking for Askhedi model...")
-        from huggingface_hub import hf_hub_download, list_repo_files
-        
-        repo_id = "Askhedi/hedi_v0_mix"
-        print(f"   Repository: {repo_id}")
-        
-        # List files in repo
-        try:
-            files = list(list_repo_files(repo_id))
-            print(f"   Files in repo: {files[:10]}")  # Show first 10 files
-            pkl_files = [f for f in files if f.endswith('.pkl')]
-            print(f"   PKL files found: {pkl_files}")
-        except Exception as e:
-            print(f"   Could not list repo files: {e}")
-        
-        print("\n2️⃣ Downloading classifier...")
-        classifier_path = hf_hub_download(
-            repo_id=repo_id,
-            filename="V1.pkl"
-        )
-        print(f"   ✓ Downloaded to: {classifier_path}")
-        
-        print("\n3️⃣ Loading classifier...")
-        ai_detection_classifier = test_pickle_file(classifier_path)
-        
-        if ai_detection_classifier:
-            print("\n4️⃣ Testing classifier...")
-            try:
-                # Create dummy features
-                dummy_features = np.random.randn(1, 100)  # Adjust size as needed
-                prediction = ai_detection_classifier.predict(dummy_features)
-                print(f"   ✓ Prediction successful: {prediction}")
-                
-                if hasattr(ai_detection_classifier, 'predict_proba'):
-                    proba = ai_detection_classifier.predict_proba(dummy_features)
-                    print(f"   ✓ Probability shape: {proba.shape}")
-                    
-            except Exception as e:
-                print(f"   ⚠️ Classifier test failed: {e}")
-                print(f"   This might be due to incorrect feature dimensions")
-        
-        print("\n✅ AI detection classifier loaded")
-        return True
-        
-    except Exception as e:
-        print(f"\n❌ AI detector loading failed: {e}")
-        traceback.print_exc()
-        return False
-
-# ============================================================
-# MAIN GRADIO APP
-# ============================================================
+    # Save to file
+    with open('radio_loader_fixed.py', 'w') as f:
+        f.write(code)
+    
+    print("\n✅ Fixed loader saved to 'radio_loader_fixed.py'")
 
-def analyze_image_debug(image):
-    """Debug version of image analysis"""
-    results = {
-        "status": "Processing",
-        "c_model": "Not loaded",
-        "ai_detection": "Not loaded",
-        "errors": []
-    }
+# Step 4: Alternative loading approach
+def try_alternative_loading():
+    """Try loading the model with different strategies"""
     
-    try:
-        if radio_l_model and radio_l_image_processor:
-            results["c_model"] = "Model loaded and ready"
-            # Add actual processing here if needed
-        else:
-            results["c_model"] = "Model not loaded"
-            
-        if ai_detection_classifier:
-            results["ai_detection"] = "Classifier loaded and ready"
-            # Add actual detection here if needed
-        else:
-            results["ai_detection"] = "Classifier not loaded"
-            
-        results["status"] = "Complete"
-        
-    except Exception as e:
-        results["errors"].append(str(e))
-        results["status"] = "Error"
+    print("\n🔄 Trying Alternative Loading Methods...")
+    
+    from transformers import AutoModel, AutoConfig
+    import transformers.modeling_utils
+    
+    repo_id = "nvidia/C-RADIOv3-B"
     
-    return json.dumps(results, indent=2)
+    # 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
+'''
+    )
+    
+    # Save patched version
+    patched_path = "dinov2_arch_patched.py"
+    with open(patched_path, 'w') as f:
+        f.write(patched_code)
+    
+    print(f"  Created patched architecture file: {patched_path}")
+    
+    print("\n✅ Alternative loading methods prepared")
 
-def create_gradio_interface():
-    """Create the Gradio interface"""
-    with gr.Blocks(title="Model Loading Debugger") as demo:
-        gr.Markdown("# Model Loading Debugger")
-        gr.Markdown("This interface helps debug model loading issues.")
+# Run all inspections
+if __name__ == "__main__":
+    try:
+        # Step 1: Inspect weights
+        state_dict, keys = inspect_model_weights()
         
-        with gr.Tab("Status"):
-            gr.Markdown("## Current Model Status")
-            status_text = gr.Textbox(
-                label="Model Status",
-                value=f"C Model: {'Loaded' if radio_l_model else 'Not loaded'}\n"
-                      f"AI Detector: {'Loaded' if ai_detection_classifier else 'Not loaded'}",
-                lines=10
-            )
-            
-            refresh_btn = gr.Button("Refresh Status")
-            
-            def refresh_status():
-                status = []
-                status.append(f"C Model: {'✓ Loaded' if radio_l_model else '✗ Not loaded'}")
-                status.append(f"Image Processor: {'✓ Loaded' if radio_l_image_processor else '✗ Not loaded'}")
-                status.append(f"AI Detector: {'✓ Loaded' if ai_detection_classifier else '✗ Not loaded'}")
-                status.append(f"\nDevice: {RADIO_DEVICE}")
-                
-                if radio_l_model:
-                    status.append(f"Model type: {type(radio_l_model).__name__}")
-                    
-                return "\n".join(status)
-            
-            refresh_btn.click(refresh_status, outputs=status_text)
+        # 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)
         
-        with gr.Tab("Test Image"):
-            image_input = gr.Image(label="Upload Test Image", type="pil")
-            analyze_btn = gr.Button("Analyze")
-            output = gr.Textbox(label="Analysis Results", lines=10)
-            
-            analyze_btn.click(analyze_image_debug, inputs=image_input, outputs=output)
     
-    return demo
-
-# ============================================================
-# MAIN EXECUTION
-# ============================================================
-
-if __name__ == "__main__":
-    print("\n" + "=" * 60)
-    print("STARTING MODEL PRELOAD")
-    print("=" * 60)
-    
-    # Suppress specific warnings if needed
-    warnings.filterwarnings("ignore", message="Couldn't find the key")
-    warnings.filterwarnings("ignore", category=UserWarning, module="sklearn")
-    
-    # Load models
-    c_model_success = preload_c_model_debug()
-    ai_detector_success = preload_ai_detector_debug()
-    
-    # Summary
-    print("\n" + "=" * 60)
-    print("LOADING SUMMARY")
-    print("=" * 60)
-    print(f"C Model: {'✅ Success' if c_model_success else '❌ Failed'}")
-    print(f"AI Detector: {'✅ Success' if ai_detector_success else '❌ Failed'}")
-    
-    # Launch Gradio
-    print("\n" + "=" * 60)
-    print("LAUNCHING GRADIO INTERFACE")
-    print("=" * 60)
-    
-    demo = create_gradio_interface()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        debug=True
-    )
+