Only try to analyse workflows after connecting to get node info

Author: Acly

ai_diffusion/custom_workflow.py

@@ -10,8 +10,8 @@
 from PyQt5.QtCore import pyqtSignal
 
 from .api import WorkflowInput
-from .comfy_workflow import ComfyWorkflow
-from .connection import Connection
+from .comfy_workflow import ComfyWorkflow, ComfyNode
+from .connection import Connection, ConnectionState
 from .image import Bounds, Image
 from .jobs import Job, JobParams, JobQueue, JobKind
 from .properties import Property, ObservableProperties
@@ -51,23 +51,29 @@ class WorkflowCollection(QAbstractListModel):
     def __init__(self, connection: Connection, folder: Path | None = None):
         super().__init__()
         self._connection = connection
-        self._workflows: list[CustomWorkflow] = []
-
         self._folder = folder or user_data_dir / "workflows"
-        for file in self._folder.glob("*.json"):
-            try:
-                self._process_file(file)
-            except Exception as e:
-                log.exception(f"Error loading workflow from {file}: {e}")
+        self._workflows: list[CustomWorkflow] = []
 
+        self._connection.state_changed.connect(self._handle_connection)
         self._connection.workflow_published.connect(self._process_remote_workflow)
-        for wf in self._connection.workflows.keys():
-            self._process_remote_workflow(wf)
+        self._handle_connection(self._connection.state)
+
+    def _handle_connection(self, state: ConnectionState):
+        if state in (ConnectionState.connected, ConnectionState.disconnected):
+            self.clear()
+
+        if state is ConnectionState.connected:
+            for file in self._folder.glob("*.json"):
+                try:
+                    self._process_file(file)
+                except Exception as e:
+                    log.exception(f"Error loading workflow from {file}: {e}")
+
+            for wf in self._connection.workflows.keys():
+                self._process_remote_workflow(wf)
 
     def _node_inputs(self):
