Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import spaces
-from ominicontrol import generate_image
+from ominicontrol import generate_image, vote_feedback
 import os
 
 
@@ -111,15 +111,62 @@ def gradio_interface():
                         show_progress="hidden",
                     )
 
-                btn = gr.Button("Generate Image")
+                btn = gr.Button("Generate Image", variant="primary")
+
+                with gr.Accordion("🏞️ Examples", open=True) as advanced_settings:
+                    examples = gr.Examples(
+                        examples=[
+                            ["examples/DistractedBoyfriend.webp", styles[0]],
+                            ["examples/steve.webp", styles[0]],
+                            ["examples/oiiai.png", styles[1]],
+                            ["examples/doge.jpg", styles[1]],
+                            ["examples/breakingbad.jpg", styles[2]],
+                            ["examples/PulpFiction.jpg", styles[3]],
+                        ],
+                        inputs=[original_image, style],
+                    )
 
             with gr.Column(elem_classes="outputPanel"):
-                output_images = gr.Image(
+                output_image = gr.Image(
                     type="pil",
-                    width=640,
-                    height=640,
+                    width=600,
+                    height=600,
                     label="Output Image",
+                    interactive=False,
+                    sources=None,
+                )
+                inference_id = gr.Textbox(
+                    visible=False,
+                    interactive=False,
                 )
+
+                # Feedback buttons
+                with gr.Column(visible=False) as feedback:
+                    gr.Markdown(
+                        """
+                        Your feedback improves the model! Please let us know how you feel about the generated image.
+                        """,
+                    )
+                    with gr.Row() as feedback_buttons:
+                        upvote = gr.Button("👍 I like it", variant="primary")
+                        downvote = gr.Button("👎 It looks bad")
+
+                        def feedback_func(feedback):
+                            def func(inputs):
+                                print(f"Feedback: {feedback}, Inference ID: {inputs}")
+                                vote_feedback(log_id=inputs, feedback=feedback)
+                                # Here you can add your feedback logging logic
+                                return gr.update(visible=False)
+
+                            return func
+
+                        upvote.click(feedback_func("1"), inference_id, feedback)
+                        downvote.click(feedback_func("0"), inference_id, feedback)
+
+                inference_id.change(
+                    lambda x: gr.update(visible=True), output_image, feedback
+                )
+
                 hint = gr.Markdown(
                     """
                     <div style="text-align: center; width: 100%;">
@@ -134,19 +181,7 @@ def gradio_interface():
                     hint,
                 )
 
-        with gr.Row():
-            examples = gr.Examples(
-                examples=[
-                    ["examples/DistractedBoyfriend.webp", styles[0]],
-                    ["examples/steve.webp", styles[0]],
-                    ["examples/oiiai.png", styles[1]],
-                    ["examples/doge.jpg", styles[1]],
-                    ["examples/breakingbad.jpg", styles[2]],
-                    ["examples/PulpFiction.jpg", styles[3]],
-                ],
-                inputs=[original_image, style],
-            )
-
+        # with gr.Row():
         btn.click(
             fn=infer,
             inputs=[
@@ -159,7 +194,10 @@ def gradio_interface():
                 seed,
                 steps,
             ],
-            outputs=output_images,
+            outputs=[
+                output_image,
+                inference_id,
+            ],
         )
 
     return demo
@@ -178,7 +216,7 @@ def infer(
     print(
         f"Style: {style}, Inference Mode: {inference_mode}, Image Guidance: {image_guidance}, Image Ratio: {image_ratio}, Use Random Seed: {use_random_seed}, Seed: {seed}"
     )
-    result_image = generate_image(
+    result_image, inference_id = generate_image(
         image=original_image,
         style=style,
         inference_mode=inference_mode,
@@ -188,7 +226,7 @@ def infer(
         seed=seed,
         steps=steps,
     )
-    return result_image
+    return result_image, inference_id
 
 
 if USE_ZERO_GPU:
@@ -196,4 +234,4 @@ if USE_ZERO_GPU:
 
 if __name__ == "__main__":
     demo = gradio_interface()
-    demo.launch(ssr_mode=False)
+    demo.launch(server_name="0.0.0.0")