Show warning message when a LoRA doesn't match the base model

Author: Acly

ai_diffusion/__init__.py

@@ -1,6 +1,6 @@
 """Generative AI plugin for Krita"""
 
-__version__ = "1.32.0"
+__version__ = "1.33.0"
 
 import importlib.util
 

ai_diffusion/cloud_client.py

@@ -189,7 +189,11 @@ async def _process_job(self, job: JobInfo):
             images = await self.receive_images(output["images"])
             pose = output.get("pose", None)
             log.info(f"{job} completed, got {len(images)} images{', got pose' if pose else ''}")
-            yield ClientMessage(ClientEvent.finished, job.local_id, 1, images, pose)
+            lora_warning = output.get("lora_warning", False)
+            if lora_warning:
+                log.warning(f"{job} encountered LoRA that could not be applied to the checkpoint")
+            error = "incompatible_lora" if lora_warning else None
+            yield ClientMessage(ClientEvent.finished, job.local_id, 1, images, pose, error=error)
 
         elif response["status"] == "FAILED":
             err_msg, err_trace = _extract_error(response, job.remote_id)

ai_diffusion/model.py

@@ -52,9 +52,15 @@ class ProgressKind(Enum):
 
 class ErrorKind(Enum):
     none = 0
-    plugin_error = 1
-    server_error = 2
-    insufficient_funds = 3
+    plugin_error = 100
+    server_error = 200
+    insufficient_funds = 201
+    warning = 300
+    incompatible_lora = 301
+
+    @property
+    def is_warning(self):
+        return self.value >= ErrorKind.warning.value
 
 
 class Error(NamedTuple):
@@ -65,6 +71,11 @@ class Error(NamedTuple):
     def __bool__(self):
         return self.kind is not ErrorKind.none
 
+    @staticmethod
+    def from_string(s: str, fallback: ErrorKind | None = None):
+        kind = ErrorKind[s] if s in ErrorKind.__members__ else fallback or ErrorKind.warning
+        return Error(kind, s)
+
 
 no_error = Error(ErrorKind.none, "")
 
@@ -463,7 +474,7 @@ def cancel(self, active=False, queued=False):
 
     def report_error(self, error: Error | str):
         if isinstance(error, str):
-            error = Error(ErrorKind.server_error, error)
+            error = Error.from_string(error, ErrorKind.server_error)
         self.error = error
         self.live.is_active = False
         self.custom.is_live = False
@@ -493,6 +504,8 @@ def handle_message(self, message: ClientMessage):
         elif message.event is ClientEvent.output:
             self.custom.show_output(message.result)
         elif message.event is ClientEvent.finished:
+            if message.error:  # successful jobs may have encountered some warnings
+                self.report_error(Error.from_string(message.error, ErrorKind.warning))
             if message.images:
                 self.jobs.set_results(job, message.images)
             if job.kind is JobKind.control_layer:

ai_diffusion/ui/widget.py

@@ -878,6 +878,8 @@ def error(self, error: Error):
         self._original_error = error.message if error else ""
         if error.kind is ErrorKind.insufficient_funds:
             self._show_payment_error(error.data)
+        elif error.kind.is_warning:
+            self._show_warning(error.kind, error.message)
         elif error:
             self._show_error(error.message)
 
@@ -897,6 +899,20 @@ def _show_error(self, text: str):
         self._copy_button.setVisible(True)
         self.setVisible(True)
 
+    def _show_warning(self, kind: ErrorKind, text: str):
+        self.reset(theme.yellow)
+        if kind is ErrorKind.incompatible_lora:
+            text = (
+                _(
+                    "Selected LoRA model could not be applied. Please make sure it is compatible with the checkpoint base model you are using."
+                )
+                + " <a href='https://docs.interstice.cloud/base-models'>"
+                + _("Lean more")
+                + "</a>"
+            )
+        self._label.setText(text)
+        self.setVisible(True)
+
     def _show_payment_error(self, data: dict[str, Any] | None):
         self.reset(theme.yellow)
         message = "Insufficient funds"