Update app.py

app.py

@@ -6,20 +6,18 @@ import io
 import torch
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from realesrgan import RealESRGANer
-import threading
 import time
+import os
 
 app = Flask(__name__)
 CORS(app)
 
 # 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
+# Initialize Real-ESRGAN upsampler at startup
+def initialize_enhancer():
+    global upsampler
     
     try:
         print("Starting model initialization...")
@@ -53,28 +51,26 @@ def initialize_enhancer_thread():
         upsampler.enhance(test_img, outscale=4, alpha_upsampler='realesrgan')
         
         print("Model initialization completed successfully")
-        initialization_status = "ready"
+        return True
         
     except Exception as e:
         error_msg = f"Model initialization failed: {str(e)}"
         print(error_msg)
-        initialization_status = "failed"
-        initialization_error = error_msg
+        return False
 
-# Start initialization in background thread
-threading.Thread(target=initialize_enhancer_thread, daemon=True).start()
+# Initialize model immediately at startup - this blocks until model is ready
+if not initialize_enhancer():
+    print("ERROR: Model failed to initialize. Server will exit.")
+    # In production, you might prefer to exit when model fails
+    # os.exit(1)
 
 @app.route('/enhance', methods=['POST'])
 def enhance_image():
-    global upsampler, initialization_status, initialization_error
+    global upsampler
     
-    # 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 upsampler is ready
+    if upsampler is None:
+        return jsonify({'error': 'Enhancement model is not initialized.'}), 500
     
     # Check if file was uploaded
     if 'file' not in request.files:
@@ -114,28 +110,25 @@ def enhance_image():
 
 @app.route('/health', methods=['GET'])
 def health_check():
-    global initialization_status, initialization_error
+    global upsampler
     
     status_info = {
-        'status': initialization_status,
+        'status': 'ready' if upsampler is not None else 'failed',
         '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():
+    global upsampler
     return jsonify({
         'message': 'Image Enhancement API',
         'endpoints': {
             'POST /enhance': 'Process images (4x upscale)',
             'GET /health': 'Service status check'
         },
-        'status': initialization_status
+        'status': 'ready' if upsampler is not None else 'not ready'
     })
 
 if __name__ == '__main__':