Update app.py

app.py

@@ -14,8 +14,6 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from prompt_check import is_unsafe_prompt
-
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 
 from diffusers import ZImagePipeline
@@ -28,10 +26,8 @@ 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 = int(os.environ.get("UNSAFE_MAX_NEW_TOKEN", "10"))
-DASHSCOPE_API_KEY = os.environ.get("DASHSCOPE_API_KEY")
+OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
 HF_TOKEN = os.environ.get("HF_TOKEN")
-UNSAFE_PROMPT_CHECK = os.environ.get("UNSAFE_PROMPT_CHECK")
 # =============================================================================
 
 
@@ -280,11 +276,11 @@ class APIPromptExpander(PromptExpander):
         try:
             from openai import OpenAI
 
-            api_key = self.api_config.get("api_key") or DASHSCOPE_API_KEY
-            base_url = self.api_config.get("base_url", "https://dashscope.aliyuncs.com/compatible-mode/v1")
+            api_key = self.api_config.get("api_key") or OPENROUTER_API_KEY
+            base_url = self.api_config.get("base_url", "https://openrouter.ai/api/v1")
 
             if not api_key:
-                print("Warning: DASHSCOPE_API_KEY not found.")
+                print("Warning: OPENROUTER_API_KEY not found.")
                 return None
 
             return OpenAI(api_key=api_key, base_url=base_url)
@@ -310,12 +306,10 @@ class APIPromptExpander(PromptExpander):
             prompt = " "
 
         try:
-            model = self.api_config.get("model", "qwen3-max-preview")
+            model = self.api_config.get("model", "google/gemini-2.5-flash")
             response = self.client.chat.completions.create(
                 model=model,
                 messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}],
-                temperature=0.7,
-                top_p=0.8,
             )
 
             content = response.choices[0].message.content
@@ -331,6 +325,8 @@ class APIPromptExpander(PromptExpander):
             else:
                 expanded_prompt = content
 
+            print(f"Original prompt: {prompt}\nFinal prompt: {expanded_prompt}")
+
             return PromptOutput(
                 status=True, prompt=expanded_prompt, seed=seed, system_prompt=system_prompt, message=content
             )
@@ -366,7 +362,7 @@ def init_app():
         pipe = None
 
     try:
-        prompt_expander = create_prompt_expander(backend="api", api_config={"model": "qwen3-max-preview"})
+        prompt_expander = create_prompt_expander(backend="api", api_config={"model": "google/gemini-2.5-flash"})
         print("Prompt expander initialized.")
     except Exception as e:
         print(f"Error initializing prompt expander: {e}")
@@ -432,52 +428,29 @@ def generate(
     else:
         new_seed = seed if seed != -1 else random.randint(1, 1000000)
 
-    class UnsafeContentError(Exception):
-        pass
-
-    try:
-        if pipe is None:
-            raise gr.Error("Model not loaded.")
-
-        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
 
-        if enhance:
-            final_prompt, _ = prompt_enhance(prompt, True)
-            print(f"Enhanced prompt: {final_prompt}")
+    if pipe is None:
+        raise gr.Error("Model not loaded.")
 
-        try:
-            resolution_str = resolution.split(" ")[0]
-        except:
-            resolution_str = "1024x1024"
-
-        image = generate_image(
-            pipe=pipe,
-            prompt=final_prompt,
-            resolution=resolution_str,
-            seed=new_seed,
-            guidance_scale=0.0,
-            num_inference_steps=int(steps + 1),
-            shift=shift,
-        )
+    final_prompt = prompt
 
-        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 = has_nsfw_concept[0]
-        if has_nsfw_concept:
-            raise UnsafeContentError("input unsafe")
+    if enhance:
+        final_prompt, msg = prompt_enhance(prompt, True)
 
-    except UnsafeContentError:
-        image = Image.open("nsfw.png")
+    try:
+        resolution_str = resolution.split(" ")[0]
+    except:
+        resolution_str = "1024x1024"
+
+    image = generate_image(
+        pipe=pipe,
+        prompt=final_prompt,
+        resolution=resolution_str,
+        seed=new_seed,
+        guidance_scale=0.0,
+        num_inference_steps=int(steps + 1),
+        shift=shift,
+    )
 
     if gallery_images is None:
         gallery_images = []
@@ -491,8 +464,8 @@ init_app()
 
 # ==================== AoTI (Ahead of Time Inductor compilation) ====================
 
-pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
-spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
+#pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
+#spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
 
 with gr.Blocks(title="Z-Image Demo") as demo:
     gr.Markdown(