import os
from flask import Flask, request, jsonify, send_file
from transformers import pipeline

app = Flask(__name__)

# Load the Hugging Face pipeline for image generation
generator = pipeline('text-to-image', model='CompVis/stable-diffusion-v1-4', device=0)  # Use device=0 if using GPU

@app.route('/generate_image', methods=['POST'])
def generate_image():
    # Get the description from the request
    description = request.json.get('description')
    
    if not description:
        return jsonify({"error": "Description is required"}), 400
    
    # Generate the image using the description
    image = generator(description)[0]["sample"]  # This will return the generated image
    
    # Save the image to a file
    output_path = "generated_image.png"
    image.save(output_path)
    
    # Return the image file
    return send_file(output_path, mimetype='image/png')

if __name__ == '__main__':
    app.run(debug=True)