Set ONEDNN_MAX_CPU_ISA=AVX2 as potential workaround for #401

Acly · Feb 26, 2024 · 69b1d83 · 69b1d83
1 parent d7f3c6d
commit 69b1d83
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/ai_diffusion/server.py b/ai_diffusion/server.py
@@ -367,6 +367,7 @@ async def start(self, port: int | None = None):
         self.state = ServerState.starting
         try:
             args = ["-su", "-Xutf8", "main.py"]
+            env = {}
             if self.backend is ServerBackend.cpu:
                 args.append("--cpu")
             elif self.backend is ServerBackend.directml:
@@ -377,7 +378,11 @@ async def start(self, port: int | None = None):
                 args += settings.server_arguments.split(" ")
             if port is not None:
                 args += ["--port", str(port)]
-            self._process = await create_process(self._python_cmd, *args, cwd=self.comfy_dir)
+            if self.backend is not ServerBackend.cpu:
+                env["ONEDNN_MAX_CPU_ISA"] = "AVX2"  # workaround for #401
+            self._process = await create_process(
+                self._python_cmd, *args, cwd=self.comfy_dir, additional_env=env
+            )
 
             assert self._process.stdout is not None
             async for line in self._process.stdout:
@@ -627,8 +632,8 @@ def _parse_common_errors(output: str, return_code: int | None):
     if "error while attempting to bind on address" in output:
         message_part = output.split("bind on address")[-1].strip()
         return (
-            f"Could not bind on address {message_part}.<br>"
-            + "<a href='https://github.com/Acly/krita-ai-diffusion/wiki/Common-Issues#error-during-server-startup-could-not-bind-on-address-only-one-usage-of-each-socket-address-is-normally-permitted'>Read more...</a>"
+            f"Could not bind on address {message_part}. "
+            + "<a href='https://github.com/Acly/krita-ai-diffusion/wiki/Common-Issues#error-during-server-startup-could-not-bind-on-address-only-one-usage-of-each-socket-address-is-normally-permitted'>More information...</a>"
         )
 
     nvidia_driver = "Found no NVIDIA driver on your system"

diff --git a/ai_diffusion/util.py b/ai_diffusion/util.py
@@ -155,7 +155,9 @@ def set_pdeathsig():
         return libc.prctl(1, signal.SIGTERM)
 
 
-async def create_process(program: str | Path, *args: str, cwd: Path | None = None):
+async def create_process(
+    program: str | Path, *args: str, cwd: Path | None = None, additional_env: dict | None = None
+):
     PIPE = asyncio.subprocess.PIPE
 
     platform_args = {}
@@ -165,6 +167,8 @@ async def create_process(program: str | Path, *args: str, cwd: Path | None = Non
         platform_args["preexec_fn"] = set_pdeathsig
 
     env = os.environ.copy()
+    if additional_env:
+        env.update(additional_env)
     if "PYTHONPATH" in env:
         del env["PYTHONPATH"]  # Krita adds its own python path, which can cause conflicts