From 69b1d832e64436bbee1c273b541040f75b962d01 Mon Sep 17 00:00:00 2001
From: Acly <aclysia@gmail.com>
Date: Mon, 26 Feb 2024 23:17:33 +0100
Subject: [PATCH] Set ONEDNN_MAX_CPU_ISA=AVX2 as potential workaround for #401

---
 ai_diffusion/server.py | 11 ++++++++---
 ai_diffusion/util.py   |  6 +++++-
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/ai_diffusion/server.py b/ai_diffusion/server.py
index ab363851a9..e2b8aaa41d 100644
--- 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
index 2afc1db174..41f9e36ac2 100644
--- 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