Use pagination for model inspection #2079 #1796

* inspect models in chunks of 8
* show some progress in the UI rather than just "Connecting..."
* timeout for one page is 1min
* total timeout is 3min
* -> issues are reported faster, and large model collections have more time
* (untested: should also support partial results)

Author: Acly

ai_diffusion/client.py

@@ -130,16 +130,13 @@ def deduce_from_filename(filename: str):
 class ClientModels:
     """Collects names of AI models the client has access to."""
 
-    checkpoints: dict[str, CheckpointInfo]
-    vae: list[str]
-    loras: list[str]
-    upscalers: list[str]
-    resources: dict[str, str | None]
-    node_inputs: ComfyObjectInfo
-
     def __init__(self) -> None:
+        self.checkpoints: dict[str, CheckpointInfo] = {}
+        self.vae: list[str] = []
+        self.loras: list[str] = []
+        self.upscalers: list[str] = []
         self.node_inputs = ComfyObjectInfo({})
-        self.resources = {}
+        self.resources: dict[str, str | None] = {}
 
     def resource(
         self, kind: ResourceKind, identifier: ControlMode | UpscalerName | str, arch: Arch
@@ -297,6 +294,14 @@ class Client(ABC):
     @abstractmethod
     async def connect(url: str, access_token: str = "") -> Client: ...
 
+    class DiscoverStatus(NamedTuple):
+        folder: str
+        current: int
+        total: int
+
+    def discover_models(self, refresh: bool) -> AsyncGenerator[DiscoverStatus, Any]:
+        raise NotImplementedError()
+
     @abstractmethod
     async def enqueue(self, work: WorkflowInput, front: bool = False) -> str: ...
 

ai_diffusion/comfy_client.py

@@ -7,6 +7,7 @@
 from enum import Enum
 from collections import deque
 from itertools import chain, product
+from time import time
 from typing import Any, Optional, Sequence
 
 from .api import WorkflowInput
@@ -164,31 +165,49 @@ async def connect(url=default_url, access_token=""):
         loras = nodes.options("LoraLoader", "lora_name")
         available_resources.update(_find_loras(loras))
 
-        # Retrieve list of checkpoints
-        checkpoints = await client.try_inspect("checkpoints")
-        diffusion_models = await client.try_inspect("diffusion_models")
-        diffusion_models.update(await client.try_inspect("unet_gguf"))
-        client._refresh_models(nodes, checkpoints, diffusion_models)
-
-        # Check supported base models and make sure there is at least one
-        client._supported_archs = {ver: client._check_workload(ver) for ver in Arch.list()}
-        supported_workloads = [
-            arch for arch, miss in client._supported_archs.items() if len(miss) == 0
-        ]
-        log.info("Supported workloads: " + ", ".join(arch.value for arch in supported_workloads))
-        if len(supported_workloads) == 0 and settings.check_server_resources:
-            raise MissingResources(client._supported_archs)
-
         # Workarounds for DirectML
         if client.device_info.type == "privateuseone":
             # OmniSR causes a crash
             for n in [2, 3, 4]:
                 id = resource_id(ResourceKind.upscaler, Arch.all, UpscalerName.fast_x(n))
                 available_resources[id] = models.default_upscaler
 
-        _ensure_supported_style(client)
         return client
 
+    async def discover_models(self, refresh: bool):
+        if refresh:
+            nodes = ComfyObjectInfo(await self._get("object_info"))
+        else:
+            nodes = self.models.node_inputs
+
+        checkpoints: dict[str, dict] = {}
+        diffusion_models: dict[str, dict] = {}
+        async for status, result in self.try_inspect("checkpoints"):
+            yield status
+            checkpoints.update(result)
+        async for status, result in self.try_inspect("diffusion_models"):
+            yield status
+            diffusion_models.update(result)
+        async for status, result in self.try_inspect("unet_gguf"):
+            yield status
+            diffusion_models.update(result)
+        self._refresh_models(nodes, checkpoints, diffusion_models)
+
+        # Check supported base models and make sure there is at least one
+        self._supported_archs = {ver: self._check_workload(ver) for ver in Arch.list()}
+        supported_workloads = [
+            arch for arch, miss in self._supported_archs.items() if len(miss) == 0
+        ]
+        log.info("Supported workloads: " + ", ".join(arch.value for arch in supported_workloads))
+        if not refresh and len(supported_workloads) == 0 and settings.check_server_resources:
+            raise MissingResources(self._supported_archs)
+
+        _ensure_supported_style(self)
+
+    async def refresh(self):
+        async for __ in self.discover_models(refresh=True):
+            pass
+
     async def _get(self, op: str, timeout: float | None = 60):
         return await self._requests.get(f"{self.url}/{op}", timeout=timeout)
 
@@ -386,14 +405,26 @@ async def disconnect(self):
                 self._unsubscribe_workflows(),
             )
 
-    async def try_inspect(self, folder_name: str) -> dict[str, Any]:
+    async def try_inspect(self, folder_name: str):
         if "gguf" in folder_name and not self.features.gguf:
-            return {}
+            return
         try:
