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