Krita 6.x / Qt6

.github/workflows/test.yml

@@ -1,5 +1,5 @@
 name: Lint, Typecheck and Test
-on: [push, pull_request]
+on: [pull_request]
 
 jobs:
   check:
@@ -35,7 +35,15 @@ jobs:
       matrix:
         os: [ubuntu-22.04, macos-14, windows-2022]
     runs-on: ${{ matrix.os }}
+    env:
+      QT_QPA_PLATFORM: offscreen
     steps:
+      # PyQt6.QtGui import requires libEGL.so.1 (even with QT_QPA_PLATFORM=offscreen)
+      - name: Packages
+        if: runner.os == 'Linux'
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y libegl1 libgl1 libxkbcommon-x11-0
       - name: Checkout
         uses: actions/checkout@v3
         with:

.vscode/launch.json

@@ -44,6 +44,16 @@
                 "tests"
             ],
             "cwd": "${workspaceFolder}",
+        },
+        {
+            "name": "Design",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${workspaceFolder}/scripts/design.py",
+            "console": "integratedTerminal",
+            "env": {
+                "PYTHONPATH": "${workspaceFolder}"
+            },
         }
     ]
 }

AGENTS.md

@@ -9,31 +9,35 @@ image generation functionality via diffusion models.
 The plugin runs within Krita's embedded Python interpreter. It may not use 3rd party
 libraries except for Qt and the websockets library.
 
-### Inference
-* `api.py` has data structures for an inference request (`WorkflowInput`) - everything relevant to image generation MUST be contained here
-* `comfy_client.py` is a HTTP/WebSocket client that connects to a ComfyUI server to fulfill requests
-* `cloud_client.py` is a HTTP client that connects to an image generation service to fulfill requests
-* `workflow.py` transforms `WorkflowInput` into ComfyUI workflows
-* `server.py` contains an installer for a ComfyUI server that will run in the background alongside the plugin
-* `resources.py` is a central location that enumerates supported AI models and extensions
-
-### UI
-The UI is separated into distinct workspaces:
-* "Generation" for launching asynchronous image jobs and viewing their results
-* "Live" for automatically generating preview output after every change
-* "Upscale" for diffusion-based super-resolution tasks
-* "Custom/Graph" for importing and running custom user ComfyUI workflows
-* "Animation" for batch-processing image frames
-
-Code is separated into:
-* `ai_diffusion/model`: Model classes which hold observable UI state and implement actions
-* `ai_diffusion/ui`: Qt widgets for the user interface
-* Persistence layer which loads/stores state in files or Krita documents
-  * `persistence.py`, `settings.py`, `files.py`, ...
-
-### Image manipulation
-Helpers and tools for Krita's objects (`document.py`, `layer.py`) and
-general image algorithms (`image.py`, `resolution.py`).
+### Architecture
+
+* `ai_diffusion/backend/` - code related to run AI models via ComfyUI
+  * `api.py` - serializable data structures for an inference request (`WorkflowInput`) - everything relevant to image generation MUST be contained here
+  * `comfy_client.py` - HTTP/WebSocket client that connects to a ComfyUI server to fulfill requests
+  * `cloud_client.py` - HTTP client that connects to an image generation service to fulfill requests
+  * `workflow.py` - transforms `WorkflowInput` into ComfyUI workflows
+  * `server.py` - installer for a ComfyUI server that will run in the background alongside the plugin
+  * `resources.py` - central location that enumerates supported AI models and extensions
+* `ai_diffusion/ui/` - Qt widgets that make up the UI
+  * `generation.py` - workspace for launching asynchronous image jobs and viewing their results in a history list
+  * `live.py` - workspace for automatically generating preview output after every change
+  * `upscale.py` - workspace for diffusion-based super-resolution tasks
+  * `custom_workflow.py` - workspace for importing and running custom user ComfyUI workflows
+  * `animation.py` - for batch-processing image frames
+* `ai_diffusion/model/` - classes which hold the observable app state and actions, often there is a 1:1 relationship between `model/` and `ui/` files
+  * `connection.py` - manages the backend clients, shared across all documents
+  * `model.py` - the document model, almost all state is kept per opened document
+  * `jobs.py` - queued and finished diffusion jobs
+  * `control.py` - state related to control layers (images used as input to ControlNet models)
+  * `region.py` - state related to regional prompts (images whose alpha coverage is used as attention masks for diffusion)
+* `ai_diffusion/` - Krita wrapper classes, persistence layer, image manipulation, and general helpers and utils
+  * `document.py` - used to interact with Krita documents
+  * `persistence.py` - stores/loads document state to .kra files
+  * `settings.py` - stores plugin settings in settings.json
+  * `style.py` - stores `Style` objects, each references a diffusion model along with parameters, LoRA, and default prompts
+  * `files.py` - lists of files, both local filesystem or remote on a server
+  * `image.py` - `Image`, `Mask`, `Extent` and `Bounds` objects, image manipulation
+  * `text.py` - text prompt processing tools
 
 
 ## Commands
