Fix function

gradio_demo.py

@@ -180,7 +180,7 @@ def stage2_process(
     input_image = upscale_image(input_image, upscale, unit_resolution=32,
                                 min_size=min_size)
 
-    samples = restore(
+    result_slider, result_gallery, restore_information, event_id = restore(
         model,
     edm_steps,
     s_stage1,
@@ -199,43 +199,8 @@ def stage2_process(
     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)]
-
-    if args.log_history:
-        os.makedirs(f'./history/{event_id[:5]}/{event_id[5:]}', exist_ok=True)
-        with open(f'./history/{event_id[:5]}/{event_id[5:]}/logs.txt', 'w') as f:
-            f.write(str(event_dict))
-        f.close()
-        Image.fromarray(input_image).save(f'./history/{event_id[:5]}/{event_id[5:]}/LQ.png')
-        for i, result in enumerate(results):
-            Image.fromarray(result).save(f'./history/{event_id[:5]}/{event_id[5:]}/HQ_{i}.png')
-
-    # All the results have the same size
-    result_height, result_width, result_channel = np.array(results[0]).shape
-
-    print('<<== stage2_process')
-    end = time.time()
-    secondes = int(end - start)
-    minutes = math.floor(secondes / 60)
-    secondes = secondes - (minutes * 60)
-    hours = math.floor(minutes / 60)
-    minutes = minutes - (hours * 60)
-    information = ("Start the process again if you want a different result. " if randomize_seed else "") + \
-    "The new image resolution is " + str(result_width) + \
-    " pixels large and " + str(result_height) + \
-    " pixels high, so a resolution of " + f'{result_width * result_height:,}' + " pixels. " + \
-    "The image(s) has(ve) been generated in " + \
-    ((str(hours) + " h, ") if hours != 0 else "") + \
-    ((str(minutes) + " min, ") if hours != 0 or minutes != 0 else "") + \
-    str(secondes) + " sec."
-    print(information)
+    return result_slider, result_gallery, restore_information, event_id
 
-    # 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,
     edm_steps,
@@ -269,12 +234,48 @@ def restore(
     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,
+    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)]
+
+    if args.log_history:
+        os.makedirs(f'./history/{event_id[:5]}/{event_id[5:]}', exist_ok=True)
+        with open(f'./history/{event_id[:5]}/{event_id[5:]}/logs.txt', 'w') as f:
+            f.write(str(event_dict))
+        f.close()
+        Image.fromarray(input_image).save(f'./history/{event_id[:5]}/{event_id[5:]}/LQ.png')
+        for i, result in enumerate(results):
+            Image.fromarray(result).save(f'./history/{event_id[:5]}/{event_id[5:]}/HQ_{i}.png')
+
+    # All the results have the same size
+    result_height, result_width, result_channel = np.array(results[0]).shape
+
+    print('<<== stage2_process')
+    end = time.time()
+    secondes = int(end - start)
+    minutes = math.floor(secondes / 60)
+    secondes = secondes - (minutes * 60)
+    hours = math.floor(minutes / 60)
+    minutes = minutes - (hours * 60)
+    information = ("Start the process again if you want a different result. " if randomize_seed else "") + \
+    "The new image resolution is " + str(result_width) + \
+    " pixels large and " + str(result_height) + \
+    " pixels high, so a resolution of " + f'{result_width * result_height:,}' + " pixels. " + \
+    "The image(s) has(ve) been generated in " + \
+    ((str(hours) + " h, ") if hours != 0 else "") + \
+    ((str(minutes) + " min, ") if hours != 0 or minutes != 0 else "") + \
+    str(secondes) + " sec."
+    print(information)
+
+    # 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
+
 def load_and_reset(param_setting):
     print('load_and_reset ==>>')
     if torch.cuda.device_count() == 0: