Add logs

gradio_demo.py

@@ -69,40 +69,38 @@ if torch.cuda.device_count() > 0:
 
 @spaces.GPU(duration=120)
 def stage1_process(input_image, gamma_correction):
+    print('Start stage1_process')
     if torch.cuda.device_count() == 0:
         gr.Warning('Set this space to GPU config to make it work.')
         return None
-    print ("stage1_process 1")
+    if input_image is None:
+        gr.Warning('Please provide an image to restore.')
+        return None
     torch.cuda.set_device(SUPIR_device)
-    print ("stage1_process 2")
     LQ = HWC3(input_image)
-    print ("stage1_process 3")
     LQ = fix_resize(LQ, 512)
-    print ("stage1_process 4")
     # stage1
     LQ = np.array(LQ) / 255 * 2 - 1
-    print ("stage1_process 5")
     LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
-    print ("stage1_process 6")
     LQ = model.batchify_denoise(LQ, is_stage1=True)
-    print ("stage1_process 7")
     LQ = (LQ[0].permute(1, 2, 0) * 127.5 + 127.5).cpu().numpy().round().clip(0, 255).astype(np.uint8)
     # gamma correction
     LQ = LQ / 255.0
-    print ("stage1_process 8")
     LQ = np.power(LQ, gamma_correction)
-    print ("stage1_process 9")
     LQ *= 255.0
-    print ("stage1_process 10")
     LQ = LQ.round().clip(0, 255).astype(np.uint8)
-    print ("stage1_process 11")
+    print('End stage1_process')
     return LQ
 
 @spaces.GPU(duration=120)
 def llave_process(input_image, temperature, top_p, qs=None):
+    print('Start llave_process')
     if torch.cuda.device_count() == 0:
         gr.Warning('Set this space to GPU config to make it work.')
         return 'Set this space to GPU config to make it work.'
+    if input_image is None:
+        gr.Warning('Please provide an image to restore.')
+        return 'Please provide an image to restore.'
     torch.cuda.set_device(LLaVA_device)
     if use_llava:
         LQ = HWC3(input_image)
@@ -110,15 +108,20 @@ def llave_process(input_image, temperature, top_p, qs=None):
         captions = llava_agent.gen_image_caption([LQ], temperature=temperature, top_p=top_p, qs=qs)
     else:
         captions = ['LLaVA is not available. Please add text manually.']
+    print('End llave_process')
     return captions[0]
 
 @spaces.GPU(duration=120)
 def stage2_process(input_image, prompt, a_prompt, n_prompt, num_samples, upscale, edm_steps, s_stage1, s_stage2,
                    s_cfg, 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):
+    print('Start stage2_process')
     if torch.cuda.device_count() == 0:
         gr.Warning('Set this space to GPU config to make it work.')
         return None, None, None, None
+    if input_image is None:
+        gr.Warning('Please provide an image to restore.')
+        return None, None, 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,
@@ -130,14 +133,12 @@ def stage2_process(input_image, prompt, a_prompt, n_prompt, num_samples, upscale
                   'model_select': model_select}
 
     if model_select != model.current_model:
+        print('load ' + model_select)
         if model_select == 'v0-Q':
-            print('load v0-Q')
             model.load_state_dict(ckpt_Q, strict=False)
-            model.current_model = 'v0-Q'
         elif model_select == 'v0-F':
-            print('load v0-F')
             model.load_state_dict(ckpt_F, strict=False)
-            model.current_model = 'v0-F'
+        model.current_model = model_select
     input_image = HWC3(input_image)
     input_image = upscale_image(input_image, upscale, unit_resolution=32,
                                 min_size=1024)
@@ -174,10 +175,11 @@ def stage2_process(input_image, prompt, a_prompt, n_prompt, num_samples, upscale
         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')
+    print('End stage2_process')
     return [input_image] + results, event_id, 3, ''
 
-@spaces.GPU(duration=120)
 def load_and_reset(param_setting):
+    print('Start load_and_reset')
     if torch.cuda.device_count() == 0:
         gr.Warning('Set this space to GPU config to make it work.')
         return None, None, None, None, None, None, None, None, None, None, None, None, None
@@ -204,6 +206,7 @@ def load_and_reset(param_setting):
         spt_linear_CFG = default_setting.spt_linear_CFG_Fidelity
     else:
         raise NotImplementedError
+    print('End load_and_reset')
     return edm_steps, s_cfg, s_stage2, s_stage1, s_churn, s_noise, a_prompt, n_prompt, color_fix_type, linear_CFG, \
         linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2
 
@@ -237,7 +240,7 @@ else:
     ⚠️SUPIR is still a research project under tested and is not yet a stable commercial product.
     
     <a href="https://arxiv.org/abs/2401.13627">Paper</a> &emsp; <a href="http://supir.xpixel.group/">Project Page</a> &emsp; <a href="https://github.com/Fanghua-Yu/SUPIR/blob/master/assets/DemoGuide.png">How to play</a> &emsp; <a href="https://huggingface.co/blog/MonsterMMORPG/supir-sota-image-upscale-better-than-magnific-ai">Local Install Guide</a>
-    <p style="background-color: blue;">For now, only the stage 2 is working (the most important one). The stage 1 and LLaVa are failing. LLaVa is disabled.</p>
+    <p style="background-color: blue;">For now, only the restoring is working (the most important one). The pre-denoising and LLaVa description are failing. LLaVa is disabled.</p>
     """