diff --git a/CHANGELOG.md b/CHANGELOG.md index 506d2df5a..4073b1809 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -68,6 +68,7 @@ But there's more than SD3: - fix sdxl "has been incorrectly initialized" - fix api face-hires - fix api ip-adapter +- fix memory exceptions with ROCm, thanks @Disty0! - cleanup image metadata - restructure api examples: `cli/api-*` - handle theme fallback when invalid theme is specified diff --git a/modules/sd_vae_approx.py b/modules/sd_vae_approx.py index e66b78011..1e5984145 100644 --- a/modules/sd_vae_approx.py +++ b/modules/sd_vae_approx.py @@ -34,19 +34,22 @@ class VAEApprox(nn.Module): def nn_approximation(sample): # Approximate NN global sd_vae_approx_model # pylint: disable=global-statement + # ROCm throws memory exceptions and crashes the GPU with it if we use approx on the GPU + device = devices.device if devices.backend != "rocm" else "cpu" + dtype = devices.dtype_vae if devices.backend != "rocm" else torch.float32 if sd_vae_approx_model is None: model_path = os.path.join(paths.models_path, "VAE-approx", "model.pt") sd_vae_approx_model = VAEApprox() if not os.path.exists(model_path): model_path = os.path.join(paths.script_path, "models", "VAE-approx", "model.pt") - approx_weights = torch.load(model_path, map_location='cpu' if devices.device.type != 'cuda' else None) + approx_weights = torch.load(model_path, map_location='cpu' if devices.device.type != 'cuda' or devices.backend == "rocm" else None) sd_vae_approx_model.load_state_dict(approx_weights) sd_vae_approx_model.eval() - sd_vae_approx_model.to(devices.device, sample.dtype) + sd_vae_approx_model.to(device, dtype) shared.log.debug(f'VAE load: type=approximate model={model_path}') try: - in_sample = sample.to(devices.device).unsqueeze(0) - sd_vae_approx_model.to(devices.device, devices.dtype) + in_sample = sample.to(device, dtype).unsqueeze(0) + sd_vae_approx_model.to(device, dtype) x_sample = sd_vae_approx_model(in_sample) x_sample = x_sample[0].detach().cpu() return x_sample