22

app.py

@@ -79,6 +79,9 @@ pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
 spaces.aoti_blocks_load(pipe.transformer.layers,
                         "zerogpu-aoti/Z-Image", variant="fa3")
 
+def prepare(prompt, input_image):
+    polished_prompt = polish_prompt(prompt)
+    return polished_prompt
 
 @spaces.GPU
 def inference(
@@ -161,9 +164,10 @@ with gr.Blocks() as demo:
                 run_button = gr.Button("Run", variant="primary")
             with gr.Column():
                 output_image = gr.Image(label="Generated image", show_label=False)
-                polished_prompt = gr.Textbox(label="Polished prompt", interactive=False)
-                with gr.Accordion("Control image", open=False):
+                
+                with gr.Accordion("Preprocessor output", open=False):
                     control_image = gr.Image(label="Control image", show_label=False)
+                    polished_prompt = gr.Textbox(label="Polished prompt", interactive=False)
 
         with gr.Accordion("Advanced Settings", open=False):
                 seed = gr.Slider(
@@ -211,23 +215,33 @@ with gr.Blocks() as demo:
         gr.Examples(examples=examples, inputs=[input_image, prompt])
         
         gr.HTML(read_file("static/footer.html"))
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=inference,
-        inputs=[
-            prompt,
-            input_image,
-            image_scale,
-            control_context_scale,
-            seed,
-            randomize_seed,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[output_image, seed],
+    run_button.click(
+        fn=prepare,
+        inputs=prompt,
+        outputs=[polished_prompt]
+        # outputs=gr.State(),  # Pass to the next function, not to UI at this step
     ).then(
-        
+        # fn=generate_image,
+        # inputs=None,   # This will automatically use the previous result
+        # outputs=output
     )
+    # gr.on(
+    #     triggers=[run_button.click, prompt.submit],
+    #     fn=inference,
+    #     inputs=[
+    #         prompt,
+    #         input_image,
+    #         image_scale,
+    #         control_context_scale,
+    #         seed,
+    #         randomize_seed,
+    #         guidance_scale,
+    #         num_inference_steps,
+    #     ],
+    #     outputs=[output_image, seed],
+    # ).then(
+        
+    # )
 
 if __name__ == "__main__":
     demo.launch(mcp_server=True, css=css)