Enhance token handling in FrozenOpenCLIPEmbedder by adding device checks and conversions. This update ensures that input tokens are correctly processed as torch tensors on the appropriate device, improving compatibility and performance across different hardware configurations.

imagedream/ldm/modules/encoders/modules.py

@@ -275,8 +275,16 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder, nn.Module):
             encoder_states = encoder_states + (x, )
         return encoder_states
     
-    def encode_with_transformer(self, text):
-        x = self.model.token_embedding(text)  # [batch_size, n_ctx, d_model]
+    def encode_with_transformer(self, tokens):
+        # Debug: print device info
+        print(f"[DEBUG] tokens type: {type(tokens)}, tokens device: {getattr(tokens, 'device', 'N/A')}")
+        print(f"[DEBUG] embedding weight device: {self.model.token_embedding.weight.device}")
+        # Ensure tokens is a torch tensor and on the correct device
+        if not isinstance(tokens, torch.Tensor):
+            tokens = torch.tensor(tokens, device=self.model.token_embedding.weight.device)
+        else:
+            tokens = tokens.to(self.model.token_embedding.weight.device)
+        x = self.model.token_embedding(tokens)  # [batch_size, n_ctx, d_model]
         x = x + self.model.positional_embedding
         x = x.permute(1, 0, 2)  # NLD -> LND
         x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)