Validate parameters before

gradio_demo.py

@@ -30,7 +30,7 @@ parser.add_argument("--no_llava", action='store_true', default=True)#False
 parser.add_argument("--use_image_slider", action='store_true', default=False)
 parser.add_argument("--log_history", action='store_true', default=False)
 parser.add_argument("--loading_half_params", action='store_true', default=True)#False
-parser.add_argument("--use_tile_vae", action='store_true', default=False)
+parser.add_argument("--use_tile_vae", action='store_true', default=True)#False
 parser.add_argument("--encoder_tile_size", type=int, default=512)
 parser.add_argument("--decoder_tile_size", type=int, default=64)
 parser.add_argument("--load_8bit_llava", action='store_true', default=False)
@@ -67,15 +67,16 @@ if torch.cuda.device_count() > 0:
     else:
         llava_agent = None
 
-@spaces.GPU(duration=120)
+def check(input_image):
+    if input_image is None:
+        raise gr.Error("Please provide an image to restore.")
+
+@spaces.GPU(duration=180)
 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
-    if input_image is None:
-        gr.Warning('Please provide an image to restore.')
-        return None
     torch.cuda.set_device(SUPIR_device)
     LQ = HWC3(input_image)
     LQ = fix_resize(LQ, 512)
@@ -92,15 +93,12 @@ def stage1_process(input_image, gamma_correction):
     print('End stage1_process')
     return LQ
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=180)
 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)
@@ -111,7 +109,7 @@ def llave_process(input_image, temperature, top_p, qs=None):
     print('End llave_process')
     return captions[0]
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=180)
 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):
@@ -119,9 +117,6 @@ def stage2_process(input_image, prompt, a_prompt, n_prompt, num_samples, upscale
     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,
@@ -279,7 +274,7 @@ with gr.Blocks(title='SUPIR') as interface:
                 qs = gr.Textbox(label="Question", info="Describe the image and its style in a very detailed manner", placeholder="The image is a realistic photography, not an art painting.")
 
             with gr.Accordion("Restoring options", open=False):
-                num_samples = gr.Slider(label="Num Samples", info="Number of generated results; I discourage to increase because the process is limited to 2 min", minimum=1, maximum=4 if not args.use_image_slider else 1
+                num_samples = gr.Slider(label="Num Samples", info="Number of generated results; I discourage to increase because the process is limited to 3 min", minimum=1, maximum=4 if not args.use_image_slider else 1
                                         , value=1, step=1)
                 upscale = gr.Slider(label="Upscale", info="The resolution increase factor", minimum=1, maximum=8, value=1, step=1)
                 edm_steps = gr.Slider(label="Steps", info="lower=faster, higher=more details", minimum=1, maximum=200, value=default_setting.edm_steps if torch.cuda.device_count() > 0 else 1, step=1)
@@ -319,10 +314,10 @@ with gr.Blocks(title='SUPIR') as interface:
                         ae_dtype = gr.Radio(['fp32', 'bf16'], label="Auto-Encoder Data Type", value="bf16",
                                             interactive=True)
                     with gr.Column():
-                        color_fix_type = gr.Radio(["None", "AdaIn", "Wavelet"], label="Color-Fix Type", value="Wavelet",
+                        color_fix_type = gr.Radio(["None", "AdaIn", "Wavelet"], label="Color-Fix Type", info="Wavelet=For JPEG artifacts", value="Wavelet",
                                                   interactive=True)
                     with gr.Column():
-                        model_select = gr.Radio(["v0-Q", "v0-F"], label="Model Selection", value="v0-Q",
+                        model_select = gr.Radio(["v0-Q", "v0-F"], label="Model Selection", info="Q=Quality, F=Fidelity", value="v0-Q",
                                                 interactive=True)
 
         with gr.Column():
@@ -352,17 +347,83 @@ with gr.Blocks(title='SUPIR') as interface:
     with gr.Row():
         gr.Markdown(claim_md)
         event_id = gr.Textbox(label="Event ID", value="", visible=False)
-
-    llave_button.click(fn=llave_process, inputs=[denoise_image, temperature, top_p, qs], outputs=[prompt])
-    denoise_button.click(fn=stage1_process, inputs=[input_image, gamma_correction],
-                         outputs=[denoise_image])
-    stage2_ips = [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]
-    diffusion_button.click(fn=stage2_process, inputs=stage2_ips, outputs=[result_gallery, event_id, fb_score, fb_text])
-    restart_button.click(fn=load_and_reset, inputs=[param_setting],
-                         outputs=[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])
-    submit_button.click(fn=submit_feedback, inputs=[event_id, fb_score, fb_text], outputs=[fb_text])
+    
+    denoise_button.click(fn = check, inputs = [
+        input_image
+    ], outputs = [], queue = False, show_progress = False).success(fn = stage1_process, inputs = [
+        input_image,
+        gamma_correction
+    ], outputs=[
+        denoise_image
+    ])
+
+    llave_button.click(fn = check, inputs = [
+        denoise_image
+    ], outputs = [], queue = False, show_progress = False).success(fn = llave_process, inputs = [
+        denoise_image,
+        temperature,
+        top_p,
+        qs
+    ], outputs = [
+        prompt
+    ])
+
+    diffusion_button.click(fn = check, inputs = [
+        input_image
+    ], outputs = [], queue = False, show_progress = False).success(fn=stage2_process, inputs = [
+        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
+    ], outputs = [
+        result_gallery,
+        event_id,
+        fb_score,
+        fb_text
+    ])
+
+    restart_button.click(fn = load_and_reset, inputs = [
+        param_setting
+    ], outputs = [
+        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
+    ])
+
+    submit_button.click(fn = submit_feedback, inputs = [
+        event_id,
+        fb_score,
+        fb_text
+    ], outputs = [
+        fb_text
+    ])
         
     interface.queue(10).launch()