Update app.py

app.py

@@ -1,4 +1,4 @@
-from flask import Flask, request, send_file
+from flask import Flask, request, send_file, jsonify
 from flask_cors import CORS
 import numpy as np
 from PIL import Image
@@ -6,13 +6,23 @@ import io
 import torch
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from realesrgan import RealESRGANer
+import threading
+import time
 
 app = Flask(__name__)
 CORS(app)
 
-# Initialize Real-ESRGAN upsampler
-def initialize_enhancer():
+# Global variables
+upsampler = None
+initialization_status = "initializing"  # "initializing", "ready", or "failed"
+initialization_error = None
+
+# Initialize Real-ESRGAN upsampler in a separate thread
+def initialize_enhancer_thread():
+    global upsampler, initialization_status, initialization_error
+    
     try:
+        print("Starting model initialization...")
         # Configuration for RealESRGAN x4v3
         model = RRDBNet(
             num_in_ch=3,
@@ -26,44 +36,69 @@ def initialize_enhancer():
         # Force CPU usage for Hugging Face compatibility
         device = torch.device('cpu')
         
-        return RealESRGANer(
+        # Initialize the upsampler
+        upsampler = RealESRGANer(
             scale=4,
             model_path='weights/realesr-general-x4v3.pth',
             model=model,
-            tile=0,        # Set to 0 for small images, increase for large images
+            tile=0,  # Set to 0 for small images, increase for large images
             tile_pad=10,
             pre_pad=0,
-            half=False,    # CPU doesn't support half precision
+            half=False,  # CPU doesn't support half precision
             device=device
         )
+        
+        # Run a small test to ensure everything is loaded
+        test_img = np.zeros((64, 64, 3), dtype=np.uint8)
+        upsampler.enhance(test_img, outscale=4, alpha_upsampler='realesrgan')
+        
+        print("Model initialization completed successfully")
+        initialization_status = "ready"
+        
     except Exception as e:
-        print(f"Initialization error: {str(e)}")
-        return None
+        error_msg = f"Model initialization failed: {str(e)}"
+        print(error_msg)
+        initialization_status = "failed"
+        initialization_error = error_msg
 
-# Global upsampler instance
-upsampler = initialize_enhancer()
+# Start initialization in background thread
+threading.Thread(target=initialize_enhancer_thread, daemon=True).start()
 
 @app.route('/enhance', methods=['POST'])
 def enhance_image():
-    if not upsampler:
-        return {'error': 'Model failed to initialize'}, 500
+    global upsampler, initialization_status, initialization_error
+    
+    # Check if model is ready
+    if initialization_status == "initializing":
+        return jsonify({'error': 'Enhancement model is still initializing. Please try again in a few minutes.'}), 503
     
+    # Check if initialization failed
+    if initialization_status == "failed":
+        return jsonify({'error': f'Model initialization failed: {initialization_error}'}), 500
+    
+    # Check if file was uploaded
     if 'file' not in request.files:
-        return {'error': 'No file uploaded'}, 400
+        return jsonify({'error': 'No file uploaded'}), 400
     
     try:
         # Read and validate image
         file = request.files['file']
         if file.filename == '':
-            return {'error': 'Empty file submitted'}, 400
-            
+            return jsonify({'error': 'Empty file submitted'}), 400
+        
+        # Process the image
         img = Image.open(file.stream).convert('RGB')
         img_array = np.array(img)
         
+        # Check image size
+        h, w = img_array.shape[:2]
+        if h > 2000 or w > 2000:
+            return jsonify({'error': 'Image too large. Maximum size is 2000x2000 pixels'}), 400
+        
         # Enhance image
         output, _ = upsampler.enhance(
             img_array,
-            outscale=4,    # 4x super-resolution
+            outscale=4,  # 4x super-resolution
             alpha_upsampler='realesrgan'
         )
         
@@ -75,22 +110,34 @@ def enhance_image():
         return send_file(img_byte_arr, mimetype='image/jpeg')
     
     except Exception as e:
-        return {'error': f'Processing error: {str(e)}'}, 500
+        return jsonify({'error': f'Processing error: {str(e)}'}), 500
 
 @app.route('/health', methods=['GET'])
 def health_check():
-    status = 'ready' if upsampler else 'unavailable'
-    return {'status': status}, 200
+    global initialization_status, initialization_error
+    
+    status_info = {
+        'status': initialization_status,
+        'timestamp': time.time()
+    }
+    
+    if initialization_status == "failed" and initialization_error:
+        status_info['error'] = initialization_error
+    
+    # Return 200 OK even if not ready, but include status in response
+    return jsonify(status_info)
 
 @app.route('/')
 def home():
-    return {
+    return jsonify({
         'message': 'Image Enhancement API',
         'endpoints': {
             'POST /enhance': 'Process images (4x upscale)',
             'GET /health': 'Service status check'
-        }
-    }, 200
+        },
+        'status': initialization_status
+    })
 
 if __name__ == '__main__':
+    # Start the Flask app
     app.run(host='0.0.0.0', port=5000)