Update app.py

app.py

@@ -4,27 +4,35 @@ from PIL import Image
 import torch
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 
-# --- Configuration ---
 # --- Configuration ---
 MODEL_NAME = "nateraw/fer-2013"
-DEVICE = "cpu"  # FORCING CPU to ensure model loads correctly on basic hardware
+# CRITICAL FIX: Explicitly set DEVICE to 'cpu' to prevent CUDA initialization errors 
+# and memory issues on default Hugging Face Spaces hardware.
+DEVICE = "cpu" 
 
 # --- Model and Processor Loading ---
+# Load model outside the prediction function for efficiency
 try:
     processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
-    model = AutoModelForImageClassification.from_pretrained(MODEL_NAME).to(DEVICE)
+    
+    # CRITICAL FIX: Load model with map_location='cpu' to prevent Out-Of-Memory (OOM) 
+    # errors during the loading process by mapping all tensors directly to CPU memory.
+    model = AutoModelForImageClassification.from_pretrained(
+        MODEL_NAME, 
+        map_location=DEVICE 
+    ).to(DEVICE)
+    
     model.eval()
     LABELS = model.config.id2label
     
     print(f"Model loaded successfully on device: {DEVICE}")
-    print(f"Available labels: {LABELS}")
 
 except Exception as e:
-    print(f"Error loading model or processor: {e}")
-    # Use a dummy function if loading fails
+    # If loading fails, ensure the error message is descriptive.
+    print(f"CRITICAL ERROR during model loading: {e}")
     processor = None
     model = None
-    LABELS = {0: "dummy_emotion"}
+    LABELS = {0: "Load_Error"}
 
 # --- Inference Function ---
 def classify_emotion(image_np: np.ndarray) -> str:
@@ -32,14 +40,13 @@ def classify_emotion(image_np: np.ndarray) -> str:
     Performs emotion classification on an input image (numpy array).
     """
     if model is None or processor is None:
-        return f"Error: Model or processor failed to load. Check logs."
+        return "System Error: Model failed to initialize. Please perform a Factory Reboot or check Space logs."
 
     try:
         # Convert numpy array (from Gradio) to PIL Image
         image = Image.fromarray(image_np).convert("RGB")
         
         # Preprocess the image
-        # The processor handles necessary resizing and normalization
         inputs = processor(images=image, return_tensors="pt").to(DEVICE)
         
         # Run inference
@@ -47,8 +54,7 @@ def classify_emotion(image_np: np.ndarray) -> str:
             outputs = model(**inputs)
             
         # Get predictions
-        logits = outputs.logits
-        probabilities = torch.nn.functional.softmax(logits, dim=-1)
+        probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
         
         # Find the dominant emotion
         confidence, predicted_class_idx = torch.max(probabilities, 1)
@@ -56,12 +62,15 @@ def classify_emotion(image_np: np.ndarray) -> str:
         dominant_emotion = LABELS[predicted_class_idx.item()]
         confidence_score = confidence.item()
 
-        # Format the result string
-        result_str = f"Emotion: **{dominant_emotion}** ({confidence_score:.2f})"
+        # Format the result with clear markdown
+        result_str = (
+            f"**Predicted Emotion:** **{dominant_emotion.upper()}**\n\n"
+            f"Confidence: {confidence_score:.2f}"
+        )
         return result_str
 
     except Exception as e:
-        return f"Prediction Error: {type(e).__name__} - {str(e)}"
+        return f"Prediction Runtime Error: {type(e).__name__} - {str(e)}"
 
 # --- Gradio Interface ---
 iface = gr.Interface(
@@ -71,21 +80,14 @@ iface = gr.Interface(
         label="Upload an image of a face"
     ),
     outputs=gr.Markdown(label="Predicted Emotion"),
-    title="Emotion Detection (PyTorch/Transformers)",
+    title="😊 PyTorch Facial Emotion Detection",
     description=(
-        "Upload an image containing a face to classify the dominant emotion. "
-        "Uses the **nateraw/fer-2013** PyTorch model from Hugging Face Transformers. "
-        "No TensorFlow or Keras dependencies."
+        "Upload an image to classify the dominant emotion. Uses the **nateraw/fer-2013** PyTorch model. "
+        "Built for stable deployment on Hugging Face Spaces."
     ),
-    examples=[
-        # Providing simple examples is good practice
-        # Note: Gradio will handle downloading and using these if they exist in the repo
-        # Since this is for a new Space, use placeholder or common sense (omit paths if files aren't included)
-    ],
     allow_flagging="never",
     theme=gr.themes.Soft()
 )
 
-# Launch the app (required for Hugging Face Spaces)
 if __name__ == "__main__":
     iface.launch()