import torch
from diffusers import StableDiffusionPipeline
import gradio as gr

# NSFW keyword filter
NSFW_KEYWORDS = [
    "nude", "naked", "sex", "porn", "erotic", "nsfw",
    "boobs", "breasts", "genital", "vagina", "penis"
]

# Load pipeline on CPU
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    use_safetensors=True
)
pipe = pipe.to("cpu")
pipe.safety_checker = lambda images, clip_input: (images, False)  # Optional: you can remove this if you want the default safety checker

# Function to check and generate image
def generate_image(prompt):
    if any(word in prompt.lower() for word in NSFW_KEYWORDS):
        raise gr.Error("❌ NSFW content detected in prompt. Please try a different one.")
    image = pipe(prompt).images[0]
    return image

# Gradio interface
demo = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(lines=1, label="Enter your prompt"),
    outputs=gr.Image(label="Generated Image"),
    title="🎨 Text-to-Image Generation",
    description="Enter a clean prompt to generate images using Stable Diffusion v1.5"
)

demo.queue()
demo.launch()