model.ae_dtype = convert_dtype(ae_dtype)

app.py

@@ -82,7 +82,12 @@ def check(input_image):
         raise gr.Error("Please provide an image to restore.")
 
 @spaces.GPU(duration=420)
-def stage1_process(input_image, gamma_correction):
+def stage1_process(
+    input_image,
+    gamma_correction,
+    diff_dtype,
+    ae_dtype
+):
     print('stage1_process ==>>')
     if torch.cuda.device_count() == 0:
         gr.Warning('Set this space to GPU config to make it work.')
@@ -93,6 +98,10 @@ def stage1_process(input_image, gamma_correction):
     # stage1
     LQ = np.array(LQ) / 255 * 2 - 1
     LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
+
+    model.ae_dtype = convert_dtype(ae_dtype)
+    model.model.dtype = convert_dtype(diff_dtype)
+
     LQ = model.batchify_denoise(LQ, is_stage1=True)
     LQ = (LQ[0].permute(1, 2, 0) * 127.5 + 127.5).cpu().numpy().round().clip(0, 255).astype(np.uint8)
     # gamma correction
@@ -318,14 +327,6 @@ def restore(
         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)
-    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,
-                  's_stage1': s_stage1, 's_stage2': s_stage2, 's_cfg': s_cfg, 'seed': seed, 's_churn': s_churn,
-                  's_noise': s_noise, 'color_fix_type': color_fix_type, 'diff_dtype': diff_dtype, 'ae_dtype': ae_dtype,
-                  'gamma_correction': gamma_correction, 'linear_CFG': linear_CFG, 'linear_s_stage2': linear_s_stage2,
-                  'spt_linear_CFG': spt_linear_CFG, 'spt_linear_s_stage2': spt_linear_s_stage2,
-                  'model_select': model_select}
 
     if model_select != model.current_model:
         print('load ' + model_select)
@@ -362,15 +363,6 @@ def restore(
         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
 
@@ -452,7 +444,7 @@ title_html = """
     The aim of SUPIR is the beauty and the illustration.
     Most of the processes only last few minutes.
     This demo can handle huge images but the process will be aborted if it lasts more than 8 min.
-    Please <a href="https://huggingface.co/spaces/Fabrice-TIERCELIN/SUPIR/discussions/new">leave a message in discussion</a> if you encounter issues.
+    Please leave a <a href="https://huggingface.co/spaces/Fabrice-TIERCELIN/SUPIR/discussions/new">message in discussion</a> if you encounter issues.
     
     <p><center><a href="https://arxiv.org/abs/2401.13627">Paper</a> &emsp; <a href="http://supir.xpixel.group/">Project Page</a> &emsp; <a href="https://huggingface.co/blog/MonsterMMORPG/supir-sota-image-upscale-better-than-magnific-ai">Local Install Guide</a></center></p>
     """
@@ -477,7 +469,7 @@ with gr.Blocks() as interface:
             gr.HTML("""
     <p style="background-color: red;"><big><big><big><b>⚠️To use SUPIR, <a href="https://huggingface.co/spaces/Fabrice-TIERCELIN/SUPIR?duplicate=true">duplicate this space</a> and set a GPU with 30 GB VRAM.</b>
     
-    You can't use SUPIR directly here because this space runs on a CPU, which is not enough for SUPIR. This is a template space. Please provide feedback if you have issues.
+    You can't use SUPIR directly here because this space runs on a CPU, which is not enough for SUPIR. Please provide <a href="https://huggingface.co/spaces/Fabrice-TIERCELIN/SUPIR/discussions/new">feedback</a> if you have issues.
     </big></big></big></p>
     """)
     gr.HTML(title_html)
@@ -677,7 +669,9 @@ with gr.Blocks() as interface:
         input_image
     ], outputs = [], queue = False, show_progress = False).success(fn = stage1_process, inputs = [
         input_image,
-        gamma_correction
+        gamma_correction,
+        diff_dtype,
+        ae_dtype
     ], outputs=[
         denoise_image,
         denoise_information