Update app.py

app.py

@@ -12,12 +12,12 @@ else:
     torch_dtype = torch.float32
 
 
-def set_client_for_session(request: gr.Request):
-    x_ip_token = request.headers['x-ip-token']
+# def set_client_for_session(request: gr.Request):
+#     x_ip_token = request.headers['x-ip-token']
 
-    # The "gradio/text-to-image" space is a ZeroGPU space
-    # return Client("stzhao/LeX-Enhancer", headers={"X-IP-Token": x_ip_token})
-    return Client("stzhao/LeX-Enhancer")
+#     # The "gradio/text-to-image" space is a ZeroGPU space
+#     # return Client("stzhao/LeX-Enhancer", headers={"X-IP-Token": x_ip_token})
+#     return Client("stzhao/LeX-Enhancer")
 
 # Load models
 def load_models():    
@@ -76,12 +76,13 @@ def generate_image(enhanced_caption, seed, num_inference_steps, guidance_scale):
     return image
 
 # @spaces.GPU(duration=130)
-def run_pipeline(image_caption, text_caption, seed, num_inference_steps, guidance_scale, enable_enhancer, client):
+def run_pipeline(image_caption, text_caption, seed, num_inference_steps, guidance_scale, enable_enhancer):
     """Run the complete pipeline from captions to final image"""
     combined_caption = f"{image_caption}, with the text on it: {text_caption}."
     
     if enable_enhancer:
         # combined_caption, enhanced_caption = generate_enhanced_caption(image_caption, text_caption)
+        client = Client("stzhao/LeX-Enhancer")
         combined_caption, enhanced_caption = prompt_enhance(client, image_caption, text_caption)
         print(f"enhanced caption:\n{enhanced_caption}")
     else:
@@ -93,7 +94,7 @@ def run_pipeline(image_caption, text_caption, seed, num_inference_steps, guidanc
 
 # Gradio interface
 with gr.Blocks() as demo:
-    client = gr.State()
+    # client = gr.State()
     gr.Markdown("# LeX-Enhancer & LeX-FLUX Demo")
     gr.Markdown("## Project Page: https://zhaoshitian.github.io/lexart/")
     gr.Markdown("Generate enhanced captions from simple image and text descriptions, then create images with LeX-FLUX")
@@ -179,11 +180,11 @@ with gr.Blocks() as demo:
     
     submit_btn.click(
         fn=run_pipeline,
-        inputs=[image_caption, text_caption, seed, num_inference_steps, guidance_scale, enable_enhancer, client],
+        inputs=[image_caption, text_caption, seed, num_inference_steps, guidance_scale, enable_enhancer],
         outputs=[output_image, combined_caption_box, enhanced_caption_box]
     )
 
-    demo.load(set_client_for_session, None, client)
+    # demo.load(set_client_for_session, None, client)
 
 if __name__ == "__main__":
     demo.launch(debug=True)