From f204d6fcec65bdf9ea9cd6ca65e559f692bce711 Mon Sep 17 00:00:00 2001
From: Jedrzej Kosinski <kosinkadink1@gmail.com>
Date: Tue, 18 Feb 2025 22:58:57 -0600
Subject: [PATCH] Added initial support for upcoming ComfyUI MultiGPU Work
 Units feature

---
 CrossAttentionPatch.py | 39 ++++++++++++++++++++++++++++++++++++++-
 IPAdapterPlus.py       | 17 +++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/CrossAttentionPatch.py b/CrossAttentionPatch.py
index 0f5c8e9..adb2079 100644
--- a/CrossAttentionPatch.py
+++ b/CrossAttentionPatch.py
@@ -8,6 +8,10 @@ class Attn2Replace:
     def __init__(self, callback=None, **kwargs):
         self.callback = [callback]
         self.kwargs = [kwargs]
+        self.multigpu_kwargs = {}
+
+    def get_multigpu_kwargs(self, device):
+        return self.multigpu_kwargs.get(device, self.kwargs)
 
     def add(self, callback, **kwargs):
         self.callback.append(callback)
@@ -21,13 +25,46 @@ class Attn2Replace:
         out = optimized_attention(q, k, v, extra_options["n_heads"])
         sigma = extra_options["sigmas"].detach().cpu()[0].item() if 'sigmas' in extra_options else 999999999.9
 
+        device_kwargs = self.get_multigpu_kwargs(q.device)
+
         for i, callback in enumerate(self.callback):
             if sigma <= self.kwargs[i]["sigma_start"] and sigma >= self.kwargs[i]["sigma_end"]:
-                out = out + callback(out, q, k, v, extra_options, **self.kwargs[i])
+                out = out + callback(out, q, k, v, extra_options, **device_kwargs[i])
 
         return out.to(dtype=dtype)
+    
+    def to(self, device, *args, **kwargs):
+        if not isinstance(device, torch.device):
+            return self
+        # if casted versions already taken care of, do nothing
+        if device in self.multigpu_kwargs:
+            return self
+        # ignore casts to CPU, as IPAdapter currently always sticks to load device(s)
+        if device == "cpu" or device == torch.device("cpu"):
+            return self
+        # for now, assume cond must be present
+        if self.kwargs[0]["cond"] is None:
+            return self
+        if self.kwargs[0]["cond"].device == device:
+            return self
+        
+        new_kwargs = []
+        for kwargs_dict in self.kwargs:
+            new_kwargs.append(kwargs_dict.copy())
+        
+        for kwargs_dict in new_kwargs:
+            for key, value in kwargs_dict.items():
+                if key == "ipadapter":
+                    value.create_multigpu_clone(device)
+                elif type(kwargs_dict[key]) == torch.Tensor:
+                    kwargs_dict[key] = value.to(device)
+
+        self.multigpu_kwargs[device] = new_kwargs
+        return self
 
 def ipadapter_attention(out, q, k, v, extra_options, module_key='', ipadapter=None, weight=1.0, cond=None, cond_alt=None, uncond=None, weight_type="linear", mask=None, sigma_start=0.0, sigma_end=1.0, unfold_batch=False, embeds_scaling='V only', **kwargs):
+    ipadapter = ipadapter.get_multigpu_clone(q.device)
+    
     dtype = q.dtype
     cond_or_uncond = extra_options["cond_or_uncond"]
     block_type = extra_options["block"][0]
diff --git a/IPAdapterPlus.py b/IPAdapterPlus.py
index ce3517a..2168904 100644
--- a/IPAdapterPlus.py
+++ b/IPAdapterPlus.py
@@ -2,6 +2,7 @@ import torch
 import os
 import math
 import folder_paths
+import copy
 
 import comfy.model_management as model_management
 from node_helpers import conditioning_set_values
@@ -71,6 +72,22 @@ class IPAdapter(nn.Module):
         self.image_proj_model.load_state_dict(ipadapter_model["image_proj"])
         self.ip_layers = To_KV(ipadapter_model["ip_adapter"], encoder_hid_proj=encoder_hid_proj, weight_kolors=weight_kolors)
 
+        self.multigpu_clones = {}
+
+    def create_multigpu_clone(self, device):
+        if device not in self.multigpu_clones:
+            orig_multigpu_clones = self.multigpu_clones
+            try:
+                self.multigpu_clones = {}
+                new_clone = copy.deepcopy(self)
+                new_clone = new_clone.to(device)
+                orig_multigpu_clones[device] = new_clone
+            finally:
+                self.multigpu_clones = orig_multigpu_clones
+
+    def get_multigpu_clone(self, device):
+        return self.multigpu_clones.get(device, self)
+
     def init_proj(self):
         image_proj_model = ImageProjModel(
             cross_attention_dim=self.cross_attention_dim,