From 30771f2695bc07a058a4a9f5d63e374c8622cefe Mon Sep 17 00:00:00 2001 From: nullquant Date: Fri, 3 May 2024 19:00:37 +0300 Subject: [PATCH] sample & timestep device --- brushnet/brushnet.py | 2 ++ brushnet_nodes.py | 8 ++++---- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/brushnet/brushnet.py b/brushnet/brushnet.py index 705e02d..5a69151 100644 --- a/brushnet/brushnet.py +++ b/brushnet/brushnet.py @@ -776,6 +776,8 @@ 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) + emb = self.time_embedding(t_emb, timestep_cond) aug_emb = None diff --git a/brushnet_nodes.py b/brushnet_nodes.py index e9c0032..5d1457b 100644 --- a/brushnet_nodes.py +++ b/brushnet_nodes.py @@ -420,20 +420,20 @@ def brushnet_inference(x, timesteps, transformer_options): cond_scale = brushnet_conditioning_scale * brushnet_keep[step] if x.shape[0] == 2: - return brushnet(x, + return brushnet(x.to(brushnet.device), encoder_hidden_states=prompt_embeds2, brushnet_cond=conditioning_latents2, - timestep = timesteps, + timestep = timesteps.to(brushnet.device), conditioning_scale=cond_scale, guess_mode=False, added_cond_kwargs=added_cond_kwargs2, return_dict=False, ) else: - return brushnet(x, + return brushnet(x.to(brushnet.device), encoder_hidden_states=prompt_embeds, brushnet_cond=conditioning_latents, - timestep = timesteps, + timestep = timesteps.to(brushnet.device), conditioning_scale=cond_scale, guess_mode=False, added_cond_kwargs=added_cond_kwargs,