From 331c2f43a0c8c0f7a03b35fa002459649a4b8f1c Mon Sep 17 00:00:00 2001
From: nullquant <merlinkun@yandex.ru>
Date: Fri, 3 May 2024 19:31:08 +0300
Subject: [PATCH] timesteps device

---
 brushnet/brushnet.py | 5 ++++-
 brushnet_nodes.py    | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/brushnet/brushnet.py b/brushnet/brushnet.py
index 5a69151..43888fb 100644
--- a/brushnet/brushnet.py
+++ b/brushnet/brushnet.py
@@ -750,6 +750,9 @@ class BrushNetModel(ModelMixin, ConfigMixin):
 
         if timestep is None and time_emb is None:
             raise ValueError(f"`timestep` and `emb` are both None")
+        
+        print("BN: sample.device", sample.device)
+        print("BN: TE.device", self.time_embedding.linear_1.weight.device)
 
         if timestep is not None:
             # 1. time
@@ -776,7 +779,7 @@ class BrushNetModel(ModelMixin, ConfigMixin):
             # there might be better ways to encapsulate this.
             t_emb = t_emb.to(dtype=sample.dtype)
 
-            print("BN t_emb.device =",t_emb.device)
+            print("t_emb.device =",t_emb.device)
 
             emb = self.time_embedding(t_emb, timestep_cond)
             aug_emb = None
diff --git a/brushnet_nodes.py b/brushnet_nodes.py
index 5d1457b..8567a44 100644
--- a/brushnet_nodes.py
+++ b/brushnet_nodes.py
@@ -87,7 +87,7 @@ class BrushNetLoader:
         brushnet_model = load_checkpoint_and_dispatch(
             brushnet_model,
             brushnet_file,
-            device_map="auto",
+            device_map="sequential",
             max_memory=None,
             offload_folder=None,
             offload_state_dict=False,