Update app/routes.py

app/routes.py

@@ -15,14 +15,15 @@ import torch
 from app.utils.model_cache import model_cache
 from app.config import MAX_FILE_SIZE, ALLOWED_EXTENSIONS
 
-# Import training functions (handle both old and new locations)
+# Import EfficientNet caption generation function
 try:
-    from training.resnet_train import visualize_attention
     from training.efficient_train import generate_caption
 except ImportError:
     # Fallback for backward compatibility (before reorganization)
-    from resnet_train import visualize_attention
-    from efficient_train import generate_caption
+    try:
+        from efficient_train import generate_caption
+    except ImportError:
+        raise ImportError("Could not import generate_caption from training.efficient_train or efficient_train")
 
 logger = logging.getLogger(__name__)
 
@@ -86,7 +87,6 @@ def health_check():
         'status': 'healthy',
         'timestamp': datetime.utcnow().isoformat(),
         'models_loaded': {
-            'resnet': model_cache.is_resnet_loaded(),
             'efficientnet': model_cache.is_efficientnet_loaded()
         }
     }), 200
@@ -94,9 +94,9 @@ def health_check():
 
 @bp.route('/ready')
 def readiness_check():
-    """Readiness check - ensures models are loaded."""
-    if not model_cache.is_resnet_loaded() and not model_cache.is_efficientnet_loaded():
-        return jsonify({'status': 'not ready', 'reason': 'models not loaded'}), 503
+    """Readiness check - ensures EfficientNet model is loaded."""
+    if not model_cache.is_efficientnet_loaded():
+        return jsonify({'status': 'not ready', 'reason': 'EfficientNet model not loaded'}), 503
     return jsonify({'status': 'ready'}), 200
 
 
@@ -160,15 +160,8 @@ def upload_file():
                     model_cache._device, 
                     max_length=64
                 )
-        else:  # resnet50
-            if not model_cache.is_resnet_loaded():
-                return jsonify({'error': 'ResNet model not available'}), 503
-            
-            encoder, decoder, vocab = model_cache.get_resnet_models()
-            
-            # Generate caption
-            with torch.no_grad():
-                caption = visualize_attention(filepath, encoder, decoder, model_cache._device)
+        else:  # resnet50 (not supported - EfficientNet only)
+            return jsonify({'error': 'ResNet model is not available. Please use EfficientNet.'}), 400
         
         inference_time = time.time() - start_time
         logger.info(f"Caption generated in {inference_time:.3f}s using {model_choice}")