assign devices.dtype early because it's needed before the model is loaded

1 files changed, 4 insertions, 0 deletions

diff --git a/modules/shared.py b/modules/shared.py
index 8ba1c569..7025a754 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -6,6 +6,7 @@ import threading
 import time

 

 import gradio as gr

+import torch

 import tqdm

 

 import modules.interrogate

@@ -63,6 +64,9 @@ cmd_opts.disable_extension_access = (cmd_opts.share or cmd_opts.listen or cmd_op
 devices.device, devices.device_interrogate, devices.device_gfpgan, devices.device_esrgan, devices.device_codeformer = \

     (devices.cpu if any(y in cmd_opts.use_cpu for y in [x, 'all']) else devices.get_optimal_device() for x in ['sd', 'interrogate', 'gfpgan', 'esrgan', 'codeformer'])

 

+devices.dtype = torch.float32 if cmd_opts.no_half else torch.float16

+devices.dtype_vae = torch.float32 if cmd_opts.no_half or cmd_opts.no_half_vae else torch.float16

+

 device = devices.device

 weight_load_location = None if cmd_opts.lowram else "cpu"