from flask import Flask, request, send_file, jsonify
from flask_cors import CORS
from huggingface_hub import InferenceClient
from PIL import Image
import io

# Initialize Flask app
app = Flask(__name__)
CORS(app)

# Your Hugging Face API Token
API_TOKEN = "your_huggingface_api_key"

# Initialize the Inference Client
client = InferenceClient(model="stabilityai/sd-x2-latent-upscaler", token=API_TOKEN)

@app.route('/upscale', methods=['POST'])
def upscale_image():
    try:
        # Retrieve the image file from the request
        image_file = request.files.get('image')
        if not image_file:
            return jsonify({"error": "No image file provided"}), 400

        # Open the image
        input_image = Image.open(image_file)

        # Use the Hugging Face Inference Client to upscale the image
        result = client.image_to_image(
            image=input_image,
            prompt="",  # Optional: Add a guiding prompt
            num_inference_steps=50,  # Adjust as needed
            guidance_scale=6.0  # Adjust as needed
        )

        # Save the upscaled image to a BytesIO object
        img_io = io.BytesIO()
        result.save(img_io, 'PNG')
        img_io.seek(0)

        return send_file(img_io, mimetype='image/png')
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# Run the app
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)