From bf50ab9dd688c52a8727f02613b9843e2edf0a0c Mon Sep 17 00:00:00 2001 From: Jason Toffaletti Date: Wed, 31 Aug 2022 00:32:07 -0700 Subject: [PATCH 1/4] changes to get dream.py working on M1 - move all device init logic to T2I.__init__ - handle m1 specific edge case with autocast device type - check torch.cuda.is_available before using cuda --- environment-mac.yaml | 2 +- ldm/simplet2i.py | 18 ++++++++++-------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/environment-mac.yaml b/environment-mac.yaml index be63a05540..42d2d5eaaf 100644 --- a/environment-mac.yaml +++ b/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 diff --git a/ldm/simplet2i.py b/ldm/simplet2i.py index 230aa74c28..cdc72f2fba 100644 --- a/ldm/simplet2i.py +++ b/ldm/simplet2i.py @@ -154,7 +154,10 @@ class T2I: 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 @@ class T2I: 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,10 @@ class T2I: callback=step_callback, ) - with scope(self.device.type), self.model.ema_scope(): + device_type = self.device.type # this returns 'mps' on M1 + if device_type != 'cuda' or device_type != 'cpu': + device_type = 'cpu' + with scope(device_type), self.model.ema_scope(): for n in trange(iterations, desc='Generating'): seed_everything(seed) image = next(images_iterator) @@ -523,17 +530,12 @@ class T2I: 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( From 66fe11014892feb9120f4d0a3b0566ad1c77497b Mon Sep 17 00:00:00 2001 From: Jason Toffaletti Date: Wed, 31 Aug 2022 16:57:39 -0700 Subject: [PATCH 2/4] default full_prevision to True for mps device --- scripts/dream.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/scripts/dream.py b/scripts/dream.py index 2911e8847a..1d6f4d2924 100755 --- a/scripts/dream.py +++ b/scripts/dream.py @@ -9,6 +9,7 @@ import sys 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', From fa98601bfb6a26fed7b411e1a7bb727635d4b305 Mon Sep 17 00:00:00 2001 From: Jason Toffaletti Date: Wed, 31 Aug 2022 19:24:23 -0700 Subject: [PATCH 3/4] better error reporting for load_model --- ldm/simplet2i.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/ldm/simplet2i.py b/ldm/simplet2i.py index cdc72f2fba..c0d2886538 100644 --- a/ldm/simplet2i.py +++ b/ldm/simplet2i.py @@ -544,12 +544,11 @@ class T2I: 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() From 09bd9fa47eb96455d29fb575e028cfaad5b39ca7 Mon Sep 17 00:00:00 2001 From: Jason Toffaletti Date: Wed, 31 Aug 2022 22:21:14 -0700 Subject: [PATCH 4/4] move autocast device selection to a function --- ldm/dream/devices.py | 8 +++++++- ldm/simplet2i.py | 6 ++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/ldm/dream/devices.py b/ldm/dream/devices.py index 240754dd36..9581abe78c 100644 --- a/ldm/dream/devices.py +++ b/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 diff --git a/ldm/simplet2i.py b/ldm/simplet2i.py index c0d2886538..ccecee1a46 100644 --- a/ldm/simplet2i.py +++ b/ldm/simplet2i.py @@ -27,7 +27,7 @@ from ldm.models.diffusion.ddim import DDIMSampler 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 @@ -315,9 +315,7 @@ class T2I: callback=step_callback, ) - device_type = self.device.type # this returns 'mps' on M1 - if device_type != 'cuda' or device_type != 'cpu': - device_type = 'cpu' + 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)