From 770d8cee5f8eaa2765a2fd45f3f4faf400b78ab7 Mon Sep 17 00:00:00 2001
From: Acly <aclysia@gmail.com>
Date: Thu, 30 Nov 2023 01:44:28 +0100
Subject: [PATCH] Refuse to upgrade or install into existing directories which
 exist and are not empty #152 This includes directories with existing Comfy
 that was not previously installed by the plugin

---
 ai_diffusion/__init__.py  |  2 +-
 ai_diffusion/server.py    | 33 ++++++++++++++-------------------
 ai_diffusion/settings.py  | 12 +++++-------
 ai_diffusion/ui/server.py | 10 +++++++++-
 4 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/ai_diffusion/__init__.py b/ai_diffusion/__init__.py
index b0fa58ae5..9dec87f52 100644
--- a/ai_diffusion/__init__.py
+++ b/ai_diffusion/__init__.py
@@ -1,6 +1,6 @@
 """Generative AI plugin for Krita using Stable Diffusion"""
 
-__version__ = "1.8.1"
+__version__ = "1.8.2"
 
 import importlib.util
 
diff --git a/ai_diffusion/server.py b/ai_diffusion/server.py
index b31a742b0..41e7698da 100644
--- a/ai_diffusion/server.py
+++ b/ai_diffusion/server.py
@@ -74,22 +74,10 @@ def check_install(self):
             self.version = None
 
         comfy_pkg = ["main.py", "nodes.py", "custom_nodes"]
-        self.comfy_dir = _find_component(comfy_pkg, [self.path, self.path / "ComfyUI"])
+        self.comfy_dir = _find_component(comfy_pkg, [self.path / "ComfyUI"])
 
         python_pkg = ["python3.dll", "python.exe"] if is_windows else ["python3", "pip3"]
-        python_search_paths = [
-            self.path / "python",
-            self.path / "venv" / "bin",
-            self.path / ".venv" / "bin",
-        ]
-        if self.comfy_dir:
-            python_search_paths += [
-                self.comfy_dir / "python",
-                self.comfy_dir / "venv" / "bin",
-                self.comfy_dir / ".venv" / "bin",
-                self.comfy_dir.parent / "python",
-                self.comfy_dir.parent / "python_embeded",
-            ]
+        python_search_paths = [self.path / "python", self.path / "venv" / "bin"]
         python_path = _find_component(python_pkg, python_search_paths)
         if python_path is None:
             self._python_cmd = _find_program("python3", "python")
@@ -341,9 +329,10 @@ def info(message: str):
             self.comfy_dir = None
             await self.install(callback)
         except Exception as e:
-            log.warning(f"Error during upgrade: {str(e)} - Restoring {upgrade_comfy_dir}")
-            shutil.rmtree(comfy_dir, ignore_errors=True)
-            shutil.move(upgrade_comfy_dir, comfy_dir)
+            if upgrade_comfy_dir.exists():
+                log.warning(f"Error during upgrade: {str(e)} - Restoring {upgrade_comfy_dir}")
+                safe_remove_dir(comfy_dir)
+                shutil.move(upgrade_comfy_dir, comfy_dir)
             raise e
 
         try:
@@ -472,10 +461,16 @@ def is_installed(self, package: str | ModelResource | CustomNode):
     def all_installed(self, packages: list[str] | list[ModelResource] | list[CustomNode]):
         return all(self.is_installed(p) for p in packages)
 
+    @property
+    def can_install(self):
+        return not self.path.exists() or (self.path.is_dir() and not any(self.path.iterdir()))
+
     @property
     def upgrade_required(self):
-        return self.state is not ServerState.not_installed and (
-            self.version is None or self.version != resources.version
+        return (
+            self.state is not ServerState.not_installed
+            and self.version is not None
+            and self.version != resources.version
         )
 
 
diff --git a/ai_diffusion/settings.py b/ai_diffusion/settings.py
index 895a2a176..a1ececc97 100644
--- a/ai_diffusion/settings.py
+++ b/ai_diffusion/settings.py
@@ -8,6 +8,7 @@
 
 from . import util
 
+
 class ServerMode(Enum):
     undefined = -1
     managed = 0
@@ -31,6 +32,7 @@ def default():
         else:
             return ServerBackend.cuda
 
+
 class PerformancePreset(Enum):
     auto = "Automatic"
     cpu = "CPU"
@@ -71,10 +73,8 @@ class Settings(QObject):
     _server_path = Setting(
         "Server Path",
         str(Path(__file__).parent / ".server"),
-        (
-            "Directory where ComfyUI is installed. At least 10GB of free disk space is required"
-            " for a full installation."
-        ),
+        "Directory where ComfyUI will be installed. At least 10GB of free disk space is required"
+        " for a full installation.",
     )
 
     server_url: str
@@ -126,9 +126,7 @@ class Settings(QObject):
     )
 
     show_control_end: bool
-    _show_control_end = Setting(
-        "Control ending step", False, "Show control ending step ratio"
-    )
+    _show_control_end = Setting("Control ending step", False, "Show control ending step ratio")
 
     history_size: int
     _history_size = Setting(
diff --git a/ai_diffusion/ui/server.py b/ai_diffusion/ui/server.py
index 0fd3ad7d7..8cb920141 100644
--- a/ai_diffusion/ui/server.py
+++ b/ai_diffusion/ui/server.py
@@ -562,7 +562,15 @@ def update(self):
 
         if self.requires_install:
             self._launch_button.setText("Install")
-            self._launch_button.setEnabled(True)
+            if not self._server.version and not self._server.can_install:
+                self._status_label.setText(
+                    "Invalid location: directory is not empty, but no previous installation was"
+                    " found"
+                )
+                self._status_label.setStyleSheet(f"color:{red};font-weight:bold")
+                self._launch_button.setEnabled(False)
+            else:
+                self._launch_button.setEnabled(True)
 
         if self._error:
             self._status_label.setText(f"<b>Error:</b> {self._error}")