1 files changed, 8 insertions, 0 deletions
diff --git a/modules/devices.py b/modules/devices.py
index 6b36622c..0100e4af 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -83,6 +83,14 @@ dtype_unet = torch.float16
 unet_needs_upcast = False
 
 
+def cond_cast_unet(input):
+    return input.to(dtype_unet) if unet_needs_upcast else input
+
+
+def cond_cast_float(input):
+    return input.float() if unet_needs_upcast else input
+
+
 def randn(seed, shape):
     torch.manual_seed(seed)
     if device.type == 'mps':