refactor devices

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/postprocess/codeformer_model.py

@@ -41,7 +41,7 @@ def setup_model(dirname):
                 from facelib.utils.face_restoration_helper import FaceRestoreHelper
                 from facelib.detection.retinaface import retinaface
                 if self.net is not None and self.face_helper is not None:
-                    self.net.to(devices.device_codeformer)
+                    self.net.to(devices.device)
                     return self.net, self.face_helper
                 model_paths = modelloader.load_models(model_path, model_url, self.cmd_dir, download_name='codeformer-v0.1.0.pth', ext_filter=['.pth'])
                 if len(model_paths) != 0:
@@ -49,14 +49,14 @@ def setup_model(dirname):
                 else:
                     shared.log.error(f"Model failed loading: type=CodeFormer model={model_path}")
                     return None, None
-                net = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(devices.device_codeformer)
+                net = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(devices.device)
                 checkpoint = torch.load(ckpt_path)['params_ema']
                 net.load_state_dict(checkpoint)
                 net.eval()
                 shared.log.info(f"Model loaded: type=CodeFormer model={ckpt_path}")
                 if hasattr(retinaface, 'device'):
-                    retinaface.device = devices.device_codeformer
-                face_helper = FaceRestoreHelper(1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', use_parse=True, device=devices.device_codeformer)
+                    retinaface.device = devices.device
+                face_helper = FaceRestoreHelper(1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', use_parse=True, device=devices.device)
                 self.net = net
                 self.face_helper = face_helper
                 return net, face_helper
@@ -74,7 +74,7 @@ def setup_model(dirname):
                 self.create_models()
                 if self.net is None or self.face_helper is None:
                     return np_image
-                self.send_model_to(devices.device_codeformer)
+                self.send_model_to(devices.device)
                 self.face_helper.clean_all()
                 self.face_helper.read_image(np_image)
                 self.face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)
@@ -82,7 +82,7 @@ def setup_model(dirname):
                 for cropped_face in self.face_helper.cropped_faces:
                     cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
                     normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
-                    cropped_face_t = cropped_face_t.unsqueeze(0).to(devices.device_codeformer)
+                    cropped_face_t = cropped_face_t.unsqueeze(0).to(devices.device)
                     try:
                         with devices.inference_context():
                             output = self.net(cropped_face_t, w=w if w is not None else shared.opts.code_former_weight, adain=True)[0] # pylint: disable=not-callable

modules/postprocess/esrgan_model.py

@@ -128,7 +128,7 @@ class UpscalerESRGAN(Upscaler):
         model = self.load_model(selected_model)
         if model is None:
             return img
-        model.to(devices.device_esrgan)
+        model.to(devices.device)
         img = esrgan_upscale(model, img)
         if opts.upscaler_unload and selected_model in self.models:
             del self.models[selected_model]
@@ -143,7 +143,7 @@ class UpscalerESRGAN(Upscaler):
         if self.models.get(info.local_data_path, None) is not None:
             log.debug(f"Upscaler cached: type={self.name} model={info.local_data_path}")
             return self.models[info.local_data_path]
-        state_dict = torch.load(info.local_data_path, map_location='cpu' if devices.device_esrgan.type == 'mps' else None)
+        state_dict = torch.load(info.local_data_path, map_location='cpu' if devices.device.type == 'mps' else None)
         log.info(f"Upscaler loaded: type={self.name} model={info.local_data_path}")
 
         if "params_ema" in state_dict:
@@ -179,7 +179,7 @@ def upscale_without_tiling(model, img):
     img = img[:, :, ::-1]
     img = np.ascontiguousarray(np.transpose(img, (2, 0, 1))) / 255
     img = torch.from_numpy(img).float()
-    img = img.unsqueeze(0).to(devices.device_esrgan)
+    img = img.unsqueeze(0).to(devices.device)
     with devices.inference_context():
         output = model(img)
     output = output.squeeze().float().cpu().clamp_(0, 1).detach().numpy()

modules/postprocess/gfpgan_model.py

@@ -16,7 +16,7 @@ def gfpgann():
     import gfpgan # pylint: disable=unused-import
     global loaded_gfpgan_model # pylint: disable=global-statement
     if loaded_gfpgan_model is not None:
-        loaded_gfpgan_model.gfpgan.to(devices.device_gfpgan)
+        loaded_gfpgan_model.gfpgan.to(devices.device)
         return loaded_gfpgan_model
     if gfpgan_constructor is None:
         return None
@@ -30,8 +30,8 @@ def gfpgann():
         shared.log.error(f"Model failed loading: type=GFPGAN model={model_file}")
         return None
     if hasattr(facexlib.detection.retinaface, 'device'):
-        facexlib.detection.retinaface.device = devices.device_gfpgan
-    model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=devices.device_gfpgan)
+        facexlib.detection.retinaface.device = devices.device
+    model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=devices.device)
     loaded_gfpgan_model = model
     shared.log.info(f"Model loaded: type=GFPGAN model={model_file}")
     return model
@@ -48,7 +48,7 @@ def gfpgan_fix_faces(np_image):
     if model is None:
         return np_image
 
-    send_model_to(model, devices.device_gfpgan)
+    send_model_to(model, devices.device)
 
     np_image_bgr = np_image[:, :, ::-1]
     _cropped_faces, _restored_faces, gfpgan_output_bgr = model.enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True)

modules/postprocess/realesrgan_model_arch.py

@@ -55,7 +55,7 @@ class RealESRGANer():
             self.device = torch.device(
                 f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device
         else:
-            self.device = devices.device_esrgan if device is None else device
+            self.device = devices.device if device is None else device
 
         if isinstance(model_path, list):
             # dni