diff --git a/environment-mac.yaml b/environment-mac.yaml index be63a055402..42d2d5eaaf5 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/dream/devices.py b/ldm/dream/devices.py index 240754dd364..9581abe78cf 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 230aa74c284..ccecee1a46a 100644 --- a/ldm/simplet2i.py +++ b/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() diff --git a/scripts/dream.py b/scripts/dream.py index 2911e8847ab..1d6f4d29241 100755 --- a/scripts/dream.py +++ b/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',