Fix ESRGAN with CPU

Artheriax · Dec 31, 2024 · 945c57c · 945c57c
1 parent 48ed3d7
commit 945c57c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -55,6 +55,7 @@ NYE refresh release with quite a few optimizatios and bug fixes...
   - hires with refiner prompt and batch processing  
   - processing with nested calls  
   - ui networks initial sort  
+  - esrgan on cpu devices  
 
 ## Update for 2024-12-24
 

diff --git a/modules/postprocess/esrgan_model.py b/modules/postprocess/esrgan_model.py
@@ -142,7 +142,7 @@ def load_model(self, path: str):
         if self.models.get(info.local_data_path, None) is not None:
             shared.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.type == 'mps' else None)
+        state_dict = torch.load(info.local_data_path, map_location='cpu' if devices.device.type in {'mps', 'cpu'} else None)
         shared.log.info(f"Upscaler loaded: type={self.name} model={info.local_data_path}")
 
         if "params_ema" in state_dict: