From cdab63d688a7fcc4caa9b04171cd3b3c1600b386 Mon Sep 17 00:00:00 2001 From: Michael Hedman Date: Mon, 19 May 2025 04:17:38 +0200 Subject: [PATCH] remove fixed float16 COLAB_FREE env var add ACE_PIPELINE_DTYPE env var dtype config set colab notebook to float16 --- acestep/pipeline_ace_step.py | 4 ++-- colab_inference.ipynb | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/acestep/pipeline_ace_step.py b/acestep/pipeline_ace_step.py index 0f6c6ef..080040e 100644 --- a/acestep/pipeline_ace_step.py +++ b/acestep/pipeline_ace_step.py @@ -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 diff --git a/colab_inference.ipynb b/colab_inference.ipynb index 53c466c..4f787a1 100644 --- a/colab_inference.ipynb +++ b/colab_inference.ipynb @@ -47,7 +47,7 @@ "#!pip uninstall ace-step -y\n", "!pip install --upgrade git+https://github.com/craftmine1000/ACE-Step-Colab-Opt.git\n", "import os\n", - "os.environ['COLAB_FREE'] = 'TRUE'" + "os.environ['ACE_PIPELINE_DTYPE'] = 'float16'" ] }, {