from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import os
from huggingface_hub import InferenceClient
from io import BytesIO
from PIL import Image  # Importing 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)

# Function to generate an image from a prompt using the specified model
def generate_image(prompt, seed=1, model="prompthero/openjourney-v4"):  # Default model if none provided
    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
@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')  # Use the provided model name or a default one

    if not prompt:
        return jsonify({"error": "Prompt is required"}), 400

    try:
        # Call the generate_image function with the dynamically provided model name
        image = generate_image(prompt, seed, model_name)

        if image:
            # Save the image to a BytesIO object
            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 an attachment
            return send_file(
                image_bytes, 
                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

# 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