-        if client := self._connection.client_if_connected:
-            return client.models.node_inputs
-        return {}
+        return self._connection.client.models.node_inputs
 
     def _create_workflow(
         self, id: str, source: WorkflowSource, graph: dict, path: Path | None = None
@@ -127,6 +133,12 @@ def remove(self, id: str):
             self._workflows.pop(idx.row())
             self.endRemoveRows()
 
+    def clear(self):
+        if len(self._workflows) > 0:
+            self.beginResetModel()
+            self._workflows.clear()
+            self.endResetModel()
+
     def set_graph(self, index: QModelIndex, graph: dict):
         wf = self._workflows[index.row()]
         wf.workflow = ComfyWorkflow.import_graph(graph, self._node_inputs())
@@ -266,20 +278,24 @@ def workflow_parameters(w: ComfyWorkflow):
             case ("ETN_Parameter", "choice"):
                 name = node.input("name", "Parameter")
                 default = node.input("default", "")
-                connected, input_name = next(w.find_connected(node.output()), (None, ""))
-                if connected:
-                    if input_type := w.input_type(connected.type, input_name):
-                        if isinstance(input_type[0], list):
-                            yield CustomParam(
-                                ParamKind.choice, name, choices=input_type[0], default=default
-                            )
+                if choices := _get_choices(w, node):
+                    yield CustomParam(ParamKind.choice, name, choices=choices, default=default)
                 else:
                     yield CustomParam(ParamKind.text, name, default=default)
             case ("ETN_Parameter", unknown_type) if unknown_type != "auto":
                 unknown = node.input("name", "?") + ": " + unknown_type
                 log.warning(f"Custom workflow has an unsupported parameter type {unknown}")
 
 
+def _get_choices(w: ComfyWorkflow, node: ComfyNode):
+    connected, input_name = next(w.find_connected(node.output()), (None, ""))
+    if connected:
+        if input_type := w.input_type(connected.type, input_name):
+            if isinstance(input_type[0], list):
+                return input_type[0]
+    return None
+
+
 ImageGenerator = Callable[[WorkflowInput | None], Awaitable[None | Literal[False] | WorkflowInput]]
 
 

ai_diffusion/root.py

@@ -40,6 +40,7 @@ def init(self):
         self._files = FileLibrary.load()
         self._workflows = WorkflowCollection(self._connection)
         self._models = []
+        self._null_model = Model(Document(), self._connection, self._workflows)
         self._recent = RecentlyUsedSync.from_settings()
         self._auto_update = AutoUpdate()
         if settings.auto_update:
@@ -96,7 +97,7 @@ def auto_update(self) -> AutoUpdate:
     def active_model(self):
         if model := self.model_for_active_document():
             return model
-        return Model(Document(), self._connection, self._workflows)
+        return self._null_model
 
     async def autostart(self, signal_server_change: Callable):
         connection = self._connection

tests/test_custom_workflow.py

@@ -3,9 +3,9 @@
 from pathlib import Path
 from PyQt5.QtCore import Qt
 
-from ai_diffusion.api import CustomWorkflowInput, ImageInput, SamplingInput
-from ai_diffusion.client import ClientModels, CheckpointInfo
-from ai_diffusion.connection import Connection
+from ai_diffusion.api import CustomWorkflowInput, ImageInput, WorkflowInput
+from ai_diffusion.client import Client, ClientModels, CheckpointInfo
+from ai_diffusion.connection import Connection, ConnectionState
 from ai_diffusion.comfy_workflow import ComfyNode, ComfyWorkflow, Output
 from ai_diffusion.custom_workflow import CustomWorkflow, WorkflowSource, WorkflowCollection
 from ai_diffusion.custom_workflow import SortedWorkflows, CustomWorkspace
@@ -19,6 +19,40 @@
 from .config import test_dir
 
 
+class MockClient(Client):
+    def __init__(self, node_inputs: dict[str, dict]):
+        self.models = ClientModels()
+        self.models.node_inputs = node_inputs
+
+    @staticmethod
+    async def connect(url: str, access_token: str = "") -> Client:
+        return MockClient({})
+
+    async def enqueue(self, work: WorkflowInput, front: bool = False) -> str:
+        return ""
+
+    async def listen(self):  # type: ignore
+        return
+
+    async def interrupt(self):
+        pass
+
+    async def clear_queue(self):
+        pass
+
+
+def create_mock_connection(
+    initial_workflows: dict[str, dict],
+    node_inputs: dict[str, dict] | None = None,
+    state: ConnectionState = ConnectionState.connected,
+):
+    connection = Connection()
+    connection._client = MockClient(node_inputs or {})
+    connection._workflows = initial_workflows
+    connection.state = state
+    return connection
+
+
 def _assert_has_workflow(
     collection: WorkflowCollection,
     name: str,
@@ -44,12 +78,13 @@ def test_collection(tmp_path: Path):
     file2_graph = {"0": {"class_type": "F2", "inputs": {}}}
     file2.write_text(json.dumps(file2_graph))
 
-    connection = Connection()
     connection_graph = {"0": {"class_type": "C1", "inputs": {}}}
     connection_workflows = {"connection1": connection_graph}
-    connection._workflows = connection_workflows
+    connection = create_mock_connection(connection_workflows, state=ConnectionState.disconnected)
 
     collection = WorkflowCollection(connection, tmp_path)
+    assert len(collection) == 0
+    connection.state = ConnectionState.connected
     assert len(collection) == 3
     _assert_has_workflow(collection, "file1", WorkflowSource.local, file1_graph, file1)
     _assert_has_workflow(collection, "file2", WorkflowSource.local, file2_graph, file2)
@@ -110,7 +145,7 @@ def make_dummy_graph(n: int = 42):
 def test_files(tmp_path: Path):
     collection_folder = tmp_path / "workflows"
 
-    collection = WorkflowCollection(Connection(), collection_folder)
+    collection = WorkflowCollection(create_mock_connection({}, {}), collection_folder)
     assert len(collection) == 0
 
     file1 = tmp_path / "file1.json"
@@ -147,9 +182,8 @@ async def dummy_generate(workflow_input):
 
 
 def test_workspace():
-    connection = Connection()
     connection_workflows = {"connection1": make_dummy_graph(42)}
-    connection._workflows = connection_workflows
+    connection = create_mock_connection(connection_workflows, {})
     workflows = WorkflowCollection(connection)
 
     jobs = JobQueue()