Spaces:
Runtime error
Runtime error
File size: 3,400 Bytes
6cd4595 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import torch
from PIL import Image
from RealESRGAN import RealESRGAN
from flask import Flask, request, jsonify, send_file
import io
import logging
# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = Flask(__name__)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logger.info(f'Using device: {device}')
model2 = RealESRGAN(device, scale=2)
model2.load_weights('weights/RealESRGAN_x2.pth', download=True)
logger.info('Model x2 loaded successfully')
model4 = RealESRGAN(device, scale=4)
model4.load_weights('weights/RealESRGAN_x4.pth', download=True)
logger.info('Model x4 loaded successfully')
model8 = RealESRGAN(device, scale=8)
model8.load_weights('weights/RealESRGAN_x8.pth', download=True)
logger.info('Model x8 loaded successfully')
def inference(image, size):
global model2, model4, model8
if torch.cuda.is_available():
torch.cuda.empty_cache()
logger.info('CUDA cache cleared')
logger.info(f'Starting inference with scale {size}')
try:
if size == '2x':
result = model2.predict(image.convert('RGB'))
elif size == '4x':
result = model4.predict(image.convert('RGB'))
else:
width, height = image.size
if width >= 5000 or height >= 5000:
return None, "The image is too large."
result = model8.predict(image.convert('RGB'))
logger.info(f'Inference completed for scale {size}')
except torch.cuda.OutOfMemoryError as e:
logger.error(f'OutOfMemoryError: {e}')
logger.info(f'Reloading model for scale {size}')
if size == '2x':
model2 = RealESRGAN(device, scale=2)
model2.load_weights('weights/RealESRGAN_x2.pth', download=False)
result = model2.predict(image.convert('RGB'))
elif size == '4x':
model4 = RealESRGAN(device, scale=4)
model4.load_weights('weights/RealESRGAN_x4.pth', download=False)
result = model4.predict(image.convert('RGB'))
else:
model8 = RealESRGAN(device, scale=8)
model8.load_weights('weights/RealESRGAN_x8.pth', download=False)
result = model8.predict(image.convert('RGB'))
logger.info(f'Model reloaded and inference completed for scale {size}')
return result, None
@app.route('/upscale', methods=['POST'])
def upscale():
if 'image' not in request.files:
logger.warning('No image uploaded')
return jsonify({"error": "No image uploaded"}), 400
image_file = request.files['image']
size = request.form.get('size', '2x')
try:
image = Image.open(image_file)
logger.info(f'Image uploaded and opened successfully')
except Exception as e:
logger.error(f'Invalid image file: {e}')
return jsonify({"error": "Invalid image file"}), 400
result, error = inference(image, size)
if error:
logger.error(f'Error during inference: {error}')
return jsonify({"error": error}), 400
img_io = io.BytesIO()
result.save(img_io, 'PNG')
img_io.seek(0)
logger.info('Image processing completed and ready to be sent back')
return send_file(img_io, mimetype='image/png')
if __name__ == '__main__':
logger.info('Starting the Flask server...')
app.run(host='0.0.0.0', port=5000) |