Rollback

gradio_demo.py

@@ -117,6 +117,7 @@ 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,
@@ -160,6 +161,7 @@ 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,
@@ -180,47 +182,6 @@ 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
@@ -328,7 +289,7 @@ def submit_feedback(event_id, fb_score, fb_text):
 
 title_html = """
     <h1><center>SUPIR</center></h1>
-    <center><big>Upscale your images up to x8 freely, without account, without watermark and download it</big></center>
+    <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>
     
     <p>This is an online demo of SUPIR, a practicing model scaling for photo-realistic image restoration.
@@ -366,6 +327,7 @@ 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 '
@@ -374,7 +336,6 @@ 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):