Fix function

gradio_demo.py

@@ -160,7 +160,6 @@ def stage2_process(
     if 1 < downscale:
         input_height, input_width, input_channel = np.array(input_image).shape
         input_image = input_image.resize((input_width // downscale, input_height // downscale), Image.LANCZOS)
-    torch.cuda.set_device(SUPIR_device)
     event_id = str(time.time_ns())
     event_dict = {'event_id': event_id, 'localtime': time.ctime(), 'prompt': prompt, 'a_prompt': a_prompt,
                   'n_prompt': n_prompt, 'num_samples': num_samples, 'upscale': upscale, 'edm_steps': edm_steps,
@@ -181,23 +180,8 @@ def stage2_process(
     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)
-    LQ *= 255.0
-    LQ = LQ.round().clip(0, 255).astype(np.uint8)
-    LQ = LQ / 255 * 2 - 1
-    LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
-    if use_llava:
-        captions = [prompt]
-    else:
-        captions = ['']
-
-    model.ae_dtype = convert_dtype(ae_dtype)
-    model.model.dtype = convert_dtype(diff_dtype)
-
     samples = restore(
         model,
-        LQ,
         captions,
     edm_steps,
     s_stage1,
@@ -255,7 +239,6 @@ def stage2_process(
 @spaces.GPU(duration=600)
 def restore(
     model,
-    LQ,
     captions,
     edm_steps,
     s_stage1,
@@ -273,6 +256,21 @@ def restore(
     spt_linear_CFG,
     spt_linear_s_stage2
 ):
+    torch.cuda.set_device(SUPIR_device)
+    LQ = np.array(input_image) / 255.0
+    LQ = np.power(LQ, gamma_correction)
+    LQ *= 255.0
+    LQ = LQ.round().clip(0, 255).astype(np.uint8)
+    LQ = LQ / 255 * 2 - 1
+    LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
+    if use_llava:
+        captions = [prompt]
+    else:
+        captions = ['']
+
+    model.ae_dtype = convert_dtype(ae_dtype)
+    model.model.dtype = convert_dtype(diff_dtype)
+
     return 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,