22

app.py

@@ -75,10 +75,10 @@ pipe.to("cuda", torch.bfloat16)
 # pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
 # spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
 
-def prepare(prompt):
+def prepare(prompt, is_polish_prompt):
+    if not is_polish_prompt: return prompt, False
     polished_prompt = polish_prompt(prompt)
-
-    return polished_prompt
+    return polished_prompt, True
 
 @spaces.GPU
 def inference(
@@ -179,7 +179,7 @@ with gr.Blocks(css=css) as demo:
                     placeholder="Enter your prompt",
                     container=False,
                 )
-
+                is_polish_prompt = gr.Checkbox(label="Polish prompt", value=True)
                 control_mode = gr.Radio(
                     choices=["Canny", "Depth", "HED", "MLSD", "Pose"],
                     value="HED",
@@ -242,8 +242,8 @@ with gr.Blocks(css=css) as demo:
 
     run_button.click(
         fn=prepare,
-        inputs=prompt,
-        outputs=[polished_prompt]
+        inputs=[prompt, is_polish_prompt],
+        outputs=[polished_prompt, is_polish_prompt]
         # outputs=gr.State(),  # Pass to the next function, not to UI at this step
     ).then(
         fn=inference,