Update src/models/sd2_sr.py

src/models/sd2_sr.py

@@ -153,16 +153,20 @@ def load_obj(path):
     return get_obj_from_str(objyaml['__class__'])(**objyaml.get("__init__", {}))
     
 
-def load_model(dtype=torch.bfloat16, device='cuda:0'):
+def load_model(dtype=torch.bfloat16, device=None):
+    # Auto-detect device if not specified
+    if device is None:
+        device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+    
     download_file(DOWNLOAD_URL, MODEL_PATH)
 
     state_dict = safetensors.torch.load_file(MODEL_PATH)
 
     config = OmegaConf.load(f'{CONFIG_FOLDER}/ddpm/v2-upsample.yaml')
 
-    unet = load_obj(f'{CONFIG_FOLDER}/unet/upsample/v2.yaml').eval().cuda()
-    vae = load_obj(f'{CONFIG_FOLDER}/vae-upsample.yaml').eval().cuda()
-    encoder = load_obj(f'{CONFIG_FOLDER}/encoders/openclip.yaml').eval().cuda()
+    unet = load_obj(f'{CONFIG_FOLDER}/unet/upsample/v2.yaml').eval().to(device)
+    vae = load_obj(f'{CONFIG_FOLDER}/vae-upsample.yaml').eval().to(device)
+    encoder = load_obj(f'{CONFIG_FOLDER}/encoders/openclip.yaml').eval().to(device)
     ddim = DDIM(config, vae, encoder, unet)
 
     extract = lambda state_dict, model: {x[len(model)+1:]:y for x,y in state_dict.items() if model in x}
@@ -193,7 +197,7 @@ def load_model(dtype=torch.bfloat16, device='cuda:0'):
         'max_noise_level': 350
     }
 
-    low_scale_model = ImageConcatWithNoiseAugmentation(**params).eval().to('cuda')
+    low_scale_model = ImageConcatWithNoiseAugmentation(**params).eval().to(device)
     low_scale_model.train = disabled_train
     for param in low_scale_model.parameters():
         param.requires_grad = False