Update app.py

app.py

@@ -125,7 +125,7 @@ def run_showui(image, query, session_id):
 
 def save_and_upload_data(image_path, query, session_id, is_example_image, votes=None):
     """Save the data to a JSON file and upload to S3."""
-    if is_example_image == "True":  # Updated to handle string values from Dropdown
+    if is_example_image:
         return
 
     votes = votes or {"upvotes": 0, "downvotes": 0}
@@ -146,7 +146,31 @@ def save_and_upload_data(image_path, query, session_id, is_example_image, votes=
 
     return data
 
-# Examples with the `is_example` flag
+def update_vote(vote_type, session_id, is_example_image):
+    """Update the vote count and re-upload the JSON file."""
+    if is_example_image:
+        return "Example image used. No vote recorded."
+
+    local_file_name = f"{session_id}.json"
+    
+    with open(local_file_name, "r") as f:
+        data = json.load(f)
+    
+    if vote_type == "upvote":
+        data["votes"]["upvotes"] += 1
+    elif vote_type == "downvote":
+        data["votes"]["downvotes"] += 1
+    
+    with open(local_file_name, "w") as f:
+        json.dump(data, f)
+    
+    upload_to_s3(local_file_name, 'altair.storage', object_name=f"ootb/{local_file_name}")
+
+    return f"Your {vote_type} has been recorded. Thank you!"
+
+with open("./assets/showui.png", "rb") as image_file:
+    base64_image = base64.b64encode(image_file.read()).decode("utf-8")
+
 examples = [
     ["./examples/app_store.png", "Download Kindle.", True],
     ["./examples/ios_setting.png", "Turn off Do not disturb.", True],
@@ -159,10 +183,34 @@ examples = [
     ["./examples/safari_google.png", "Click on search bar.", True],
 ]
 
-def build_demo():
-    with gr.Blocks() as demo:
+def check_if_example(query):
+    """Check if the query matches an example query."""
+    print(f"Checking query: {query}")
+    for _, example_query, is_example in examples:
+        print(f"Comparing with example: {example_query}")
+        if query.strip() == example_query.strip():
+            print("Match found!")
+            return is_example
+    print("No match found.")
+    return False
+
+def build_demo(embed_mode, concurrency_count=1):
+    with gr.Blocks(title="ShowUI Demo", theme=gr.themes.Default()) as demo:
         state_image_path = gr.State(value=None)
         state_session_id = gr.State(value=None)
+        state_is_example_image = gr.State(value=False)
+
+        if not embed_mode:
+            gr.HTML(
+                f"""
+                <div style="text-align: center; margin-bottom: 20px;">
+                    <div style="display: flex; justify-content: center;">
+                        <img src="data:image/png;base64,{base64_image}" alt="ShowUI" width="320" style="margin-bottom: 10px;"/>
+                    </div>
+                    <p>ShowUI is a lightweight vision-language-action model for GUI agents.</p>
+                </div>
+                """
+            )
 
         with gr.Row():
             with gr.Column(scale=3):
@@ -174,58 +222,69 @@ def build_demo():
                 )
                 submit_btn = gr.Button(value="Submit", variant="primary")
 
-                # Examples component
                 gr.Examples(
                     examples=[[e[0], e[1]] for e in examples],
                     inputs=[imagebox, textbox],
-                    outputs=[textbox],  # Only update the query textbox
-                    examples_per_page=3,
-                )
-
-                # Add a hidden dropdown to pass the `is_example` flag
-                is_example_dropdown = gr.Dropdown(
-                    choices=["True", "False"],
-                    value="False",
-                    visible=False,
-                    label="Is Example Image",
-                )
-
-                def set_is_example(query):
-                    # Find the example and return its `is_example` flag
-                    for _, example_query, is_example in examples:
-                        if query.strip() == example_query.strip():
-                            return str(is_example)  # Return as string for Dropdown compatibility
-                    return "False"
-
-                textbox.change(
-                    set_is_example,
-                    inputs=[textbox],
-                    outputs=[is_example_dropdown],
+                    outputs=[state_is_example_image],
+                    fn=lambda img, query: check_if_example(query),
+                    examples_per_page=3
                 )
 
             with gr.Column(scale=8):
                 output_img = gr.Image(type="pil", label="Output Image")
                 output_coords = gr.Textbox(label="Clickable Coordinates")
 
-                with gr.Row(equal_height=True):
+                with gr.Row(elem_id="action-buttons", equal_height=True):
+                    upvote_btn = gr.Button(value="Looks good!", variant="secondary")
+                    downvote_btn = gr.Button(value="Too bad!", variant="secondary")
                     clear_btn = gr.Button(value="🗑️ Clear", interactive=True)
 
-            # Submit button logic
+            def on_submit(image, query, is_example_image):
+                if image is None:
+                    raise ValueError("No image provided. Please upload an image before submitting.")
+                
+                session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
+                
+                result_image, click_coords, image_path = run_showui(image, query, session_id)
+                
+                save_and_upload_data(image_path, query, session_id, is_example_image)
+                
+                return result_image, click_coords, image_path, session_id, is_example_image
+
             submit_btn.click(
-                lambda image, query, is_example: on_submit(image, query, is_example),
-                inputs=[imagebox, textbox, is_example_dropdown],
-                outputs=[output_img, output_coords, state_image_path, state_session_id],
+                on_submit,
+                [imagebox, textbox, state_is_example_image],
+                [output_img, output_coords, state_image_path, state_session_id, state_is_example_image],
             )
 
-            # Clear button logic
             clear_btn.click(
-                lambda: (None, None, None, None),
+                lambda: (None, None, None, None, None),
                 inputs=None,
-                outputs=[imagebox, textbox, output_img, output_coords, state_image_path, state_session_id],
+                outputs=[imagebox, textbox, output_img, output_coords, state_image_path, state_session_id, state_is_example_image],
+                queue=False
+            )
+
+            upvote_btn.click(
+                lambda session_id, is_example_image: update_vote("upvote", session_id, is_example_image),
+                inputs=[state_session_id, state_is_example_image],
+                outputs=[],
+                queue=False
+            )
+
+            downvote_btn.click(
+                lambda session_id, is_example_image: update_vote("downvote", session_id, is_example_image),
+                inputs=[state_session_id, state_is_example_image],
+                outputs=[],
+                queue=False
             )
 
     return demo
 
 if __name__ == "__main__":
-    demo = build_demo()
-    demo.launch()
+    demo = build_demo(embed_mode=False)
+    demo.queue(api_open=False).launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        ssr_mode=False,
+        debug=True,
+    )