diff --git a/src/custom_widgets/chat_widget.py b/src/custom_widgets/chat_widget.py
index 1a4d6f0..dc31bd5 100644
--- a/src/custom_widgets/chat_widget.py
+++ b/src/custom_widgets/chat_widget.py
@@ -70,6 +70,7 @@ class chat(Gtk.ScrolledWindow):
         )
         self.messages = {}
         self.welcome_screen = None
+        self.regenerate_button = None
         self.busy = False
 
     def stop_message(self):
@@ -180,7 +181,19 @@ class chat(Gtk.ScrolledWindow):
                     messages_dict[message_id]['files'] = files
         return messages_dict
 
-
+    def show_regenerate_button(self, msg:message):
+        if self.regenerate_button:
+            self.remove(self.regenerate_button)
+        self.regenerate_button = Gtk.Button(
+            child=Adw.ButtonContent(
+                icon_name='update-symbolic',
+                label=_('Regenerate Response')
+            ),
+            css_classes=["suggested-action"],
+            halign=3
+        )
+        self.regenerate_button.connect('clicked', lambda *_: msg.action_buttons.regenerate_message())
+        self.container.append(self.regenerate_button)
 
 class chat_tab(Gtk.ListBoxRow):
     __gtype_name__ = 'AlpacaChatTab'
diff --git a/src/custom_widgets/message_widget.py b/src/custom_widgets/message_widget.py
index 423f02e..ec7dacd 100644
--- a/src/custom_widgets/message_widget.py
+++ b/src/custom_widgets/message_widget.py
@@ -342,12 +342,13 @@ class action_buttons(Gtk.Box):
         message_element = self.get_parent()
         if not chat.busy:
             message_element.set_text()
-            message_element.container.remove(message_element.footer)
+            if message_element.footer:
+                message_element.container.remove(message_element.footer)
             message_element.remove_overlay(self)
             message_element.action_buttons = None
             history = window.convert_history_to_ollama(chat)[:list(chat.messages).index(message_element.message_id)]
             data = {
-                "model": window.get_current_model(1),
+                "model": window.model_selector.get_model(),
                 "messages": history,
                 "options": {"temperature": window.model_tweaks["temperature"], "seed": window.model_tweaks["seed"]},
                 "keep_alive": f"{window.model_tweaks['keep_alive']}m"
diff --git a/src/window.py b/src/window.py
index f56c0b2..34de4b7 100644
--- a/src/window.py
+++ b/src/window.py
@@ -591,34 +591,27 @@ Generate a title following these rules:
         self.stop_button.set_visible(not send)
         self.send_button.set_visible(send)
 
-    def run_message(self, data, message_element:message_widget.message):
+    def run_message(self, data:dict, message_element:message_widget.message):
         logger.debug("Running message")
         chat = message_element.get_parent().get_parent().get_parent().get_parent()
         chat.busy = True
         self.chat_list_box.set_sensitive(False)
         if chat.welcome_screen:
             chat.welcome_screen.set_visible(False)
+        if chat.regenerate_button:
+            chat.container.remove(chat.regenerate_button)
         self.switch_send_stop_button(False)
         if self.regenerate_button:
             GLib.idle_add(self.chat_list_box.get_current_chat().remove, self.regenerate_button)
         try:
-            print(data)
             response = connection_handler.stream_post(f"{connection_handler.URL}/api/chat", data=json.dumps(data), callback=lambda data, message_element=message_element: GLib.idle_add(message_element.update_message, data))
             if response.status_code != 200:
-                print(response)
                 raise Exception('Network Error')
         except Exception as e:
+            chat.busy = False
+            GLib.idle_add(message_element.add_action_buttons)
+            GLib.idle_add(chat.show_regenerate_button, message_element)
             GLib.idle_add(self.connection_error)
-            self.regenerate_button = Gtk.Button(
-                child=Adw.ButtonContent(
-                    icon_name='update-symbolic',
-                    label=_('Regenerate Response')
-                ),
-                css_classes=["suggested-action"],
-                halign=3
-            )
-            #GLib.idle_add(self.chat_list_box.get_current_chat().append, self.regenerate_button)
-            #self.regenerate_button.connect('clicked', lambda button, message_id=message_id, bot_message_box=self.bot_message_box, bot_message_button_container=self.bot_message_button_container : self.regenerate_message(message_id, bot_message_box, bot_message_button_container))
 
     def pull_model_update(self, data, model_name):
         if 'error' in data: