Up to date dev version

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,24 +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(
+    result_slider, result_gallery, restore_information, event_id = restore(
         model,
-        LQ,
-        captions,
     edm_steps,
     s_stage1,
     s_churn,
@@ -216,6 +199,48 @@ def stage2_process(
     spt_linear_s_stage2
     )
 
+    return result_slider, result_gallery, restore_information, event_id
+
+@spaces.GPU(duration=600)
+def restore(
+    model,
+    edm_steps,
+    s_stage1,
+    s_churn,
+    s_noise,
+    s_cfg,
+    s_stage2,
+    seed,
+    num_samples,
+    a_prompt,
+    n_prompt,
+    color_fix_type,
+    linear_CFG,
+    linear_s_stage2,
+    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)
+
+    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,
+                                    use_linear_CFG=linear_CFG, use_linear_control_scale=linear_s_stage2,
+                                    cfg_scale_start=spt_linear_CFG, control_scale_start=spt_linear_s_stage2)
+
     x_samples = (einops.rearrange(samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().round().clip(
         0, 255).astype(np.uint8)
     results = [x_samples[i] for i in range(num_samples)]
@@ -252,33 +277,6 @@ def stage2_process(
     # Only one image can be shown in the slider
     return [noisy_image] + [results[0]], gr.update(format = output_format, value = [noisy_image] + results), gr.update(value = information, visible = True), event_id
 
-@spaces.GPU(duration=600)
-def restore(
-    model,
-    LQ,
-    captions,
-    edm_steps,
-    s_stage1,
-    s_churn,
-    s_noise,
-    s_cfg,
-    s_stage2,
-    seed,
-    num_samples,
-    a_prompt,
-    n_prompt,
-    color_fix_type,
-    linear_CFG,
-    linear_s_stage2,
-    spt_linear_CFG,
-    spt_linear_s_stage2
-):
-    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,
-                                    use_linear_CFG=linear_CFG, use_linear_control_scale=linear_s_stage2,
-                                    cfg_scale_start=spt_linear_CFG, control_scale_start=spt_linear_s_stage2)
-
 def load_and_reset(param_setting):
     print('load_and_reset ==>>')
     if torch.cuda.device_count() == 0:
@@ -346,15 +344,10 @@ title_html = """
 
 claim_md = """
 ## **Piracy**
-
 The images are not stored but the logs are saved during a month.
-
 ## **Terms of use**
-
 By using this service, users are required to agree to the following terms: The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research. Please submit a feedback to us if you get any inappropriate answer! We will collect those to keep improving our models. For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
-
 ## **License**
-
 The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/Fanghua-Yu/SUPIR) of SUPIR.
 """
 
@@ -552,4 +545,4 @@ with gr.Blocks(title="SUPIR") as interface:
         fb_text
     ])
         
-    interface.queue(10).launch()
+    interface.queue(10).launch()