app.py

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

# Initialize the Flask app
app = Flask(__name__)
CORS(app)  # 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)

# Hardcoded negative prompt
HARDCODED_NEGATIVE_PROMPT = """
ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, 
poorly drawn hands, poorly drawn face, mutation, deformed, blurry, bad anatomy, 
bad proportions, extra limbs, cloned face, disfigured, gross proportions, 
malformed limbs, missing arms, missing legs, extra arms, extra legs, 
mutated hands, fused fingers, too many fingers, long neck
"""

@app.route('/')
def home():
    return "Welcome to the Image Background Remover!"

# Simple content moderation function
def is_prompt_explicit(prompt):
    explicit_keywords = ["sexual", "nudity", "erotic", "explicit", "porn", "pornographic", "xxx", "hentai", "fetish", "sex", "sensual", "nude", "strip", "stripping", "adult", "lewd", "provocative", "obscene", "vulgar", "intimacy", "intimate", "lust", "arouse", "seductive", "seduction", "kinky", "bdsm", "dominatrix", "bondage", "hardcore", "softcore", "topless", "bottomless", "threesome", "orgy", "incest", "taboo", "masturbation", "genital", "penis", "vagina", "breast", "boob", "nipple", "butt", "anal", "oral", "ejaculation", "climax", "moan", "foreplay", "intercourse", "naked", "exposed", "suicide", "self-harm", "overdose", "poison", "hang", "end life", "kill myself", "noose", "depression", "hopeless", "worthless", "die", "death", "harm myself"]  # Add more keywords as needed
    for keyword in explicit_keywords:
        if keyword.lower() in prompt.lower():
            return True
    return False

# Function to generate an image from a text prompt
def generate_image(prompt, negative_prompt=None, height=512, width=512, model="stabilityai/stable-diffusion-2-1", num_inference_steps=50, guidance_scale=7.5, seed=None):
    try:
        # Generate the image using Hugging Face's inference API with additional parameters
        image = client.text_to_image(
            prompt=prompt, 
            negative_prompt=HARDCODED_NEGATIVE_PROMPT,
            height=height, 
            width=width, 
            model=model,
            num_inference_steps=num_inference_steps,  # Control the number of inference steps
            guidance_scale=guidance_scale,  # Control the guidance scale
            seed=seed  # Control the seed for reproducibility
        )
        return image  # Return the generated image
    except Exception as e:
        print(f"Error generating image: {str(e)}")
        return None

# Flask route for the API endpoint to generate an image
@app.route('/generate_image', methods=['POST'])
def generate_api():
    data = request.get_json()

    # Extract required fields from the request
    prompt = data.get('prompt', '')
    negative_prompt = data.get('negative_prompt', None)
    height = data.get('height', 1024)  # Default height
    width = data.get('width', 720)  # Default width
    num_inference_steps = data.get('num_inference_steps', 50)  # Default number of inference steps
    guidance_scale = data.get('guidance_scale', 7.5)  # Default guidance scale
    model_name = data.get('model', 'stabilityai/stable-diffusion-2-1')  # Default model
    seed = data.get('seed', None)  # Seed for reproducibility, default is None

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

    try:
        # Check for explicit content
        if is_prompt_explicit(prompt):
            # Return the pre-defined "thinkgood.png" image
            return send_file(
                "thinkgood.jpeg",
                mimetype='image/png',
                as_attachment=False,
                download_name='thinkgood.png'
            )

        # Call the generate_image function with the provided parameters
        image = generate_image(prompt, negative_prompt, height, width, model_name, num_inference_steps, guidance_scale, seed)

        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 a response
            return send_file(
                img_byte_arr, 
                mimetype='image/png', 
                as_attachment=False,  # 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__":
    app.run(host='0.0.0.0', port=7860)  # Run directly if needed for testing