Update app.py

app.py

@@ -1,42 +1,62 @@
-from flask import Flask, request, jsonify, send_file
+from flask import Flask, request, send_file
 from flask_cors import CORS
-import os
+import numpy as np
 from PIL import Image
 import io
 import torch
-from gfpgan import GFPGANer
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from realesrgan import RealESRGANer
 
 app = Flask(__name__)
 CORS(app)
 
-model = GFPGANer(
-    model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.3.pth',
-    upscale=2,
-    arch='clean',
-    channel_multiplier=2,
-    bg_upsampler=None,
-    device='cpu'
+# Initialize model
+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
 )
 
 @app.route('/enhance', methods=['POST'])
-def enhance():
-    if 'image' not in request.files:
-        return jsonify({'error': 'No image uploaded'}), 400
-
-    img_file = request.files['image']
-    img = Image.open(img_file.stream).convert('RGB')
-
+def enhance_image():
+    if 'file' not in request.files:
+        return {'error': 'No file uploaded'}, 400
+    
+    file = request.files['file']
+    
     try:
-        _, _, restored_img = model.enhance(np.array(img), has_aligned=False, only_center_face=True, paste_back=True)
-        result_img = Image.fromarray(restored_img)
-
-        img_io = io.BytesIO()
-        result_img.save(img_io, 'JPEG')
-        img_io.seek(0)
-
-        return send_file(img_io, mimetype='image/jpeg')
+        # 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)
+        
+        # Convert back to PIL Image
+        result_img = Image.fromarray(output)
+        
+        # Prepare response
+        img_byte_arr = io.BytesIO()
+        result_img.save(img_byte_arr, format='JPEG', quality=95)
+        img_byte_arr.seek(0)
+        
+        return send_file(img_byte_arr, mimetype='image/jpeg')
+    
     except Exception as e:
-        return jsonify({'error': str(e)}), 500
+        return {'error': str(e)}, 500
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    return {'status': 'healthy'}, 200
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5000)