From f52c45800cf48d28fbb00dc9666891b032a6712e Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 7 Nov 2024 19:42:59 +0200
Subject: [PATCH] allow 2nd GPU as FasterCache device

experimental
---
 nodes.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index 9f613d9..f386156 100644
--- a/nodes.py
+++ b/nodes.py
@@ -755,7 +755,7 @@ class CogVideoXFasterCache:
                 "start_step": ("INT", {"default": 15, "min": 0, "max": 1024, "step": 1}),
                 "hf_step": ("INT", {"default": 30, "min": 0, "max": 1024, "step": 1}),
                 "lf_step": ("INT", {"default": 40, "min": 0, "max": 1024, "step": 1}),
-                "cache_device": (["main_device", "offload_device"], {"default": "main_device", "tooltip": "The device to use for the cache, main_device is on GPU and uses a lot of VRAM"}),
+                "cache_device": (["main_device", "offload_device", "cuda:1"], {"default": "main_device", "tooltip": "The device to use for the cache, main_device is on GPU and uses a lot of VRAM"}),
             },
         }
 
@@ -767,11 +767,13 @@ class CogVideoXFasterCache:
     def args(self, start_step, hf_step, lf_step, cache_device):
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
+        if cache_device == "cuda:1":
+            device = torch.device("cuda:1")
         fastercache = {
             "start_step" : start_step,
             "hf_step" : hf_step,
             "lf_step" : lf_step,
-            "cache_device" : device if cache_device == "main_device" else offload_device
+            "cache_device" : device if cache_device != "offload_device" else offload_device
         }
         return (fastercache,)