Update app.py

app.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 import json
 import logging
 import os
@@ -5,13 +6,12 @@ import random
 import re
 import sys
 import warnings
-from dataclasses import dataclass
 
+from PIL import Image
+from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
 import gradio as gr
 import spaces
 import torch
-from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
-from PIL import Image
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from prompt_check import is_unsafe_prompt
@@ -28,8 +28,10 @@ MODEL_PATH = os.environ.get("MODEL_PATH", "Tongyi-MAI/Z-Image-Turbo")
 ENABLE_COMPILE = os.environ.get("ENABLE_COMPILE", "true").lower() == "true"
 ENABLE_WARMUP = os.environ.get("ENABLE_WARMUP", "true").lower() == "true"
 ATTENTION_BACKEND = os.environ.get("ATTENTION_BACKEND", "flash_3")
+UNSAFE_MAX_NEW_TOKEN = os.environ.get("UNSAFE_MAX_NEW_TOKEN", 10)
 DASHSCOPE_API_KEY = os.environ.get("DASHSCOPE_API_KEY")
 HF_TOKEN = os.environ.get("HF_TOKEN")
+UNSAFE_PROMPT_CHECK = os.environ.get("UNSAFE_PROMPT_CHECK")
 # =============================================================================
 
 
@@ -308,7 +310,7 @@ class APIPromptExpander(PromptExpander):
             if json_start != -1:
                 json_end = content.find("```", json_start + 7)
                 try:
-                    json_str = content[json_start + 7: json_end].strip()
+                    json_str = content[json_start + 7 : json_end].strip()
                     data = json.loads(json_str)
                     expanded_prompt = data.get("revised_prompt", content)
                 except:
@@ -377,8 +379,15 @@ def prompt_enhance(prompt, enable_enhance):
 
 @spaces.GPU
 def generate(
-    prompt, resolution="1024x1024 ( 1:1 )", seed=42, steps=9, shift=3.0, random_seed=True, gallery_images=None,
-    enhance=False, progress=gr.Progress(track_tqdm=True)
+    prompt,
+    resolution="1024x1024 ( 1:1 )",
+    seed=42,
+    steps=9,
+    shift=3.0,
+    random_seed=True,
+    gallery_images=None,
+    enhance=False,
+    progress=gr.Progress(track_tqdm=True),
 ):
     """
     Generate an image using the Z-Image model based on the provided prompt and settings.
@@ -417,9 +426,15 @@ def generate(
         if pipe is None:
             raise gr.Error("Model not loaded.")
 
-        has_nsfw_concept = is_unsafe_prompt(pipe.text_encoder, pipe.tokenizer, prompt)
-        if has_nsfw_concept:
-            raise UnsafeContentError("input unsafe")
+        has_unsafe_concept = is_unsafe_prompt(
+            pipe.text_encoder,
+            pipe.tokenizer,
+            system_prompt=UNSAFE_PROMPT_CHECK,
+            user_prompt=prompt,
+            max_new_token=UNSAFE_MAX_NEW_TOKEN,
+        )
+        if has_unsafe_concept:
+            raise UnsafeContentError("Input unsafe")
 
         final_prompt = prompt
 
@@ -443,9 +458,7 @@ def generate(
         )
 
         safety_checker_input = pipe.safety_feature_extractor([image], return_tensors="pt").pixel_values.cuda()
-        _, has_nsfw_concept = pipe.safety_checker(
-            images=[torch.zeros(1)], clip_input=safety_checker_input
-        )
+        _, has_nsfw_concept = pipe.safety_checker(images=[torch.zeros(1)], clip_input=safety_checker_input)
         has_nsfw_concept = has_nsfw_concept[0]
         if has_nsfw_concept:
             raise UnsafeContentError("input unsafe")
@@ -493,8 +506,9 @@ with gr.Blocks(title="Z-Image Demo") as demo:
                 res_cat = gr.Dropdown(value=1024, choices=choices, label="Resolution Category")
 
                 initial_res_choices = RES_CHOICES["1024"]
-                resolution = gr.Dropdown(value=initial_res_choices[0], choices=RESOLUTION_SET,
-                                         label="Width x Height (Ratio)")
+                resolution = gr.Dropdown(
+                    value=initial_res_choices[0], choices=RESOLUTION_SET, label="Width x Height (Ratio)"
+                )
 
             with gr.Row():
                 seed = gr.Number(label="Seed", value=42, precision=0)
@@ -512,12 +526,16 @@ with gr.Blocks(title="Z-Image Demo") as demo:
 
         with gr.Column(scale=1):
             output_gallery = gr.Gallery(
-                label="Generated Images", columns=2, rows=2, height=600, object_fit="contain", format="png",
-                interactive=False
+                label="Generated Images",
+                columns=2,
+                rows=2,
+                height=600,
+                object_fit="contain",
+                format="png",
+                interactive=False,
             )
             used_seed = gr.Textbox(label="Seed Used", interactive=False)
 
-
     def update_res_choices(_res_cat):
         if str(_res_cat) in RES_CHOICES:
             res_choices = RES_CHOICES[str(_res_cat)]
@@ -525,7 +543,6 @@ with gr.Blocks(title="Z-Image Demo") as demo:
             res_choices = RES_CHOICES["1024"]
         return gr.update(value=res_choices[0], choices=res_choices)
 
-
     res_cat.change(update_res_choices, inputs=res_cat, outputs=resolution, api_visibility="private")
 
     # PE enhancement button (Temporarily disabled)
@@ -542,8 +559,8 @@ with gr.Blocks(title="Z-Image Demo") as demo:
         api_visibility="public",
     )
 
-css = '''
+css = """
 .fillable{max-width: 1230px !important}
-'''
+"""
 if __name__ == "__main__":
     demo.launch(css=css, mcp_server=True)