Select model without GPU

app.py

@@ -210,6 +210,19 @@ def restore_in_Xmin(
 
     denoise_image = HWC3(np.array(denoise_image))
 
+    torch.cuda.set_device(SUPIR_device)
+
+    if model_select != model.current_model:
+        print('load ' + model_select)
+        if model_select == 'v0-Q':
+            model.load_state_dict(ckpt_Q, strict=False)
+        elif model_select == 'v0-F':
+            model.load_state_dict(ckpt_F, strict=False)
+        model.current_model = model_select
+
+    model.ae_dtype = convert_dtype(ae_dtype)
+    model.model.dtype = convert_dtype(diff_dtype)
+
     # Allocation
     if allocation == 1:
         return restore_in_1min(
@@ -325,16 +338,6 @@ def restore(
     start = time.time()
     print('restore ==>>')
 
-    torch.cuda.set_device(SUPIR_device)
-
-    if model_select != model.current_model:
-        print('load ' + model_select)
-        if model_select == 'v0-Q':
-            model.load_state_dict(ckpt_Q, strict=False)
-        elif model_select == 'v0-F':
-            model.load_state_dict(ckpt_F, strict=False)
-        model.current_model = model_select
-
     input_image = upscale_image(input_image, upscale, unit_resolution=32, min_size=min_size)
     LQ = np.array(input_image) / 255.0
     LQ = np.power(LQ, gamma_correction)
@@ -344,9 +347,6 @@ def restore(
     LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
     captions = ['']
 
-    model.ae_dtype = convert_dtype(ae_dtype)
-    model.model.dtype = convert_dtype(diff_dtype)
-
     samples = model.batchify_sample(LQ, captions, num_steps=edm_steps, restoration_scale=s_stage1, s_churn=s_churn,
                                     s_noise=s_noise, cfg_scale=s_cfg, control_scale=s_stage2, seed=seed,
                                     num_samples=num_samples, p_p=a_prompt, n_p=n_prompt, color_fix_type=color_fix_type,
@@ -379,7 +379,7 @@ def restore(
     " pixels high, so a resolution of " + f'{result_width * result_height:,}' + " pixels."
     print(information)
     unique_name = str(uuid.uuid4()) + ".png"
-    Image.fromarray(results[0]).save(unique_name)
+    results[0].save(unique_name)
     print(unique_name)
 
     # Only one image can be shown in the slider