import spaces
import gradio as gr
import torch
from diffusers import AutoPipelineForText2Image
from transformers import CLIPImageProcessor
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker

# Load safety checker & feature extractor
safety_checker = StableDiffusionSafetyChecker.from_pretrained(
    "CompVis/stable-diffusion-safety-checker",
    torch_dtype=torch.bfloat16
).to("cuda")

feature_extractor = CLIPImageProcessor.from_pretrained(
    "openai/clip-vit-base-patch32"
)

# Load Z-Image-Turbo pipeline
pipe = AutoPipelineForText2Image.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=torch.bfloat16
).to("cuda")

# Blacklist of prohibited keywords for prompt pre-filtering
BLOCKED_WORDS = [
    "nsfw", "nude", "nudity", "naked", "porn", "xxx", "erotic",
    "gore", "blood", "bloody", "kill", "murder", "corpse",
    "violence", "decapitation", "torture", "suicide", "mutilation",
    "dead body", "severed", "slaughter"
]

def check_prompt_safety(prompt: str) -> bool:
    prompt_lower = prompt.lower()
    return any(bad_word in prompt_lower for bad_word in BLOCKED_WORDS)

@spaces.GPU(duration=25)
def generate_image(prompt: str, seed: int, height: int, width: int):
    # Layer 1: Prompt Blacklist Filter
    if check_prompt_safety(prompt):
        raise gr.Error("Safety Alert: Prompt contains restricted NSFW/NSFL keywords. Generation blocked.")

    generator = torch.Generator("cuda").manual_seed(int(seed))

    # 8-step Turbo inference
    output = pipe(
        prompt=prompt,
        height=height,
        width=width,
        num_inference_steps=8,
        guidance_scale=0.0,
        generator=generator
    )

    image = output.images[0]

    # Layer 2: Output Image Safety Filter
    safety_inputs = feature_extractor(images=image, return_tensors="pt").to("cuda")
    _, has_nsfw = safety_checker(
        images=[image],
        clip_input=safety_inputs.pixel_values.to(torch.bfloat16)
    )

    if has_nsfw[0]:
        raise gr.Error("Safety Alert: The generated image violated safety policies. Blocked.")

    return image

with gr.Blocks(title="Z-Image-Turbo Safe Generator") as demo:
    gr.Markdown("# 🚀 Z-Image-Turbo Safe Generator")
    gr.Markdown("Fast 8-step generation with built-in dual-layer NSFW/NSFL content moderation.")

    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="A futuristic cyberpunk street in Tokyo at sunset, 8k resolution..."
            )
            seed_input = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=2147483647,
                step=1,
                value=42
            )
            with gr.Row():
                height_input = gr.Dropdown(
                    label="Height",
                    choices=[512, 768, 1024],
                    value=1024
                )
                width_input = gr.Dropdown(
                    label="Width",
                    choices=[512, 768, 1024],
                    value=1024
                )
            generate_btn = gr.Button("Generate Image", variant="primary")
        with gr.Column():
            image_output = gr.Image(label="Result", type="pil")

    generate_btn.click(
        fn=generate_image,
        inputs=[prompt_input, seed_input, height_input, width_input],
        outputs=image_output
    )

if __name__ == "__main__":
    demo.launch()