Reduce ZERO space usage

gradio_demo.py

@@ -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,
@@ -161,7 +160,6 @@ def stage2_process(
     if 1 < downscale:
         input_height, input_width, input_channel = np.array(input_image).shape
         input_image = input_image.resize((input_width // downscale, input_height // downscale), Image.LANCZOS)
-    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,
@@ -182,6 +180,47 @@ def stage2_process(
     input_image = upscale_image(input_image, upscale, unit_resolution=32,
                                 min_size=min_size)
 
+    result_slider, result_gallery, restore_information, event_id = restore(
+        model,
+    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 result_slider, result_gallery, restore_information, event_id
+
+@spaces.GPU(duration=540)
+def restore(
+    model,
+    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
+):
+    torch.cuda.set_device(SUPIR_device)
     LQ = np.array(input_image) / 255.0
     LQ = np.power(LQ, gamma_correction)
     LQ *= 255.0
@@ -289,7 +328,7 @@ 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>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.
@@ -327,7 +366,6 @@ with gr.Blocks(title="SUPIR") as interface:
     input_image = gr.Image(label="Input", show_label=True, type="numpy", height=600, elem_id="image-input")
     with gr.Group():
         prompt = gr.Textbox(label="Image description for LlaVa", value="", placeholder="A person, walking, in a town, Summer, photorealistic", lines=3, visible=False)
-        upscale = gr.Radio([["x1", 1], ["x2", 2], ["x3", 3], ["x4", 4], ["x5", 5], ["x6", 6], ["x7", 7], ["x8", 8]], label="Upscale factor", info="Resolution x1 to x8", value=2, interactive=True)
         a_prompt = gr.Textbox(label="Image description",
                               info="Help the AI understand what the image represents; describe as much as possible",
                               value='Cinematic, High Contrast, highly detailed, taken using a Canon EOS R '
@@ -336,6 +374,7 @@ with gr.Blocks(title="SUPIR") as interface:
                                     'hyper sharpness, perfect without deformations.',
                               lines=3)
         a_prompt_hint = gr.HTML("You can use a <a href='"'https://huggingface.co/spaces/MaziyarPanahi/llava-llama-3-8b'"'>LlaVa space</a> to auto-generate the description of your image.")
+        upscale = gr.Radio([["x1", 1], ["x2", 2], ["x3", 3], ["x4", 4], ["x5", 5], ["x6", 6], ["x7", 7], ["x8", 8]], label="Upscale factor", info="Resolution x1 to x8", value=2, interactive=True)
         output_format = gr.Radio([["*.png", "png"], ["*.webp", "webp"], ["*.jpeg", "jpeg"], ["*.gif", "gif"], ["*.bmp", "bmp"]], label="Image format for result", info="File extention", value="png", interactive=True)
 
     with gr.Accordion("Pre-denoising (optional)", open=False):