From 6c6ae28bf5fd1e8bc3e8f64a3430b6f29f338f77 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Tue, 4 Oct 2022 12:32:22 +0300
Subject: send all three of GFPGAN's and codeformer's models to CPU memory
 instead of just one for #1283

---
 modules/devices.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

(limited to 'modules/devices.py')

diff --git a/modules/devices.py b/modules/devices.py
index ff82f2f6..12aab665 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -1,3 +1,5 @@
+import contextlib
+
 import torch
 
 # has_mps is only available in nightly pytorch (for now), `getattr` for compatibility
@@ -57,3 +59,11 @@ def randn_without_seed(shape):
 
     return torch.randn(shape, device=device)
 
+
+def autocast():
+    from modules import shared
+
+    if dtype == torch.float32 or shared.cmd_opts.precision == "full":
+        return contextlib.nullcontext()
+
+    return torch.autocast("cuda")
-- 
cgit v1.2.3