ZERO space only for AI

gradio_demo.py

@@ -79,7 +79,7 @@ def check(input_image):
 def reset_feedback():
     return 3, ''
 
-@spaces.GPU(duration=540)
+@spaces.GPU(duration=600)
 def stage1_process(input_image, gamma_correction):
     print('stage1_process ==>>')
     if torch.cuda.device_count() == 0:
@@ -101,7 +101,7 @@ def stage1_process(input_image, gamma_correction):
     print('<<== stage1_process')
     return LQ, gr.update(visible = True)
 
-@spaces.GPU(duration=540)
+@spaces.GPU(duration=600)
 def llave_process(input_image, temperature, top_p, qs=None):
     print('llave_process ==>>')
     if torch.cuda.device_count() == 0:
@@ -117,7 +117,6 @@ def llave_process(input_image, temperature, top_p, qs=None):
     print('<<== llave_process')
     return captions[0]
 
-@spaces.GPU(duration=540)
 def stage2_process(
     noisy_image,
     denoise_image,
@@ -196,11 +195,26 @@ def stage2_process(
     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)
+    samples = 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
+    )
 
     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)
@@ -238,6 +252,33 @@ 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:
@@ -289,15 +330,15 @@ def submit_feedback(event_id, fb_score, fb_text):
 
 title_html = """
     <h1><center>SUPIR</center></h1>
-    <big><center>Upscale your images up to x8 freely, without account, without watermark and download it</center></big>
-    <center>🤸<big><big><big><big><big><big><big><big>🤸</big></big></big></big></big></big></big></big></center>
+    <center><big>Upscale your images up to x8 freely, without account, without watermark and download it</big></center>
+    <center><big><big>🤸<big><big><big><big><big><big>🤸</big></big></big></big></big></big></big></big></center>
     
     <p>This is an online demo of SUPIR, a practicing model scaling for photo-realistic image restoration.
     It is still a research project under tested and is not yet a stable commercial product.
     LlaVa is not integrated in this demo. The content added by SUPIR is imagination, not real-world information.
     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 9 min.
+    This demo can handle huge images but the process will be aborted if it lasts more than 10 min.
     
     <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://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></center></p>
     """