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

app = Flask(__name__)
CORS(app)  # Enable CORS for all routes
client = InferenceClient()

@app.route('/')
def home():
    return "Welcome to the Image Background Remover!"

@app.route('/generate-image', methods=['POST'])
def generate_image():
    data = request.json  # Get the JSON data from the request
    base64_image = data['image']  # Get the base64 image string
    prompt = data['prompt']  # Get the prompt

    # Decode the base64 image
    image_data = base64.b64decode(base64_image)
    image = Image.open(io.BytesIO(image_data))

    # Generate image using the InferenceClient
    generated_image = client.image_to_image(image, prompt=prompt)

    # Save the generated image to a BytesIO object
    img_byte_arr = io.BytesIO()
    generated_image.save(img_byte_arr, format='PNG')
    img_byte_arr.seek(0)  # Move the cursor to the beginning of the BytesIO object

    # Send the generated image back to the client
    return send_file(img_byte_arr, mimetype='image/png')


if __name__ == "__main__":
    app.run(host='0.0.0.0', port=7860)  # Run directly if needed for testing