-            return await self._get(f"api/etn/model_info/{folder_name}", timeout=120)
+            log.info(f"Inspecting models at {self.url}/api/etn/model_info/{folder_name}")
+            start, timeout = time(), 180
+            offset, total = 0, 100
+            while offset < total and (time() - start) < timeout:
+                r = await self._get(f"api/etn/model_info/{folder_name}?offset={offset}&limit=8")
+                if "_meta" not in r:  # server doesn't support pagination
+                    yield (Client.DiscoverStatus(folder_name, len(r), len(r)), r)
+                    return
+                total = r["_meta"]["total"]
+                del r["_meta"]
+                yield (Client.DiscoverStatus(folder_name, offset + len(r), total), r)
+                offset += 8
+            if offset < total:
+                log.warning(f"Timeout while inspecting models, received {offset}/{total} entries")
         except NetworkError as e:
             log.error(f"Error while inspecting models in {folder_name}: {str(e)}")
-            return {}
 
     @property
     def queued_count(self):
@@ -403,16 +434,6 @@ def queued_count(self):
     def is_executing(self):
         return self._active is not None
 
-    async def refresh(self):
-        nodes, checkpoints, diffusion_models, diffusion_gguf = await asyncio.gather(
-            self._get("object_info"),
-            self.try_inspect("checkpoints"),
-            self.try_inspect("diffusion_models"),
-            self.try_inspect("unet_gguf"),
-        )
-        diffusion_models.update(diffusion_gguf)
-        self._refresh_models(ComfyObjectInfo(nodes), checkpoints, diffusion_models)
-
     def _refresh_models(
         self, nodes: ComfyObjectInfo, checkpoints: dict | None, diffusion_models: dict | None
     ):

ai_diffusion/connection.py

@@ -20,18 +20,22 @@ class ConnectionState(Enum):
     connected = 2
     error = 3
 
-    auth_missing = 4
-    auth_requesting = 5
-    auth_pending = 6
-    auth_error = 7
+    discover_models = 10
+
+    auth_missing = 20
+    auth_requesting = 21
+    auth_pending = 22
+    auth_error = 23
 
 
 class Connection(QObject, ObservableProperties):
     state = Property(ConnectionState.disconnected)
     error = Property("")
+    progress = Property((1, 1))
 
     state_changed = pyqtSignal(ConnectionState)
     error_changed = pyqtSignal(str)
+    progress_changed = pyqtSignal(tuple)
     models_changed = pyqtSignal()
     message_received = pyqtSignal(ClientMessage)
     workflow_published = pyqtSignal(str)
@@ -90,6 +94,9 @@ async def _connect(self, url: str, mode: ServerMode, access_token=""):
                 self._client = await CloudClient.connect(CloudClient.default_api_url, access_token)
             else:
                 self._client = await ComfyClient.connect(url)
+                self.state = ConnectionState.discover_models
+                async for status in self._client.discover_models(refresh=False):
+                    self.progress = (status.current, status.total)
                 self.missing_resources = self._client.missing_resources
 
             apply_performance_preset(settings, self._client.device_info)

ai_diffusion/resources.py

@@ -42,7 +42,7 @@ class CustomNode(NamedTuple):
         "External Tooling Nodes",
         "comfyui-tooling-nodes",
         "https://github.com/Acly/comfyui-tooling-nodes",
-        "fbf99f2a08368211513bcc3a3ab5b8bc62a476b1",
+        "a088a2dde2ca4a0db9193ac1c01ac77c8b31a70d",
         ["ETN_LoadImageCache", "ETN_SaveImageCache", "ETN_Translate"],
     ),
     CustomNode(

ai_diffusion/ui/settings.py

@@ -299,6 +299,7 @@ def __init__(self, server: Server):
 
         root.connection.state_changed.connect(self.update_server_status)
         root.connection.error_changed.connect(self.update_server_status)
+        root.connection.progress_changed.connect(self.update_server_status)
         self.update_server_status()
 
     @property
@@ -352,7 +353,7 @@ def _connect(self):
     def update_server_status(self):
         connection = root.connection
         self._cloud_widget.update_connection_state(connection.state)
-        self._connect_button.setEnabled(connection.state != ConnectionState.connecting)
+        self._connect_button.setEnabled(True)
         self._client_id.setVisible(False)
         if connection.state == ConnectionState.connected:
             self._connection_status.setText(_("Connected"))
@@ -362,6 +363,12 @@ def update_server_status(self):
         elif connection.state == ConnectionState.connecting:
             self._connection_status.setText(_("Connecting"))
             self._connection_status.setStyleSheet(f"color: {yellow}; font-weight:bold")
+            self._connect_button.setEnabled(False)
+        elif connection.state == ConnectionState.discover_models:
+            progress = f" ({connection.progress[0]}/{connection.progress[1]})"
+            self._connection_status.setText(_("Discovering models") + progress)
+            self._connection_status.setStyleSheet(f"color: {yellow}; font-weight:bold")
+            self._connect_button.setEnabled(False)
         elif connection.state == ConnectionState.disconnected:
             self._connection_status.setText(_("Disconnected"))
             self._connection_status.setStyleSheet(f"color: {grey}; font-style:italic")