Try to reinstall after a failed upgrade

* After failed upgrade the UI would try to start the old backup server, which likely doesn't work
* When restoring a backup, also restore the old version file

Author: Acly

ai_diffusion/server.py

@@ -144,6 +144,10 @@ def check_install(self):
             self.state = ServerState.missing_resources
         elif update_required:
             self.state = ServerState.update_required
+        elif self.version == "incomplete":
+            # At this point everything seems to be there, but the version file indicates
+            # that installation did not complete successfully. Best bet is to try and re-install.
+            self.state = ServerState.update_required
         else:
             self.state = ServerState.stopped
         self.missing_resources += missing_sd15 + missing_sdxl + missing_flux
@@ -402,7 +406,8 @@ def info(message: str):
             log.info(message)
             callback(InstallationProgress("Upgrading", message=message))
 
-        info(f"Starting upgrade from {self.version} to {resources.version}")
+        old_version = self.version
+        info(f"Starting upgrade from {old_version} to {resources.version}")
         comfy_dir = self.comfy_dir
         upgrade_dir = self.path / f"upgrade-{resources.version}"
         upgrade_comfy_dir = upgrade_dir / "ComfyUI"
@@ -430,6 +435,7 @@ def info(message: str):
                 log.warning(f"Error during upgrade: {str(e)} - Restoring {upgrade_comfy_dir}")
                 safe_remove_dir(comfy_dir)
                 shutil.move(upgrade_comfy_dir, comfy_dir)
+                self._version_file.write_text(f"{old_version} {self.backend.name}")
             raise e
 
         try:

ai_diffusion/ui/server.py

@@ -275,13 +275,11 @@ def _package_name(self, package: str | ModelResource | CustomNode):
 
 
 class ServerWidget(QWidget):
-    _server: Server
-    _error = ""
-    _packages: dict[str, PackageGroupWidget]
-
     def __init__(self, srv: Server, parent=None):
         super().__init__(parent)
         self._server = srv
+        self._error = ""
+        self._packages: dict[str, PackageGroupWidget]
 
         layout = QVBoxLayout(self)
 
@@ -755,9 +753,10 @@ def update_ui(self):
             self._launch_button.setEnabled(False)
             self._manage_button.setEnabled(False)
         elif state is ServerState.update_required:
-            self._status_label.setText(
-                _("Upgrade required") + f": v{self._server.version} -> v{resources.version}"
-            )
+            text = _("Upgrade required") + f": v{self._server.version} -> v{resources.version}"
+            if self._server.version == "incomplete":
+                text = _("Previous installation is incomplete")
+            self._status_label.setText(text)
             self._status_label.setStyleSheet(f"color:{yellow};font-weight:bold")
             self._launch_button.setText(_("Upgrade"))
         elif state is ServerState.stopped: