changes to get dream.py working on M1

environment-mac.yaml

@@ -28,7 +28,7 @@ dependencies:
     - kornia==0.6.0
     - -e git+https://github.com/openai/CLIP.git@main#egg=clip
     - -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
-    - -e git+https://github.com/lstein/k-diffusion.git@master#egg=k-diffusion
+    - -e git+https://github.com/Birch-san/k-diffusion.git@mps#egg=k_diffusion
     - -e .
 variables:
   PYTORCH_ENABLE_MPS_FALLBACK: 1

ldm/dream/devices.py

@@ -8,4 +8,10 @@ def choose_torch_device() -> str:
         return 'mps'
     return 'cpu'
 
-
+def choose_autocast_device(device) -> str:
+    '''Returns an autocast compatible device from a torch device'''
+    device_type = device.type # this returns 'mps' on M1
+    # autocast only supports cuda or cpu
+    if device_type != 'cuda' or device_type != 'cpu':
+        return 'cpu'
+    return device_type

ldm/simplet2i.py

@@ -27,7 +27,7 @@
 from ldm.models.diffusion.plms import PLMSSampler
 from ldm.models.diffusion.ksampler import KSampler
 from ldm.dream.pngwriter import PngWriter
-from ldm.dream.devices import choose_torch_device
+from ldm.dream.devices import choose_autocast_device, choose_torch_device
 
 """Simplified text to image API for stable diffusion/latent diffusion
 
@@ -154,7 +154,10 @@ def __init__(
         self.model = None     # empty for now
         self.sampler = None
         self.latent_diffusion_weights = latent_diffusion_weights
-        self.device = device
+        if device == 'cuda' and not torch.cuda.is_available():
+            device = choose_torch_device()
+            print("cuda not available, using device", device)
+        self.device = torch.device(device)
 
         # for VRAM usage statistics
         self.session_peakmem = torch.cuda.max_memory_allocated() if self.device == 'cuda' else None
@@ -279,7 +282,8 @@ def process_image(image,seed):
             self._set_sampler()
 
         tic = time.time()
-        torch.cuda.torch.cuda.reset_peak_memory_stats()
+        if torch.cuda.is_available():
+            torch.cuda.torch.cuda.reset_peak_memory_stats()
         results = list()
 
         try:
@@ -311,7 +315,8 @@ def process_image(image,seed):
                     callback=step_callback,
                 )
 
-            with scope(self.device.type), self.model.ema_scope():
+            device_type = choose_autocast_device(self.device)
+            with scope(device_type), self.model.ema_scope():
                 for n in trange(iterations, desc='Generating'):
                     seed_everything(seed)
                     image = next(images_iterator)
@@ -523,17 +528,12 @@ def _new_seed(self):
         self.seed = random.randrange(0, np.iinfo(np.uint32).max)
         return self.seed
 
-    def _get_device(self):
-        device_type = choose_torch_device()
-        return torch.device(device_type)
-
     def load_model(self):
         """Load and initialize the model from configuration variables passed at object creation time"""
         if self.model is None:
             seed_everything(self.seed)
             try:
                 config = OmegaConf.load(self.config)
-                self.device = self._get_device()
                 model = self._load_model_from_config(config, self.weights)
                 if self.embedding_path is not None:
                     model.embedding_manager.load(
@@ -542,12 +542,11 @@ def load_model(self):
                 self.model = model.to(self.device)
                 # model.to doesn't change the cond_stage_model.device used to move the tokenizer output, so set it here
                 self.model.cond_stage_model.device = self.device
-            except AttributeError:
+            except AttributeError as e:
                 import traceback
-                print(
-                    'Error loading model. Only the CUDA backend is supported', file=sys.stderr)
+                print(f'Error loading model. {str(e)}', file=sys.stderr)
                 print(traceback.format_exc(), file=sys.stderr)
-                raise SystemExit
+                raise SystemExit from e
 
             self._set_sampler()
 

scripts/dream.py

@@ -9,6 +9,7 @@
 import copy
 import warnings
 import time
+from ldm.dream.devices import choose_torch_device
 import ldm.dream.readline
 from ldm.dream.pngwriter import PngWriter, PromptFormatter
 from ldm.dream.server import DreamServer, ThreadingDreamServer
@@ -347,6 +348,8 @@ def create_argv_parser():
         dest='full_precision',
         action='store_true',
         help='Use slower full precision math for calculations',
+        # MPS only functions with full precision, see https://github.com/lstein/stable-diffusion/issues/237
+        default=choose_torch_device() == 'mps',
     )
     parser.add_argument(
         '-g',