Update app.py

app.py

@@ -1,34 +1,394 @@
-def preload_models():
-    """Preload models at startup to improve response time"""
-    global radio_l_image_processor, radio_l_model
+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
+# ============================================================
+
+print("=" * 60)
+print("DEBUGGING APP.PY - Model Loading Diagnostics")
+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}")
 
-    print("🔄 Preloading C model (4GB)...")
+# 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__}")
+    
+    # Get state dict
     try:
-        hf_repo = os.getenv('MODEL_REPO', 'fallback')
-        if hf_repo and hf_repo != 'fallback':
-            from transformers import AutoModel, CLIPImageProcessor
-            
-            # Load the model first to inspect it
-            radio_l_model = AutoModel.from_pretrained(hf_repo, trust_remote_code=True)
+        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}")
             
-            # Debug: Print available keys
-            state_dict = radio_l_model.state_dict()
-            print("Available keys in model (first 10):")
-            for i, key in enumerate(list(state_dict.keys())[:10]):
-                print(f"  {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}'")
             
-            # Check for blocks.0.ls1 related keys
-            ls1_keys = [k for k in state_dict.keys() if 'ls1' in k]
-            if ls1_keys:
-                print(f"Found ls1 keys: {ls1_keys[:5]}")
+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
+# ============================================================
+
+def preload_c_model_debug():
+    """Debug version of C model preloading"""
+    global radio_l_image_processor, radio_l_model
+    
+    print("\n" + "=" * 60)
+    print("LOADING C MODEL")
+    print("=" * 60)
+    
+    hf_repo = os.getenv('MODEL_REPO', 'fallback')
+    print(f"Repository: {hf_repo}")
+    
+    if not hf_repo or hf_repo == 'fallback':
+        print("⚠️ No MODEL_REPO environment variable set")
+        return False
+    
+    try:
+        print("\n1️⃣ Importing transformers...")
+        from transformers import AutoModel, CLIPImageProcessor, AutoConfig
+        print("   ✓ Imports successful")
+        
+        # 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}")
+        
+        # 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")
+        
+        # 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__}")
+            
+        except Exception as e:
+            print(f"   ✗ Standard loading failed: {e}")
+            print("   Trying with force_download...")
+            
+            # 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")
+        
+        # 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("✅ C model preloaded successfully!")
-            return True
+        radio_l_model.eval()
+        print("   ✓ Model moved and set to eval mode")
+        
+        # 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
+        
     except Exception as e:
-        print(f"⚠️ Could not preload C model: {repr(e)}")
-        import traceback
+        print(f"\n❌ C model loading failed completely: {e}")
         traceback.print_exc()
-    return False
+        return False
+
+def preload_ai_detector_debug():
+    """Debug version of AI detector preloading"""
+    global ai_detection_classifier
+    
+    print("\n" + "=" * 60)
+    print("LOADING AI DETECTION CLASSIFIER")
+    print("=" * 60)
+    
+    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
+# ============================================================
+
+def analyze_image_debug(image):
+    """Debug version of image analysis"""
+    results = {
+        "status": "Processing",
+        "c_model": "Not loaded",
+        "ai_detection": "Not loaded",
+        "errors": []
+    }
+    
+    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"
+    
+    return json.dumps(results, indent=2)
+
+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.")
+        
+        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)
+        
+        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
+    )