remove fixed float16 COLAB_FREE env var

add ACE_PIPELINE_DTYPE env var dtype config set colab notebook to float16
2025-05-19 04:17:38 +02:00
parent e2efc727e9
commit cdab63d688
2 changed files with 3 additions and 3 deletions
@@ -129,8 +129,8 @@ class ACEStepPipeline:
            self.dtype = torch.float16
        if device.type == "mps":
            self.dtype = torch.float32
-        if 'COLAB_FREE' in os.environ and len(os.environ['COLAB_FREE']):
-            self.dtype = torch.float16
+        if 'ACE_PIPELINE_DTYPE' in os.environ and len(os.environ['ACE_PIPELINE_DTYPE']):
+            self.dtype = getattr(torch, os.environ['ACE_PIPELINE_DTYPE'])
        self.device = device
        self.loaded = False
        self.torch_compile = torch_compile