@@ -56,6 +60,13 @@ Run tests with these priorities:
 4. *only after changing cloud client* (includes inference): `pytest tests/test_workflow.py --cloud`
 5. *only after changing installer* (slow tests): `pytest tests/test_server.py --test-install`
 
+### UI Tests
+
+There are no dedicated tests for UI. The following command can be used as a quick check
+that the UI code runs without errors:
+```
+python scripts/design.py --exit
+```
 
 ## Code Guidelines
 

ai_diffusion/__init__.py

@@ -18,8 +18,8 @@
 
     if not getattr(krita, "IS_MOCK", False):
         krita_ver = krita.Krita.instance().version()
-        if not krita_ver.startswith("5"):
-            raise ImportError(f"This Plugin is for Krita 5.x, but you are using Krita {krita_ver}.")
+        if not krita_ver.startswith("6"):
+            raise ImportError(f"This Plugin is for Krita 6.x, but you are using Krita {krita_ver}.")
 
         from .extension import AIToolsExtension as AIToolsExtension
 

ai_diffusion/backend/client.py

@@ -9,7 +9,7 @@
 from enum import Enum
 from typing import Any, Generic, NamedTuple, TypeVar
 
-from PyQt5.QtCore import QObject, pyqtSignal
+from PyQt6.QtCore import QObject, pyqtSignal
 
 from ..files import FileFormat, FileLibrary
 from ..image import ImageCollection, Point

ai_diffusion/backend/comfy_workflow.py

@@ -269,7 +269,7 @@ def _add_image(self, image: Image):
         return id
 
     def _add_image_hashed(self, image: Image):
-        data = image.to_bytes()
+        data = image.to_bytes().data()
         hash = zlib.crc32(data)
         id = f"{hash:08x}"
         self.image_data[id] = data

ai_diffusion/backend/network.py

@@ -8,22 +8,27 @@
 from pathlib import Path
 from typing import NamedTuple
 
-from PyQt5.QtCore import QBuffer, QByteArray, QFile, QUrl
-from PyQt5.QtNetwork import QNetworkAccessManager, QNetworkReply, QNetworkRequest, QSslError
+from PyQt6.QtCore import QBuffer, QByteArray, QFile, QUrl
+from PyQt6.QtNetwork import QNetworkAccessManager, QNetworkReply, QNetworkRequest, QSslError
 
 from ..localization import translate as _
 from ..util import client_logger as log
 
 
 class NetworkError(Exception):
-    code: int
+    code: int | QNetworkReply.NetworkError
     message: str
     url: str
     status: int | None = None
     data: dict | None = None
 
     def __init__(
-        self, code: int, msg: str, url: str, status: int | None = None, data: dict | None = None
+        self,
+        code: int | QNetworkReply.NetworkError,
+        msg: str,
+        url: str,
+        status: int | None = None,
+        data: dict | None = None,
     ):
         self.code = code
         self.message = msg
