Merge pull request #386 from frantic/save-env-vars

Don't override environment variables
Fannovel16 · Jul 5, 2024 · 3c1d9de · 3c1d9de
2 parents 6f1ba1c + feb8aec
commit 3c1d9de
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/__init__.py b/__init__.py
@@ -16,7 +16,7 @@
 #Enable CPU fallback for ops not being supported by MPS like upsample_bicubic2d.out
 #https://github.com/pytorch/pytorch/issues/77764
 #https://github.com/Fannovel16/comfyui_controlnet_aux/issues/2#issuecomment-1763579485
-os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = '1' 
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = os.getenv("PYTORCH_ENABLE_MPS_FALLBACK", '1')
 
 
 def load_nodes():

diff --git a/utils.py b/utils.py
@@ -48,10 +48,10 @@
     USE_SYMLINKS = False
     ORT_PROVIDERS = ["CUDAExecutionProvider", "DirectMLExecutionProvider", "OpenVINOExecutionProvider", "ROCMExecutionProvider", "CPUExecutionProvider", "CoreMLExecutionProvider"]
 
-os.environ['AUX_ANNOTATOR_CKPTS_PATH'] = annotator_ckpts_path
-os.environ['AUX_TEMP_DIR'] = str(TEMP_DIR)
-os.environ['AUX_USE_SYMLINKS'] = str(USE_SYMLINKS)
-os.environ['AUX_ORT_PROVIDERS'] = str(",".join(ORT_PROVIDERS))
+os.environ['AUX_ANNOTATOR_CKPTS_PATH'] = os.getenv('AUX_ANNOTATOR_CKPTS_PATH', annotator_ckpts_path)
+os.environ['AUX_TEMP_DIR'] = os.getenv('AUX_TEMP_DIR', str(TEMP_DIR))
+os.environ['AUX_USE_SYMLINKS'] = os.getenv('AUX_USE_SYMLINKS', str(USE_SYMLINKS))
+os.environ['AUX_ORT_PROVIDERS'] = os.getenv('AUX_ORT_PROVIDERS', str(",".join(ORT_PROVIDERS)))
 
 log.info(f"Using ckpts path: {annotator_ckpts_path}")
 log.info(f"Using symlinks: {USE_SYMLINKS}")