Enable CUDA in managed server tests if available

Author: Acly

tests/test_client.py

@@ -5,6 +5,7 @@
 from ai_diffusion import eventloop, resources
 from ai_diffusion.api import WorkflowInput, WorkflowKind, LoraInput
 from ai_diffusion.api import CheckpointInput, ImageInput, SamplingInput, ConditioningInput
+from ai_diffusion.platform_tools import get_cuda_devices
 from ai_diffusion.resources import ControlMode
 from ai_diffusion.network import NetworkError
 from ai_diffusion.image import Extent
@@ -19,12 +20,17 @@
 
 @pytest.fixture(scope="session")
 def comfy_server(qtapp):
-    server = Server(str(server_dir), ServerBackend.cpu)
+    backend = ServerBackend.cpu
+    if len(get_cuda_devices()) > 0:
+        backend = ServerBackend.cuda
+
+    server = Server(str(server_dir), backend)
     assert server.state is ServerState.stopped, (
         f"Expected server installation at {server_dir}. To create the default installation run"
         " `pytest tests/test_server.py --test-install`"
     )
-    yield qtapp.run(server.start(port=8189))
+    qtapp.run(server.start(port=8189))
+    yield server
     qtapp.run(server.stop())
 
 
@@ -47,9 +53,10 @@ async def main():
 
 
 @pytest.mark.parametrize("cancel_point", ["after_enqueue", "after_start", "after_sampling"])
-def test_cancel(qtapp, comfy_server, cancel_point):
+def test_cancel(qtapp, comfy_server: Server, cancel_point):
     async def main():
-        client = await ComfyClient.connect(comfy_server)
+        assert comfy_server.url is not None
+        client = await ComfyClient.connect(comfy_server.url)
         async for _ in client.discover_models(refresh=False):
             pass
         job_id = None
@@ -64,7 +71,7 @@ async def main():
                 assert msg.event is not ClientEvent.finished
                 assert msg.job_id == job_id or msg.job_id == ""
                 if not job_id:
-                    job_id = await client.enqueue(make_default_work(steps=200))
+                    job_id = await client.enqueue(make_default_work(steps=1000))
                     assert client.queued_count == 1
                 if not interrupted:
                     if cancel_point == "after_enqueue":
@@ -99,13 +106,14 @@ async def main():
     qtapp.run(main())
 
 
-def test_disconnect(qtapp, comfy_server):
+def test_disconnect(qtapp, comfy_server: Server):
     async def listen(client: ComfyClient):
         async for msg in client.listen():
             assert msg.event is ClientEvent.connected
 
     async def main():
-        client = await ComfyClient.connect(comfy_server)
+        assert comfy_server.url is not None
+        client = await ComfyClient.connect(comfy_server.url)
         task = eventloop._loop.create_task(listen(client))
         task.cancel()
         with pytest.raises(asyncio.CancelledError):
@@ -156,16 +164,23 @@ def check_resolve_sd_version(client: ComfyClient, arch: Arch):
     assert resolve_arch(style, None) == arch
 
 
-def test_info(pytestconfig, qtapp, comfy_server):
+def check_nunchaku(server: Server, client: ComfyClient):
+    if server.backend is ServerBackend.cuda:
+        assert "NunchakuFluxDiTLoader" in client.models.node_inputs.nodes
+
+
+def test_info(pytestconfig, qtapp, comfy_server: Server):
     async def main():
-        client = await ComfyClient.connect(comfy_server)
+        assert comfy_server.url is not None
+        client = await ComfyClient.connect(comfy_server.url)
         async for _ in client.discover_models(refresh=False):
             pass
         check_client_info(client)
         await client.refresh()
         check_client_info(client)
         check_resolve_sd_version(client, Arch.sd15)
         # check_resolve_sd_version(client, Arch.sdxl) # no SDXL checkpoint in default installation
+        check_nunchaku(comfy_server, client)
 
     qtapp.run(main())
 

tests/test_server.py

@@ -7,6 +7,7 @@
 import sys
 
 from ai_diffusion import network, server, resources
+from ai_diffusion.platform_tools import get_cuda_devices
 from ai_diffusion.style import Arch
 from ai_diffusion.server import Server, ServerState, ServerBackend, InstallationProgress
 from ai_diffusion.server import model_dirs
@@ -48,6 +49,12 @@ def clear_test_server():
     server_dir.mkdir(exist_ok=True)
 
 
+def get_backend():
+    if len(get_cuda_devices()) > 0:
+        return ServerBackend.cuda
+    return ServerBackend.cpu
+
+
 def test_install_and_run(qtapp, pytestconfig, local_download_server):
     """Test installing and running ComfyUI server from scratch.
     * Takes a while, only runs with --test-install
@@ -62,7 +69,7 @@ def test_install_and_run(qtapp, pytestconfig, local_download_server):
 
     clear_test_server()
 
-    server = Server(str(server_dir), ServerBackend.cpu)
+    server = Server(str(server_dir), get_backend())
     assert server.state in [ServerState.not_installed, ServerState.missing_resources]
 
     last_stage = ""
@@ -110,7 +117,7 @@ def test_run_external(qtapp, pytestconfig):
     if not (server_dir / "ComfyUI").exists():
         pytest.skip("ComfyUI installation not found")
 
-    server = Server(str(server_dir), ServerBackend.cpu)
+    server = Server(str(server_dir), get_backend())
     assert server.has_python
     assert server.state in [ServerState.stopped, ServerState.missing_resources]
 
@@ -129,7 +136,7 @@ def test_extra_model_dirs(pytestconfig):
     if not pytestconfig.getoption("--test-install"):
         pytest.skip("Only runs with --test-install")
 
-    server = Server(str(server_dir), ServerBackend.cpu)
+    server = Server(str(server_dir), get_backend())
     # Requires test_install_and_run to setup the server
     assert server.state in [ServerState.stopped]
 
@@ -141,7 +148,7 @@ def test_verify_and_fix(qtapp, pytestconfig, local_download_server):
     if not pytestconfig.getoption("--test-install"):
         pytest.skip("Only runs with --test-install")
 
-    server = Server(str(server_dir), ServerBackend.cpu)
+    server = Server(str(server_dir), get_backend())
     # Requires test_install_and_run to setup the server
     assert server.state in [ServerState.stopped]
 
@@ -179,7 +186,7 @@ def test_uninstall(qtapp, pytestconfig, local_download_server):
     temp_server_dir = server_dir.parent / "temp-server"
     if temp_server_dir.exists():
         shutil.rmtree(temp_server_dir, ignore_errors=True)
-    server = Server(str(temp_server_dir), ServerBackend.cpu)
+    server = Server(str(temp_server_dir), get_backend())
     assert server.state is ServerState.not_installed
 
     def handle_progress(report: InstallationProgress):