Update app.py

app.py

@@ -625,34 +625,39 @@ with gr.Blocks() as iface:
     initial_state = gr.State()
     
     def create_buttons(button_texts):
-        with button_container:
-            button_container.clear()
-            buttons = [gr.Button(text) for text in button_texts]
+        buttons = []
+        for text in button_texts:
+            button = gr.Button(text)
+            button.click(
+                show_details,
+                inputs=[button, output, initial_state],
+                outputs=[output, back_button, initial_state]
+            )
+            buttons.append(button)
         return buttons
 
+    def update_ui(explanation, image, button_texts, show_back, state):
+        button_container.clear()
+        buttons = create_buttons(button_texts)
+        return explanation, image, *buttons, gr.update(visible=show_back), state
+
     input_image.change(
         predict,
         inputs=input_image,
         outputs=[output, output_image, button_container, back_button, initial_state]
-    )
-    
-    def on_button_click(button_text, output, state):
-        return show_details(button_text, output, state)
-    
-    button_container.select(
-        on_button_click,
-        inputs=["selected", output, initial_state],
-        outputs=[output, back_button, initial_state]
+    ).then(
+        update_ui,
+        inputs=[output, output_image, button_container, back_button, initial_state],
+        outputs=[output, output_image] + [button_container] * 9 + [back_button, initial_state]
     )
     
     def go_back(state):
-        buttons = create_buttons(state["buttons"])
-        return state["explanation"], gr.update(visible=state["show_back"]), state
+        return update_ui(state["explanation"], None, state["buttons"], state["show_back"], state)
 
     back_button.click(
         go_back,
         inputs=[initial_state],
-        outputs=[output, back_button, initial_state]
+        outputs=[output, output_image] + [button_container] * 9 + [back_button, initial_state]
     )
     
     gr.Examples(