Update app.py

app.py

@@ -16,7 +16,7 @@ app = Flask(__name__)
 CORS(app)
 
 # CONFIG
-DEVICE = torch.device('cpu')
+DEVICE = torch.device('cpu') # Hugging Face Free Tier is CPU
 MODEL_PATH = "best_model.pth"
 MODEL_NAME = "nvidia/segformer-b2-finetuned-ade-512-512"
 NUM_CLASSES = 6
@@ -26,8 +26,12 @@ print("Loading model...")
 model = SegformerForSemanticSegmentation.from_pretrained(
     MODEL_NAME, num_labels=NUM_CLASSES, ignore_mismatched_sizes=True
 )
-checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
+
+# --- FIX IS HERE ---
+checkpoint = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=False)
 state_dict = checkpoint['model_state_dict']
+
+# Fix key names (remove 'module.' if trained on multi-GPU)
 new_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
 model.load_state_dict(new_state_dict)
 model.to(DEVICE)
@@ -83,7 +87,5 @@ def predict():
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 
-# --- CRITICAL CHANGE FOR HUGGING FACE ---
 if __name__ == '__main__':
-    # Hugging Face runs on port 7860
     app.run(host='0.0.0.0', port=7860)