diff --git a/backend/memory_management.py b/backend/memory_management.py
index 746711cb..900731e5 100644
--- a/backend/memory_management.py
+++ b/backend/memory_management.py
@@ -827,7 +827,7 @@ def unet_dtype(device: torch.device = None, model_params: int = 0, supported_dty
     return torch.float32
 
 
-def inference_cast(weight_dtype: torch.device, inference_device: torch.device, supported_dtypes: list[torch.dtype] = [torch.float16, torch.bfloat16, torch.float32]) -> torch.dtype:
+def inference_cast(weight_dtype: torch.dtype, inference_device: torch.device, supported_dtypes: list[torch.dtype] = [torch.float16, torch.bfloat16, torch.float32]) -> torch.dtype:
     if weight_dtype == torch.float32:
         return weight_dtype
 
diff --git a/modules/devices.py b/modules/devices.py
index 3227094c..df003a58 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -1,101 +1,66 @@
 import contextlib
+
 import torch
+
 from backend import memory_management
 
+cpu: torch.device = torch.device("cpu")
+fp8: bool = False
+device: torch.device = memory_management.get_torch_device()
+device_gfpgan = device_esrgan = device_codeformer = device
+dtype_vae: torch.dtype = memory_management.vae_dtype()
+dtype_unet: torch.dtype = memory_management.unet_dtype()
+dtype_inference: torch.dtype = memory_management.inference_cast(dtype_unet, device)
+dtype: torch.dtype = torch.float32 if dtype_unet is torch.float32 else torch.float16
+unet_needs_upcast: bool = False
+
 
 def has_xpu() -> bool:
-    return memory_management.xpu_available
+    return memory_management.is_device_xpu(device)
 
 
 def has_mps() -> bool:
-    return memory_management.mps_mode()
+    return memory_management.is_device_mps(device)
 
 
-def cuda_no_autocast(device_id=None) -> bool:
-    return False
+def get_cuda_device_id() -> int:
+    return device.index
 
 
-def get_cuda_device_id():
-    return memory_management.get_torch_device().index
+def get_cuda_device_string() -> str:
+    return str(device)
 
 
-def get_cuda_device_string():
-    return str(memory_management.get_torch_device())
+def get_optimal_device_name() -> str:
+    return device.type
 
 
-def get_optimal_device_name():
-    return memory_management.get_torch_device().type
+def get_optimal_device() -> torch.device:
+    return device
 
 
-def get_optimal_device():
-    return memory_management.get_torch_device()
-
-
-def get_device_for(task):
-    return memory_management.get_torch_device()
+def get_device_for(*args, **kwargs) -> torch.device:
+    return device
 
 
 def torch_gc():
     memory_management.soft_empty_cache()
 
 
-def torch_npu_set_device():
-    return
-
-
-def enable_tf32():
-    return
-
-
-cpu: torch.device = torch.device("cpu")
-fp8: bool = False
-device: torch.device = memory_management.get_torch_device()
-device_gfpgan: torch.device = memory_management.get_torch_device()  # will be managed by memory management system
-device_esrgan: torch.device = memory_management.get_torch_device()  # will be managed by memory management system
-device_codeformer: torch.device = memory_management.get_torch_device()  # will be managed by memory management system
-dtype: torch.dtype = torch.float32 if memory_management.unet_dtype() is torch.float32 else torch.float16
-dtype_vae: torch.dtype = memory_management.vae_dtype()
-dtype_unet: torch.dtype = memory_management.unet_dtype()
-dtype_inference: torch.dtype = memory_management.unet_dtype()
-unet_needs_upcast = False
-
-
-def cond_cast_unet(input):
-    return input
-
-
-def cond_cast_float(input):
-    return input
-
-
-nv_rng = None
-patch_module_list = []
-
-
-def manual_cast_forward(target_dtype):
-    return
-
-
-@contextlib.contextmanager
-def manual_cast(target_dtype):
-    return
-
-
-def autocast(disable=False):
+def autocast(*args, **kwargs):
     return contextlib.nullcontext()
 
 
-def without_autocast(disable=False):
+def without_autocast(*args, **kwargs):
     return contextlib.nullcontext()
 
 
-class NansException(Exception):
-    pass
+# class NansException(Exception):
+#     pass
 
 
-def test_for_nans(x, where):
-    return
-
-
-def first_time_calculation():
-    return
+def test_for_nans(x: torch.Tensor, *args, **kwargs):
+    if torch.isnan(x).any():
+        memory_management.logger.warning("Encountered NaN in Latent" + "; Try --disable-sage" if memory_management.sage_enabled() else "")
+        x.nan_to_num_(nan=0.0, posinf=1.0, neginf=0.0)
+        # raise NansException
diff --git a/modules/processing.py b/modules/processing.py
index a8ce9831..10eff1c9 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -1036,9 +1036,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
 
             for i, x_sample in enumerate(x_samples_ddim):
                 p.batch_index = i
-                if torch.isnan(x_sample).any():
-                    logger.warning("Encountered NaN in Latent\nIf you are using SageAttention, try --disable-sage")
-                    x_sample.nan_to_num_(nan=0.0, posinf=1.0, neginf=0.0)
                 x_sample = 255.0 * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
                 x_sample = x_sample.astype(np.uint8)
                 if _is_video: