Add Mac OS support for Apple Silicon (#770)

* Made an install script and auto updates env for mac

* GPU sensors and initial training working for MAC. Still WIP.

* Switch dataloader to single threaded until I can work around some mac pickeling issues.

* Get quantization working on mac

* Fix mac exclusive imports so they don't break other builds.

* Add mac instructions to the UI

toolkit/basic.py

@@ -9,7 +9,11 @@ def value_map(inputs, min_in, max_in, min_out, max_out):
 
 
 def flush(garbage_collect=True):
-    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    # if is mps, also clear the mps cache
+    if torch.backends.mps.is_available():
+        torch.mps.empty_cache()
     if garbage_collect:
         gc.collect()
 

toolkit/config_modules.py

@@ -7,6 +7,7 @@ import torch
 import torchaudio
 
 from toolkit.prompt_utils import PromptEmbeds
+from torchao.quantization.quant_primitives import _DTYPE_TO_BIT_WIDTH
 
 ImgExt = Literal['jpg', 'png', 'webp']
 
@@ -668,6 +669,12 @@ class ModelConfig:
             self.qtype = "float8"
         if self.layer_offloading and self.qtype_te == "qfloat8":
             self.qtype_te = "float8"
+            
+        # Mac mps only works with torachao uint
+        if torch.backends.mps.is_available() and self.qtype == "qfloat8":
+            self.qtype = "int8"
+        if torch.backends.mps.is_available() and self.qtype_te == "qfloat8":
+            self.qtype_te = "int8"
         
         # 0 is off and 1.0 is 100% of the layers
         self.layer_offloading_transformer_percent = kwargs.get("layer_offloading_transformer_percent", 1.0)

toolkit/data_loader.py

@@ -29,6 +29,9 @@ import platform
 def is_native_windows():
     return platform.system() == "Windows" and platform.release() != "2"
 
+def is_macos():
+    return platform.system() == "Darwin"
+
 if TYPE_CHECKING:
     from toolkit.stable_diffusion_model import StableDiffusion
     
@@ -678,7 +681,7 @@ def get_dataloader_from_datasets(
 
     dataloader_kwargs = {}
     
-    if is_native_windows():
+    if is_native_windows() or is_macos():
         dataloader_kwargs['num_workers'] = 0
     else:
         dataloader_kwargs['num_workers'] = dataset_config_list[0].num_workers

toolkit/models/base_model.py

@@ -41,6 +41,7 @@ from torchvision.transforms import functional as TF
 from toolkit.accelerator import get_accelerator, unwrap_model
 from typing import TYPE_CHECKING
 from toolkit.print import print_acc
+from toolkit.basic import flush
 
 if TYPE_CHECKING:
     from toolkit.lora_special import LoRASpecialNetwork
@@ -90,11 +91,6 @@ class BlankNetwork:
         pass
 
 
-def flush():
-    torch.cuda.empty_cache()
-    gc.collect()
-
-
 UNET_IN_CHANNELS = 4  # Stable Diffusion の in_channels は 4 で固定。XLも同じ。
 # VAE_SCALE_FACTOR = 8  # 2 ** (len(vae.config.block_out_channels) - 1) = 8
 

toolkit/stable_diffusion_model.py

@@ -70,6 +70,7 @@ from typing import TYPE_CHECKING
 from toolkit.print import print_acc
 from diffusers import FluxFillPipeline
 from transformers import AutoModel, AutoTokenizer, Gemma2Model, Qwen2Model, LlamaModel
+from toolkit.basic import flush
 
 if TYPE_CHECKING:
     from toolkit.lora_special import LoRASpecialNetwork
@@ -118,11 +119,6 @@ class BlankNetwork:
         pass
 
 
-def flush():
-    torch.cuda.empty_cache()
-    gc.collect()
-
-
 UNET_IN_CHANNELS = 4  # Stable Diffusion の in_channels は 4 で固定。XLも同じ。
 # VAE_SCALE_FACTOR = 8  # 2 ** (len(vae.config.block_out_channels) - 1) = 8
 

toolkit/util/quantize.py

@@ -8,6 +8,7 @@ from torchao.quantization.quant_api import (
     quantize_ as torchao_quantize_,
     Float8WeightOnlyConfig,
     UIntXWeightOnlyConfig,
+    Int8WeightOnlyConfig
 )
 from optimum.quanto import freeze
 from tqdm import tqdm
@@ -41,6 +42,7 @@ torchao_qtypes = {
     "uint6": UIntXWeightOnlyConfig(torch.uint6),
     "uint7": UIntXWeightOnlyConfig(torch.uint7),
     "uint8": UIntXWeightOnlyConfig(torch.uint8),
+    "int8": Int8WeightOnlyConfig(),
     "float8": Float8WeightOnlyConfig(),
 }