File size: 2,445 Bytes
3af9bb2 8be12ab 3af9bb2 21441f0 c23f4d4 8be12ab c23f4d4 ec20299 3af9bb2 c23f4d4 3af9bb2 a40654e c23f4d4 3af9bb2 c23f4d4 3af9bb2 ec20299 c23f4d4 3af9bb2 3806071 3af9bb2 3806071 3af9bb2 c23f4d4 33fcfe8 c23f4d4 3af9bb2 c23f4d4 3af9bb2 1ff4547 f9c8cf7 8be12ab ec20299 |
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 |
from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import os
from huggingface_hub import InferenceClient
from io import BytesIO # For converting image to bytes
from PIL import Image # Import Pillow for image processing
# Initialize the Flask app
myapp = Flask(__name__)
CORS(myapp) # Enable CORS for all routes
# Initialize the InferenceClient with your Hugging Face token
HF_TOKEN = os.environ.get("HF_TOKEN") # Ensure to set your Hugging Face token in the environment
client = InferenceClient(token=HF_TOKEN)
@myapp.route('/')
def home():
return "Welcome to the Image Background Remover!"
# Function to generate an image from a prompt
def generate_image(prompt, seed=1, model="prompthero/openjourney-v4"):
try:
# Generate the image using Hugging Face's inference API with the given model
result_image = client.text_to_image(prompt=prompt, seed=seed, model=model)
return result_image
except Exception as e:
print(f"Error generating image: {str(e)}")
return None
# Flask route for the API endpoint
@myapp.route('/generate_image', methods=['POST'])
def generate_api():
data = request.get_json()
# Extract required fields from the request
prompt = data.get('prompt', '')
seed = data.get('seed', 1)
model_name = data.get('model', 'prompthero/openjourney-v4') # Default model
if not prompt:
return jsonify({"error": "Prompt is required"}), 400
try:
# Call the generate_image function with the custom model name
image = generate_image(prompt, seed, model_name)
if image:
# Save the image to a BytesIO object to send as response
image_bytes = BytesIO()
image.save(image_bytes, format='PNG')
image_bytes.seek(0) # Go to the start of the byte stream
# Send the generated image as a response with a download option
return send_file(image_bytes, mimetype='image/png', as_attachment=True, download_name='generated_image.png')
else:
return jsonify({"error": "Failed to generate image"}), 500
except Exception as e:
print(f"Error in generate_api: {str(e)}") # Log the error
return jsonify({"error": str(e)}), 500
# Add this block to make sure your app runs when called
if __name__ == "__main__":
myapp.run(host='0.0.0.0', port=7860) # Run directly if needed for testing |