@@ -37,7 +42,7 @@ def __str__(self):
 
     @staticmethod
     def from_reply(reply: QNetworkReply):
-        code: QNetworkReply.NetworkError = reply.error()  # type: ignore (bug in PyQt5-stubs)
+        code = reply.error()
         url = reply.url().toString()
         status = reply.attribute(QNetworkRequest.Attribute.HttpStatusCodeAttribute)
         if reply.isReadable():
@@ -96,7 +101,10 @@ def set_auth(self, bearer: str):
 
     def _prepare_request(self, url: str, timeout: float | None = None, bearer: str | None = None):
         request = QNetworkRequest(QUrl(url))
-        request.setAttribute(QNetworkRequest.FollowRedirectsAttribute, True)
+        request.setAttribute(
+            QNetworkRequest.Attribute.RedirectPolicyAttribute,
+            QNetworkRequest.RedirectPolicy.NoLessSafeRedirectPolicy,
+        )
         bearer_token = bearer or self._bearer_token
         if bearer_token:
             request.setRawHeader(b"Authorization", f"Bearer {bearer_token}".encode())
@@ -110,7 +118,7 @@ def http(
         self,
         method,
         url: str,
-        data: dict | QByteArray | None = None,
+        data: dict | QByteArray | bytes | None = None,
         timeout: float | None = None,
         bearer: str | None = None,
     ):
