Update app.py

app.py

@@ -10,43 +10,48 @@ from realesrgan import RealESRGANer
 app = Flask(__name__)
 CORS(app)
 
-# Initialize model
+# Initialize model with explicit device allocation
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
-upsampler = RealESRGANer(
-    scale=4,
-    model_path='weights/realesr-general-x4v3.pth',
-    model=model,
-    tile=0,
-    tile_pad=10,
-    pre_pad=0,
-    half=False,
-    device=device
-)
+
+try:
+    upsampler = RealESRGANer(
+        scale=4,
+        model_path='weights/realesr-general-x4v3.pth',
+        model=model,
+        tile=400,
+        tile_pad=10,
+        pre_pad=0,
+        half=False,
+        device=device
+    )
+except Exception as e:
+    print(f"Model initialization failed: {str(e)}")
+    upsampler = None
 
 @app.route('/enhance', methods=['POST'])
 def enhance_image():
+    if not upsampler:
+        return {'error': 'Model not initialized'}, 500
+    
     if 'file' not in request.files:
         return {'error': 'No file uploaded'}, 400
     
-    file = request.files['file']
-    
     try:
-        # Read image
-        input_img = Image.open(file.stream).convert('RGB')
-        
-        # Convert to numpy array
-        img_array = np.array(input_img)
-        
-        # Enhance image
-        output, _ = upsampler.enhance(img_array, outscale=4)
+        file = request.files['file']
+        img = Image.open(file.stream).convert('RGB')
+        img_array = np.array(img)
         
-        # Convert back to PIL Image
-        result_img = Image.fromarray(output)
+        # Process image
+        output, _ = upsampler.enhance(
+            img_array,
+            outscale=4,
+            alpha_upsampler='realesrgan'
+        )
         
-        # Prepare response
+        # Convert to bytes
         img_byte_arr = io.BytesIO()
-        result_img.save(img_byte_arr, format='JPEG', quality=95)
+        Image.fromarray(output).save(img_byte_arr, format='JPEG')
         img_byte_arr.seek(0)
         
         return send_file(img_byte_arr, mimetype='image/jpeg')
@@ -56,7 +61,8 @@ def enhance_image():
 
 @app.route('/health', methods=['GET'])
 def health_check():
-    return {'status': 'healthy'}, 200
+    status = 'ready' if upsampler else 'initializing'
+    return {'status': status}, 200
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5000)