ROCm fix memory exceptions

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

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