From 801ef2e3f0dabd1d0d1e3be7752921441ef3e8b2 Mon Sep 17 00:00:00 2001
From: doctorpangloss <@hiddenswitch.com>
Date: Thu, 23 May 2024 22:51:44 -0700
Subject: [PATCH] Fix interrupt messaging, add AMD and Intel Dockerfiles

---
 Dockerfile        | 4 ++++
 amd.Dockerfile    | 1 +
 comfy/cmd/main.py | 6 +++---
 intel.Dockerfile  | 1 +
 4 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 923ccb02b..7f1a4fe72 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,4 +2,8 @@ FROM nvcr.io/nvidia/pytorch:24.04-py3
 RUN pip install --no-cache --no-build-isolation git+https://github.com/hiddenswitch/ComfyUI.git
 EXPOSE 8188
 WORKDIR /workspace
+# tries to address https://github.com/pytorch/pytorch/issues/104801
+# and issues reported by importing nodes_canny
+ENV PYTORCH_CUDA_ALLOC_CONF="backend:cudaMallocAsync"
+RUN comfyui --quick-test-for-ci --cpu --cwd /workspace
 CMD ["/usr/local/bin/comfyui", "--listen"]
diff --git a/amd.Dockerfile b/amd.Dockerfile
index 21d80fb24..c16257701 100644
--- a/amd.Dockerfile
+++ b/amd.Dockerfile
@@ -2,4 +2,5 @@ FROM rocm/pytorch:rocm6.0.2_ubuntu22.04_py3.10_pytorch_2.1.2
 RUN pip install --no-cache --no-build-isolation git+https://github.com/hiddenswitch/ComfyUI.git
 EXPOSE 8188
 WORKDIR /workspace
+RUN comfyui --quick-test-for-ci --cpu --cwd /workspace
 CMD ["/usr/local/bin/comfyui", "--listen"]
diff --git a/comfy/cmd/main.py b/comfy/cmd/main.py
index 48886eef5..834c48941 100644
--- a/comfy/cmd/main.py
+++ b/comfy/cmd/main.py
@@ -213,11 +213,11 @@ async def main():
     try:
         await run(server, address=args.listen, port=args.port, verbose=not args.dont_print_server,
                   call_on_start=call_on_start)
-    except asyncio.CancelledError:
+    except (asyncio.CancelledError, KeyboardInterrupt):
+        logging.debug("\nStopped server")
+    finally:
         if distributed:
             await q.close()
-        logging.debug("\nStopped server")
-
     cleanup_temp()
 
 
diff --git a/intel.Dockerfile b/intel.Dockerfile
index 5b0d09bd1..c3bc34a46 100644
--- a/intel.Dockerfile
+++ b/intel.Dockerfile
@@ -2,4 +2,5 @@ FROM intel/intel-optimized-pytorch:2.3.0-pip-base
 RUN pip install --no-cache --no-build-isolation git+https://github.com/hiddenswitch/ComfyUI.git
 EXPOSE 8188
 WORKDIR /workspace
+RUN comfyui --quick-test-for-ci --cpu --cwd /workspace
 CMD ["/usr/local/bin/comfyui", "--listen"]
\ No newline at end of file