change to float16 model dtype

2025-05-16 18:02:48 +02:00
parent 9b425b6bfa
commit 048db5c12f
1 changed files with 2 additions and 0 deletions
@@ -129,6 +129,8 @@ class ACEStepPipeline:
            self.dtype = torch.float16
        if device.type == "mps":
            self.dtype = torch.float32
+        if len(os.environ['COLAB_FREE']):
+            self.dtype = torch.float16
        self.device = device
        self.loaded = False
        self.torch_compile = torch_compile