Update app.py

app.py

@@ -247,37 +247,62 @@ def create_interface(hf_token: str):
         hf_token=hf_token
     )
 
-    return gr.ChatInterface(
-        partial(chat.respond),
-        additional_inputs=[
-            gr.Slider(
-                minimum=1,
-                maximum=16343,
-                value=16343,
-                step=1,
-                label="Max new tokens"
-            ),
-            gr.Slider(
-                minimum=0.1,
-                maximum=4.0,
-                value=0.7,
-                step=0.1,
-                label="Temperature"
-            ),
-            gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=0.95,
-                step=0.05,
-                label="Top-p (nucleus sampling)"
-            ),
-        ],
-        css="""
-        @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600&display=swap');
-        body {
-            font-family: 'Inter', sans-serif;
-        }
-        """
+    # Define the inputs and outputs for the interface
+    inputs = [
+        gr.Textbox(label="User Input", lines=2),
+        gr.Slider(
+            minimum=1,
+            maximum=10,
+            step=0.5,
+            value=3,
+            label="Temperature"
+        ),
+        gr.Slider(
+            minimum=1,
+            maximum=10,
+            step=0.5,
+            value=3,
+            label="Top P"
+        ),
+        gr.Slider(
+            minimum=1,
+            maximum=3000,
+            value=2000,
+            label="Max Tokens"
+        ),
+        gr.Checkbox(label="Enable Image Generation"),
+    ]
+    outputs = [
+        gr.Chatbot(label="Assistant Response", type="messages"),
+        gr.Image(label="Generated Image", type="pil", visible=False),
+    ]
+
+    def wrap_response(message, history, temp, top_p, max_tokens, image_generation_enabled):
+        """Wrapper to handle the interface call."""
+        responses = chat.respond(
+            message,
+            history,
+            temperature=temp,
+            top_p=top_p,
+            max_tokens=max_tokens,
+        )
+
+        response_text = ""
+        image = None
+        for response in responses:
+            if isinstance(response, str):
+                response_text = response
+            elif isinstance(response, bytes):
+                image = response
+
+        return response_text, image
+
+    return gr.Interface(
+        fn=wrap_response,
+        inputs=inputs,
+        outputs=outputs,
+        live=True,
+        allow_flagging="never",
     )
 
 if __name__ == "__main__":