File size: 2,581 Bytes
26c97ce 8be12ab 3af9bb2 7b1766e 26c97ce 8be12ab c23f4d4 ec20299 3af9bb2 c23f4d4 3af9bb2 26c97ce 3af9bb2 7b1766e c23f4d4 3af9bb2 26c97ce ec20299 c23f4d4 3af9bb2 7b1766e 3af9bb2 26c97ce 3af9bb2 26c97ce c23f4d4 7b1766e c23f4d4 7b1766e 26c97ce 33fcfe8 7b1766e 26c97ce 7b1766e 3af9bb2 c23f4d4 3af9bb2 1ff4547 26c97ce 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 64 65 66 67 68 69 |
from flask import Flask, request, jsonify, send_file, render_template
from flask_cors import CORS
import os
from huggingface_hub import InferenceClient
from io import BytesIO
from PIL import Image
# 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)
# Function to generate an image from a text prompt
def generate_image(prompt, seed=1, model="prompthero/openjourney-v4"):
try:
# Generate the image using Hugging Face's inference API
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 to generate an image based on a text prompt
@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
img_byte_arr = BytesIO()
image.save(img_byte_arr, format='PNG') # Convert the image to PNG
img_byte_arr.seek(0) # Move to the start of the byte stream
# Send the generated image as an attachment
return send_file(
img_byte_arr,
mimetype='image/png',
as_attachment=True, # Send the file as an attachment
download_name='generated_image.png' # The file name for download
)
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
# Flask route for the UI
@myapp.route('/')
def index():
return render_template('index.html') # Serve the HTML page
# 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 |