mdztxi2 / app.py
Geek7's picture
Update app.py
cd230ae verified
raw
history blame
1.2 kB
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='JPEG')
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/jpeg')
if __name__ == "__main__":
app.run(debug=True)