bfloat16

Author: JonathannJacobs

predict.py

@@ -33,7 +33,7 @@ def setup(self) -> None:
         self.pipe = FluxControlNetPipeline.from_pretrained(
             MODEL_CACHE,
             controlnet=controlnet,
-            torch_dtype=torch.float16
+            torch_dtype=torch.bfloat16
         ).to("cuda")
 
         # Quantize transformer
@@ -62,7 +62,7 @@ def setup(self) -> None:
                 self.inpaint_pipe.transformer
             ]:
                 if hasattr(component, "to"):
-                    component.to(dtype=torch.float16)
+                    component.to(dtype=torch.bfloat16)
 
     def predict(
         self,