From bf42e687fddc120f9ca81e0f17d9a5e5bc9a6015 Mon Sep 17 00:00:00 2001
From: doctorpangloss <@hiddenswitch.com>
Date: Wed, 3 Jan 2024 14:29:16 -0800
Subject: [PATCH] Fix more nodes

---
 comfy_extras/nodes/nodes_upscale_model.py |  2 +-
 setup.py                                  | 20 ++++----------------
 2 files changed, 5 insertions(+), 17 deletions(-)

diff --git a/comfy_extras/nodes/nodes_upscale_model.py b/comfy_extras/nodes/nodes_upscale_model.py
index dd63f1ae6..5678d428b 100644
--- a/comfy_extras/nodes/nodes_upscale_model.py
+++ b/comfy_extras/nodes/nodes_upscale_model.py
@@ -19,7 +19,7 @@ class UpscaleModelLoader:
         model_path = folder_paths.get_full_path("upscale_models", model_name)
         sd = utils.load_torch_file(model_path, safe_load=True)
         if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
-            sd = comfy.utils.state_dict_prefix_replace(sd, {"module.":""})
+            sd = utils.state_dict_prefix_replace(sd, {"module.":""})
         out = model_loading.load_state_dict(sd).eval()
         return (out, )
 
diff --git a/setup.py b/setup.py
index f73ae4bea..e69d6146c 100644
--- a/setup.py
+++ b/setup.py
@@ -28,23 +28,16 @@ version = '0.0.1'
 """
 The package index to the torch built with AMD ROCm.
 """
-amd_torch_index = "https://download.pytorch.org/whl/rocm5.4.2"
+amd_torch_index = "https://download.pytorch.org/whl/rocm5.6"
 
 """
 The package index to torch built with CUDA.
 Observe the CUDA version is in this URL.
 """
-nvidia_torch_index = "https://download.pytorch.org/whl/cu118"
+nvidia_torch_index = "https://download.pytorch.org/whl/cu121"
 
 """
 The package index to torch built against CPU features.
-This includes macOS MPS support.
-"""
-cpu_torch_index_nightlies = "https://download.pytorch.org/whl/nightly/cpu"
-
-"""
-The package index to torch built against CPU features.
-Non-nightlies are selected when building Linux on arm64.
 """
 cpu_torch_index = "https://download.pytorch.org/whl/cpu"
 
@@ -110,22 +103,17 @@ def _is_linux_arm64():
 def dependencies() -> List[str]:
     _dependencies = open(os.path.join(os.path.dirname(__file__), "requirements.txt")).readlines()
     # todo: also add all plugin dependencies
-    _alternative_indices = [amd_torch_index, nvidia_torch_index, cpu_torch_index_nightlies]
+    _alternative_indices = [amd_torch_index, nvidia_torch_index]
     session = PipSession()
 
-    gpu_accelerated = False
     index_urls = ['https://pypi.org/simple']
     # prefer nvidia over AMD because AM5/iGPU systems will have a valid ROCm device
     if _is_nvidia():
         index_urls += [nvidia_torch_index]
-        gpu_accelerated = True
     elif _is_amd():
         index_urls += [amd_torch_index]
-        gpu_accelerated = True
-    elif _is_linux_arm64():
-        index_urls += [cpu_torch_index]
     else:
-        index_urls += [cpu_torch_index_nightlies]
+        index_urls += [cpu_torch_index]
 
     if len(index_urls) == 1:
         return _dependencies