Code out of GPU compute

app.py

@@ -100,7 +100,7 @@ def stage2_process(*args, **kwargs):
         return restore_in_Xmin(*args, **kwargs)
     except Exception as e:
         print('Exception of type ' + str(type(e)))
-        if type(e).__name__ == 'gradio.exceptions.Error':
+        if type(e).__name__ == "<class 'gradio.exceptions.Error'>":
             print('Exception of name ' + type(e).__name__)
         raise e
 
@@ -201,6 +201,12 @@ def restore_in_Xmin(
         input_height, input_width, input_channel = denoise_image.shape
         denoise_image = denoise_image.resize((input_width // downscale, input_height // downscale), Image.LANCZOS)
 
+    if torch.cuda.device_count() == 0:
+        gr.Warning('Set this space to GPU config to make it work.')
+        return [noisy_image, denoise_image], [denoise_image], None, None
+
+    denoise_image = HWC3(np.array(denoise_image))
+
     # Allocation
     if allocation == 1:
         return restore_in_1min(
@@ -243,39 +249,39 @@ def restore_in_Xmin(
             noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, output_format, allocation
         )
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=59)
 def restore_in_1min(*args, **kwargs):
     return restore(*args, **kwargs)
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=119)
 def restore_in_2min(*args, **kwargs):
     return restore(*args, **kwargs)
 
-@spaces.GPU(duration=180)
+@spaces.GPU(duration=179)
 def restore_in_3min(*args, **kwargs):
     return restore(*args, **kwargs)
 
-@spaces.GPU(duration=240)
+@spaces.GPU(duration=239)
 def restore_in_4min(*args, **kwargs):
     return restore(*args, **kwargs)
 
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=299)
 def restore_in_5min(*args, **kwargs):
     return restore(*args, **kwargs)
 
-@spaces.GPU(duration=360)
+@spaces.GPU(duration=359)
 def restore_in_6min(*args, **kwargs):
     return restore(*args, **kwargs)
 
-@spaces.GPU(duration=420)
+@spaces.GPU(duration=419)
 def restore_in_7min(*args, **kwargs):
     return restore(*args, **kwargs)
 
-@spaces.GPU(duration=480)
+@spaces.GPU(duration=479)
 def restore_in_8min(*args, **kwargs):
     return restore(*args, **kwargs)
 
-@spaces.GPU(duration=540)
+@spaces.GPU(duration=539)
 def restore_in_9min(*args, **kwargs):
     return restore(*args, **kwargs)
 
@@ -316,9 +322,6 @@ def restore(
     start = time.time()
     print('restore ==>>')
 
-    if torch.cuda.device_count() == 0:
-        gr.Warning('Set this space to GPU config to make it work.')
-        return [input_image] * 2, [input_image], None, None
     torch.cuda.set_device(SUPIR_device)
 
     if model_select != model.current_model:
@@ -328,7 +331,6 @@ def restore(
         elif model_select == 'v0-F':
             model.load_state_dict(ckpt_F, strict=False)
         model.current_model = model_select
-    input_image = HWC3(np.array(input_image))
     input_image = upscale_image(input_image, upscale, unit_resolution=32,
                                 min_size=min_size)