app.py

from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import os
import subprocess
from huggingface_hub import InferenceClient
from io import BytesIO
from PIL import Image
import re

# 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
NEGATIVE_PROMPT_FINGERS = """2D,missing fingers, extra fingers, elongated fingers, fused fingers, 
mutated fingers, poorly drawn fingers, disfigured fingers, 
too many fingers, deformed hands, extra hands, malformed hands, 
blurry hands, disproportionate fingers"""

# Define the list of explicit keywords

EXPLICIT_KEYWORDS = [
    "sexual", "sex", "boobs", "boob", "breasts", "cleavage", "penis", "phallus", "porn", "pornography", "hentai", 
    "fetish", "nude", "nudity", "provocative", "obscene", "vulgar", "intimate", "kinky", "hardcore", 
    "threesome", "orgy", "masturbation", "masturbate", "genital", "genitals", "vagina", "vaginal", 
    "anus", "anal", "butt", "buttocks", "butthole", "ass", "prostate", "erection", "cum", "ejaculation", 
    "sperm", "semen", "naked", "bare", "lingerie", "thong", "striptease", "stripper", 
    "seductive", "sensual", "explicit", "lewd", "taboo", "NSFW", "bdsm", "dominatrix", "submission", 
    "intercourse", "penetration", "orgasm", "fuck", "fucking", "fuckers", "fucker", "slut", "whore", 
    "prostitute", "hooker", "escort", "camgirl", "camwhore", "sugar daddy", "sugar baby", "adult content", 
    "sexually explicit", "arousal", "lust", "depraved", "hardcore porn", "softcore", "erotic", "erotica",
    "roleplay", "incest", "taboo", "voyeur", "exhibitionist", "peeping", "dildo", "sex toy", "vibrator", 
    "suicide", "self-harm", "depression", "kill myself", "worthless", "abuse", "violence", "rape", 
    "sexual violence", "molestation", "pedophilia", "child porn", "underage", "illegal content"
]

# Function to scan the entire prompt for explicit keywords
def scan_prompt(prompt, keywords):
    # Create a regex pattern to match any keyword (case insensitive)
    pattern = r'\b(?:' + '|'.join(re.escape(keyword) for keyword in keywords) + r')\b'
    # Find all matches in the prompt
    matches = re.findall(pattern, prompt, flags=re.IGNORECASE)
    # Return True if matches are found, and the list of matched keywords
    return bool(matches), matches

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

# 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=NEGATIVE_PROMPT_FINGERS,
            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 using scan_prompt
        is_nsfw, found_keywords = scan_prompt(prompt, EXPLICIT_KEYWORDS)
        if is_nsfw:
            print(f"Explicit keywords found: {found_keywords}")
            # Return the pre-defined "nsfw.jpg" image
            return send_file(
                "nsfw.jpg",
                mimetype='image/png',
                as_attachment=False,
                download_name='nsfw_detected.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__":
    subprocess.Popen(["python", "wk.py"])  # Start awake.py

    app.run(host='0.0.0.0', port=7860)  # Run directly if needed for testing