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