Less useless downcasting.

comfyanonymous · Dec 4, 2023 · be3468d · be3468d
1 parent ca82ade
commit be3468d
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py
@@ -84,12 +84,16 @@ def __init__(self, version="openai/clip-vit-large-patch14", device="cpu", max_le
 
         self.inner_name = inner_name
         if dtype is not None:
-            self.transformer.to(dtype)
             inner_model = getattr(self.transformer, self.inner_name)
             if hasattr(inner_model, "embeddings"):
-                inner_model.embeddings.to(torch.float32)
+                embeddings_bak = inner_model.embeddings.to(torch.float32)
+                inner_model.embeddings = None
+                self.transformer.to(dtype)
+                inner_model.embeddings = embeddings_bak
             else:
-                self.transformer.set_input_embeddings(self.transformer.get_input_embeddings().to(torch.float32))
+                previous_inputs = self.transformer.get_input_embeddings().to(torch.float32, copy=True)
+                self.transformer.to(dtype)
+                self.transformer.set_input_embeddings(previous_inputs)
 
         self.max_length = max_length
         if freeze: