Ruff: manual fixes

Author: Acly

ai_diffusion/client.py

@@ -5,6 +5,7 @@
 from abc import ABC, abstractmethod
 from collections import deque
 from collections.abc import AsyncGenerator, Callable, Iterable
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any, Generic, NamedTuple, TypeVar
 
@@ -23,6 +24,10 @@
 from .util import ensure
 
 
+class ServerError(Exception):
+    pass
+
+
 class ClientEvent(Enum):
     progress = 0
     finished = 1
@@ -245,7 +250,7 @@ def resource(
         id = ResourceId(kind, arch, identifier)
         model = self.find(id)
         if model is None:
-            raise Exception(f"{id.name} not found")
+            raise KeyError(f"{id.name} not found")
         return model
 
     def find(self, id: ResourceId):
@@ -389,10 +394,11 @@ def from_list(data: list):
         return [TranslationPackage.from_dict(item) for item in data]
 
 
-class ClientFeatures(NamedTuple):
+@dataclass(frozen=True)
+class ClientFeatures:
     ip_adapter: bool = True
     translation: bool = True
-    languages: list[TranslationPackage] = []
+    languages: list[TranslationPackage] = field(default_factory=list)
     max_upload_size: int = 0
     max_control_layers: int = 1000
     gguf: bool = False

ai_diffusion/cloud_client.py

@@ -8,7 +8,7 @@
 from collections.abc import Iterable
 from copy import copy
 from dataclasses import dataclass, field
-from datetime import datetime
+from datetime import datetime, timezone
 from enum import Enum
 from typing import Any
 
@@ -23,6 +23,7 @@
     ClientModels,
     DeviceInfo,
     News,
+    ServerError,
     TranslationPackage,
     User,
     loras_to_upload,
@@ -122,9 +123,9 @@ async def sign_in(self):
         yield sign_in_url
 
         auth_confirm = await self._post("auth/confirm", {"client_id": client_id})
-        time = datetime.now()
+        time = datetime.now(timezone.utc)
         while auth_confirm["status"] == "not-found":
-            if (datetime.now() - time).seconds > 300:
+            if (datetime.now(timezone.utc) - time).seconds > 300:
                 raise TimeoutError(_("Sign-in attempt timed out after 5 minutes"))
             await asyncio.sleep(2)
             auth_confirm = await self._post("auth/confirm", {"client_id": client_id})
@@ -403,9 +404,11 @@ async def _upload_lora(self, lora: File):
                 yield sent / max(total, 1)
         except NetworkError as e:
             log.error(f"LoRA model upload failed [{e.status}]: {e.message}")
-            raise Exception(_("Connection error during upload of LoRA model") + f" {lora.name}")
+            raise ServerError(
+                _("Connection error during upload of LoRA model") + f" {lora.name}"
+            ) from e
         except Exception as e:
-            raise Exception(_("Error during upload of LoRA model") + f" {lora.name}") from e
+            raise ServerError(_("Error during upload of LoRA model") + f" {lora.name}") from e
 
     async def receive_images(self, images: dict):
         offsets = images.get("offsets")

ai_diffusion/comfy_client.py

@@ -26,6 +26,7 @@
     MissingResources,
     OutputBatchMode,
     Quantization,
+    ServerError,
     SharedWorkflow,
     TextOutput,
     TranslationPackage,
@@ -174,7 +175,7 @@ async def connect(url=default_url, access_token=""):
                 pass
         except Exception as e:
             msg = _("Could not establish websocket connection at") + f" {wsurl}: {e!s}"
-            raise Exception(msg)
+            raise ServerError(msg)
 
         # Check custom nodes
         log.info("Checking for required custom nodes...")
@@ -615,7 +616,7 @@ async def upload_loras(self, work: WorkflowInput, local_job_id: str):
 
                 await self.refresh()
             except Exception as e:
-                raise Exception(_("Error during upload of LoRA model") + f" {file.path}: {e!s}")
+                raise ServerError(_("Error during upload of LoRA model") + f" {file.path}: {e!s}")
 
     async def _get_active_job(self, job_id: str):
         if self._active_job and self._active_job.id == job_id:
@@ -859,7 +860,7 @@ def _ensure_supported_style(client: Client):
         if checkpoint is None:
             log.warning("No checkpoints found for any of the supported workloads!")
             if len(client.models.checkpoints) == 0:
-                raise Exception(_("No diffusion model checkpoints found"))
+                raise RuntimeError(_("No diffusion model checkpoints found"))
             return
         log.info(f"No supported styles found, creating default style with checkpoint {checkpoint}")
         default = next((s for s in Styles.list() if s.filename == "default.json"), None)

ai_diffusion/comfy_workflow.py

@@ -346,7 +346,7 @@ def sampler_custom_advanced(
         start_at_step=0,
         cfg=7.0,
         seed=-1,
-        extent=Extent(1024, 1024),
+        extent: Extent | None = None,
     ):
         self.sample_count += steps - start_at_step
 
@@ -373,7 +373,12 @@ def sampler_custom_advanced(
         )[1]
 
     def scheduler_sigmas(
-        self, model: Output, scheduler="normal", steps=20, arch=Arch.sdxl, extent=Extent(1024, 1024)
+        self,
+        model: Output,
+        scheduler="normal",
+        steps=20,
+        arch=Arch.sdxl,
+        extent: Extent | None = None,
     ):
         if scheduler in ("align_your_steps", "ays"):
             assert arch is Arch.sd15 or arch.is_sdxl_like
@@ -418,6 +423,7 @@ def scheduler_sigmas(
                 beta=0.5,
             )
         elif scheduler == "flux2":
+            extent = extent or Extent(1024, 1024)
             return self.add(
                 "Flux2Scheduler",
                 output_count=1,

ai_diffusion/custom_workflow.py

@@ -124,7 +124,7 @@ def _process_file(self, file: Path):
             graph = json.load(f)
             self._process_workflow(file.stem, WorkflowSource.local, graph, file)
 
-    def rowCount(self, parent=QModelIndex()):
+    def rowCount(self, parent: QModelIndex | None = None):
         return len(self._workflows)
 
     def data(self, index: QModelIndex, role: int = 0):

ai_diffusion/document.py

@@ -332,10 +332,9 @@ def _selection_is_entire_document(selection: krita.Selection, extent: Extent):
 
 
 class PoseLayers:
-    _layers: dict[str, Pose] = {}
-    _timer = QTimer()
-
     def __init__(self):
+        self._layers: dict[str, Pose] = {}
+        self._timer = QTimer()
         self._timer.setInterval(500)
         self._timer.timeout.connect(self.update)
         self._timer.start()

ai_diffusion/eventloop.py

@@ -38,13 +38,13 @@ def schedule():
 
 
 def stop():
-    global _timer, _loop
+    global _timer
     try:
         _timer.stop()
         _timer = None
         _loop.stop()
         _loop.close()
-    except Exception:
+    except Exception:  # noqa
         pass
 
 

ai_diffusion/extension.py

@@ -16,11 +16,10 @@
 
 
 class AIToolsExtension(Extension):
-    _actions: dict[str, QAction] = {}
-    _settings_dialog: SettingsDialog
-
     def __init__(self, parent):
         super().__init__(parent)
+        self._actions: dict[str, QAction] = {}
+
         log.info(f"Extension initialized, Version: {__version__}, Python: {sys.version}")
 
         extension_dir = Path(__file__).parent

ai_diffusion/files.py

@@ -99,7 +99,7 @@ def __init__(self, database: Path | None = None, parent=None):
         self._files: list[File] = []
         self.load()
 
-    def rowCount(self, parent=QModelIndex()):
+    def rowCount(self, parent: QModelIndex | None = None):
         return len(self._files)
 
     def data(self, index: QModelIndex, role: int = 0):

ai_diffusion/image.py

@@ -378,7 +378,7 @@ def from_bytes(data: QBuffer | QByteArray | memoryview | bytes, format: str | No
         if loader.read(img):
             return Image(img)
         else:
-            raise Exception(f"Failed to load image from buffer: {loader.errorString()}")
+            raise RuntimeError(f"Failed to load image from buffer: {loader.errorString()}")
 
     @staticmethod
     def from_pil(pil_image):
@@ -561,7 +561,7 @@ def write(
                 global _qt_supports_webp
                 _qt_supports_webp = False
                 self.write(buffer, format.no_webp_fallback)
-            raise Exception(f"Failed to write image to buffer: {writer.errorString()} {info}")
+            raise RuntimeError(f"Failed to write image to buffer: {writer.errorString()} {info}")
 
     def to_bytes(self, format=ImageFileFormat.png):
         byte_array = QByteArray()
@@ -604,7 +604,7 @@ def save(
         fmt = format or ImageFileFormat.from_extension(filepath)
         file = QFile(str(filepath))
         if not file.open(QFile.OpenModeFlag.WriteOnly):
-            raise Exception(f"Failed to open {filepath} for writing: {file.errorString()}")
+            raise RuntimeError(f"Failed to open {filepath} for writing: {file.errorString()}")
         try:
             self.write(file, fmt, quality)
         finally: