make default init in torch.bfloat16

models/pipeline.py

@@ -13,11 +13,13 @@ class SwittiPipeline:
     text_encoder_path = "openai/clip-vit-large-patch14"
     text_encoder_2_path = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
 
-    def __init__(self, switti, vae, text_encoder, text_encoder_2, device):
-        self.switti = switti
-        self.vae = vae
-        self.text_encoder = text_encoder
-        self.text_encoder_2 = text_encoder_2
+    def __init__(self, switti, vae, text_encoder, text_encoder_2, device,
+                dtype=torch.bfloat16,
+                ):
+        self.switti = switti.to(dtype)
+        self.vae = vae.to(dtype)
+        self.text_encoder = text_encoder.to(dtype)
+        self.text_encoder_2 = text_encoder_2.to(dtype)
 
         self.switti.eval()
         self.vae.eval()