Update app.py

app.py

@@ -6,10 +6,11 @@ import gradio as gr
 import os
 import sys
 
-# Check if running on Hugging Face
-IS_HUGGINGFACE = os.environ.get("SPACE_ID") is not None
+print("Starting AI Image Detector...")
+print(f"Working directory: {os.getcwd()}")
+print(f"Files in directory: {os.listdir('.')}")
 
-# Set up device - important to set it right for each environment
+# Set up device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
@@ -20,9 +21,8 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-# Load the model
 def load_model():
-    print("Loading model...")
+    print("Creating model architecture...")
     # Create model architecture
     model = models.efficientnet_v2_s(weights=None)
     
@@ -37,46 +37,42 @@ def load_model():
         nn.Linear(512, 2)
     )
     
-    # Load the state dict - handle both local and HF paths
-    try:
-        model_paths = [
-            'best_model_improved.pth',  # Local path
-            '/repository/best_model_improved.pth',  # Hugging Face path
-            os.path.join(os.path.dirname(os.path.abspath(__file__)), 'best_model_improved.pth')  # Absolute path
-        ]
-        
-        model_loaded = False
-        for model_path in model_paths:
-            if os.path.exists(model_path):
-                print(f"Loading model from: {model_path}")
+    # Try to load from multiple possible locations
+    possible_paths = [
+        "best_model_improved.pth",
+        "pytorch_model.bin",
+        "/repository/best_model_improved.pth",
+        "/repository/pytorch_model.bin",
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), "best_model_improved.pth"),
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), "pytorch_model.bin")
+    ]
+    
+    model_loaded = False
+    for model_path in possible_paths:
+        if os.path.exists(model_path):
+            print(f"Loading model from: {model_path}")
+            try:
                 model.load_state_dict(torch.load(model_path, map_location=device))
                 model_loaded = True
                 break
-        
-        if not model_loaded:
-            print(f"Model not found in any of the expected locations: {model_paths}")
-            print(f"Current directory contents: {os.listdir('.')}")
-            if os.path.exists('/repository'):
-                print(f"Repository directory contents: {os.listdir('/repository')}")
-            raise FileNotFoundError("Model file not found")
-            
-        model.to(device)
-        model.eval()
-        print("Model loaded successfully")
-        return model
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        raise
+            except Exception as e:
+                print(f"Error loading from {model_path}: {e}")
+    
+    if not model_loaded:
+        print("WARNING: Could not load model weights. Using untrained model.")
+    
+    model.to(device)
+    model.eval()
+    return model
 
 # Global model variable
 model = None
 
-# Inference function
 def predict_image(img):
     global model
     
     if img is None:
-        return {"error": "No image provided"}, "Error: No image provided", "Please upload an image"
+        return {"Error": "No image provided"}, "Error: Please upload an image"
     
     try:
         # Load model if not already loaded
@@ -105,14 +101,17 @@ def predict_image(img):
         # Determine classification
         classification = "Real Image" if prediction == 0 else "AI-Generated Image"
         confidence = real_prob if prediction == 0 else ai_prob
+        confidence_text = f"Confidence: {confidence:.2f}%"
         
-        return result, classification, f"Confidence: {confidence:.2f}%"
+        return result, classification + " - " + confidence_text
     
     except Exception as e:
+        import traceback
         print(f"Error during prediction: {e}")
-        return {"error": str(e)}, f"Error: {str(e)}", "Failed to analyze image"
+        traceback.print_exc()
+        return {"error": str(e)}, f"Error: {str(e)}"
 
-# Define Gradio interface
+# Define Gradio interface - simplified for Hugging Face
 def create_interface():
     with gr.Blocks(title="AI Image Detector", theme=gr.themes.Soft()) as interface:
         gr.Markdown("# AI Image Detector")
@@ -124,15 +123,14 @@ def create_interface():
                 analyze_btn = gr.Button("Analyze Image", variant="primary")
             
             with gr.Column():
-                result_label = gr.Label(label="Prediction")
-                classification = gr.Textbox(label="Classification")
-                confidence = gr.Textbox(label="Confidence")
+                result_label = gr.Label(label="Prediction Probabilities")
+                classification = gr.Textbox(label="Classification Result")
         
         # Set up the click event
         analyze_btn.click(
             fn=predict_image, 
             inputs=input_image, 
-            outputs=[result_label, classification, confidence]
+            outputs=[result_label, classification]
         )
         
         gr.Markdown("### How It Works")
@@ -147,18 +145,5 @@ def create_interface():
 
 # Launch the interface
 if __name__ == "__main__":
-    try:
-        print("Starting AI Image Detector application...")
-        interface = create_interface()
-        
-        # Different launch options based on environment
-        if IS_HUGGINGFACE:
-            print("Running on Hugging Face, launching with share=False")
-            interface.launch(share=False)
-        else:
-            print("Running locally, launching with share=True")
-            interface.launch(share=True)
-            
-    except Exception as e:
-        print(f"Error starting application: {e}")
-        sys.exit(1)
+    interface = create_interface()
+    interface.launch()