Spaces:
Running on Zero
Running on Zero
File size: 3,490 Bytes
616fef0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | 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() |