@@ -158,7 +166,10 @@ async def upload(self, url: str, data: QByteArray | bytes, sha256: str | None =
         assert isinstance(data, QByteArray)
 
         request = QNetworkRequest(QUrl(url))
-        request.setAttribute(QNetworkRequest.Attribute.FollowRedirectsAttribute, True)
+        request.setAttribute(
+            QNetworkRequest.Attribute.RedirectPolicyAttribute,
+            QNetworkRequest.RedirectPolicy.NoLessSafeRedirectPolicy,
+        )
         if sha256:
             request.setRawHeader(b"x-amz-checksum-sha256", sha256.encode("utf-8"))
         request.setHeader(
@@ -192,7 +203,7 @@ def download(self, url: str, timeout: float | None = None):
         buffer.open(QBuffer.OpenModeFlag.WriteOnly)
 
         def write(bytes_received, bytes_total):
-            buffer.write(reply.readAll())
+            buffer.write(reply.readAll().data())
 
         future = asyncio.get_running_loop().create_future()
         tracker = Request(url, future, buffer)
@@ -212,14 +223,14 @@ def _upload_progress(self, bytes_sent: int, bytes_total: int):
     def _finished(self, reply: QNetworkReply):
         future = None
         try:
-            code = reply.error()  # type: ignore (bug in PyQt5-stubs)
+            code = reply.error()
             tracker = self._requests[reply]
             future = tracker.future
             if future.cancelled():
                 return  # operation was cancelled, discard result
             if code == QNetworkReply.NetworkError.NoError:
                 if tracker.buffer is not None:
-                    tracker.buffer.write(reply.readAll())
+                    tracker.buffer.write(reply.readAll().data())
                     future.set_result(tracker.buffer.data())
                 else:
                     content_type = reply.header(QNetworkRequest.KnownHeaders.ContentTypeHeader)
@@ -297,13 +308,17 @@ def _write_file_chunks(file: QFile, reply: QNetworkReply):
 
 async def _try_download(network: QNetworkAccessManager, url: str, path: Path):
     out_file = QFile(str(path) + ".part")
-    if not out_file.open(QFile.ReadWrite | QFile.Append):  # type: ignore
+    rwa = QFile.OpenModeFlag.ReadOnly | QFile.OpenModeFlag.WriteOnly | QFile.OpenModeFlag.Append
+    if not out_file.open(rwa):
         raise RuntimeError(
             _("Error during download: could not open {path} for writing", path=out_file.fileName())
         )
 
     request = QNetworkRequest(QUrl(_map_host(url)))
-    request.setAttribute(QNetworkRequest.FollowRedirectsAttribute, True)
+    request.setAttribute(
+        QNetworkRequest.Attribute.RedirectPolicyAttribute,
+        QNetworkRequest.RedirectPolicy.NoLessSafeRedirectPolicy,
+    )
     if out_file.size() > 0:
         log.info(f"Found {path}.part, resuming download from {out_file.size()} bytes")
         request.setRawHeader(b"Range", f"bytes={out_file.size()}-".encode())
@@ -325,9 +340,9 @@ def handle_finished():
         out_file.close()
         if finished_future.cancelled():
             return  # operation was cancelled, discard result
-        if reply.error() == QNetworkReply.NetworkError.NoError:  # type: ignore (bug in PyQt5-stubs)
+        if reply.error() == QNetworkReply.NetworkError.NoError:
             finished_future.set_result(path)
-        elif reply.attribute(QNetworkRequest.HttpStatusCodeAttribute) == 416:
+        elif reply.attribute(QNetworkRequest.Attribute.HttpStatusCodeAttribute) == 416:
             # 416 = Range Not Satisfiable
             finished_future.set_exception(NetworkError(416, "Resume not supported", url))
         else:

ai_diffusion/backend/server.py

@@ -11,7 +11,7 @@
 from pathlib import Path
 from typing import NamedTuple
 
-from PyQt5.QtNetwork import QNetworkAccessManager
+from PyQt6.QtNetwork import QNetworkAccessManager
 
 from .. import eventloop
 from ..localization import translate as _

ai_diffusion/document.py

@@ -7,7 +7,7 @@
 
 import krita
 from krita import Krita
-from PyQt5.QtCore import QByteArray, QObject, QTimer, pyqtSignal
+from PyQt6.QtCore import QByteArray, QObject, QTimer, QUuid, pyqtSignal
 
 from .image import Bounds, Extent, Image, Mask
 from .layer import Layer, LayerManager, LayerType
@@ -342,7 +342,7 @@ def _selection_is_entire_document(selection: krita.Selection, extent: Extent):
 
 class PoseLayers:
     def __init__(self):
-        self._layers: dict[str, Pose] = {}
+        self._layers: dict[QUuid, Pose] = {}
         self._timer = QTimer()
         self._timer.setInterval(500)
         self._timer.timeout.connect(self.update)

ai_diffusion/eventloop.py

@@ -1,7 +1,7 @@
 import asyncio
 from collections.abc import Callable
 
-from PyQt5.QtCore import QTimer
+from PyQt6.QtCore import QTimer
 
 _loop = asyncio.new_event_loop()
 _timer = QTimer()

ai_diffusion/extension.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 
 from krita import DockWidgetFactory, DockWidgetFactoryBase, Extension, Krita, Window  # type: ignore
-from PyQt5.QtWidgets import QAction
+from PyQt6.QtGui import QAction
 
 from . import __version__, eventloop
 from .model.model import Workspace
@@ -88,5 +88,9 @@ def createActions(self, window):
 
 Krita.instance().addExtension(AIToolsExtension(Krita.instance()))
 Krita.instance().addDockWidgetFactory(
-    DockWidgetFactory("imageDiffusion", DockWidgetFactoryBase.DockRight, ImageDiffusionWidget)  # type: ignore
+    DockWidgetFactory(
+        "imageDiffusion",
+        DockWidgetFactoryBase.DockPosition.DockRight,  # type: ignore
+        ImageDiffusionWidget,
+    )
 )

ai_diffusion/files.py

@@ -7,8 +7,8 @@
 from pathlib import Path
 from typing import Any, NamedTuple, cast
 
-from PyQt5.QtCore import QAbstractListModel, QModelIndex, QSortFilterProxyModel, Qt
-from PyQt5.QtGui import QIcon
+from PyQt6.QtCore import QAbstractListModel, QModelIndex, QSortFilterProxyModel, Qt
+from PyQt6.QtGui import QIcon
 
 from .util import client_logger as log
 from .util import encode_json, read_json_with_comments, user_data_dir