diff --git a/comfy/sd.py b/comfy/sd.py index eb8379cc7..b96366e72 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -577,7 +577,9 @@ class VAE: if pixel_samples is None: pixel_samples = torch.empty((samples_in.shape[0],) + tuple(out.shape[1:]), device=self.output_device) pixel_samples[x:x+batch_number] = out - except model_management.OOM_EXCEPTION: + except Exception as ex: + if not model_management.is_oom_exception(ex): + raise logging.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.") dims = samples_in.ndim - 2 if dims == 1 or self.extra_1d_channel is not None: