From 2819b48e9838c8d704f3ca2ebedfeba2647dfcd6 Mon Sep 17 00:00:00 2001
From: Badis <110173477+BadisG@users.noreply.github.com>
Date: Mon, 25 Nov 2024 10:13:12 +0100
Subject: [PATCH] Fix the start_percent > 0 error

---
 flux/layers.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/flux/layers.py b/flux/layers.py
index 4fbd6bf..4274387 100644
--- a/flux/layers.py
+++ b/flux/layers.py
@@ -57,8 +57,9 @@ class DoubleStreamBlockIPA(nn.Module):
 
         txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
 
-        ip_hidden_states = self.ip_adapter(self.num_heads, img_q, self.image_emb, t).to(self.device)
+        ip_hidden_states = self.ip_adapter(self.num_heads, img_q, self.image_emb, t)
         if ip_hidden_states is not None:
+            ip_hidden_states = ip_hidden_states.to(self.device)
             img_attn = img_attn + ip_hidden_states
 
         # calculate the img bloks
@@ -115,12 +116,13 @@ class SingleStreamBlockIPA(nn.Module):
         # compute attention
         attn = attention(q, k, v, pe=pe)
 
-        ip_hidden_states = self.ip_adapter(self.num_heads, q, self.image_emb, t).to(self.device)
+        ip_hidden_states = self.ip_adapter(self.num_heads, q, self.image_emb, t)
         if ip_hidden_states is not None:
+            ip_hidden_states = ip_hidden_states.to(self.device)
             attn = attn + ip_hidden_states
         # compute activation in mlp stream, cat again and run second linear layer
         output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
         x += mod.gate * output
         if x.dtype == torch.float16:
             x = torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
-        return x
\ No newline at end of file
+        return x