This PR sets allocation per second

Click on _Merge_ to add this feature

app.py

@@ -73,7 +73,7 @@ def reset():
         None,
         None,
         "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.",
-        "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
+        "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, pixel, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
         1,
         1024,
         1,
@@ -96,12 +96,13 @@ def reset():
         0.,
         "v0-Q",
         "input",
-        6
+        179
     ]
 
-def check(input_image):
+def check_and_update(input_image):
     if input_image is None:
         raise gr.Error("Please provide an image to restore.")
+    return gr.update(visible = True)
 
 @spaces.GPU(duration=420)
 def stage1_process(
@@ -272,88 +273,43 @@ def restore_in_Xmin(
     model.ae_dtype = convert_dtype(ae_dtype)
     model.model.dtype = convert_dtype(diff_dtype)
 
-    # Allocation
-    if allocation == 1:
-        return restore_in_1min(
-            noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, 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, output_format, allocation
-        )
-    if allocation == 2:
-        return restore_in_2min(
-            noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, 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, output_format, allocation
-        )
-    if allocation == 3:
-        return restore_in_3min(
-            noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, 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, output_format, allocation
-        )
-    if allocation == 4:
-        return restore_in_4min(
-            noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, 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, output_format, allocation
-        )
-    if allocation == 5:
-        return restore_in_5min(
-            noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, 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, output_format, allocation
-        )
-    if allocation == 7:
-        return restore_in_7min(
-            noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, 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, output_format, allocation
-        )
-    if allocation == 8:
-        return restore_in_8min(
-            noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, 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, output_format, allocation
-        )
-    if allocation == 9:
-        return restore_in_9min(
-            noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, 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, output_format, allocation
-        )
-    if allocation == 10:
-        return restore_in_10min(
-            noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, 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, output_format, allocation
-        )
-    else:
-        return restore_in_6min(
-            noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, 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, output_format, allocation
-        )
-
-@spaces.GPU(duration=59)
-def restore_in_1min(*args, **kwargs):
-    return restore_on_gpu(*args, **kwargs)
-
-@spaces.GPU(duration=119)
-def restore_in_2min(*args, **kwargs):
-    return restore_on_gpu(*args, **kwargs)
-
-@spaces.GPU(duration=179)
-def restore_in_3min(*args, **kwargs):
-    return restore_on_gpu(*args, **kwargs)
-
-@spaces.GPU(duration=239)
-def restore_in_4min(*args, **kwargs):
-    return restore_on_gpu(*args, **kwargs)
-
-@spaces.GPU(duration=299)
-def restore_in_5min(*args, **kwargs):
-    return restore_on_gpu(*args, **kwargs)
-
-@spaces.GPU(duration=359)
-def restore_in_6min(*args, **kwargs):
-    return restore_on_gpu(*args, **kwargs)
-
-@spaces.GPU(duration=419)
-def restore_in_7min(*args, **kwargs):
-    return restore_on_gpu(*args, **kwargs)
-
-@spaces.GPU(duration=479)
-def restore_in_8min(*args, **kwargs):
-    return restore_on_gpu(*args, **kwargs)
-
-@spaces.GPU(duration=539)
-def restore_in_9min(*args, **kwargs):
-    return restore_on_gpu(*args, **kwargs)
+    return restore_on_gpu(
+        noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, 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, output_format, allocation
+    )
 
-@spaces.GPU(duration=599)
-def restore_in_10min(*args, **kwargs):
-    return restore_on_gpu(*args, **kwargs)
+def get_duration(
+    noisy_image,
+    input_image,
+    prompt,
+    a_prompt,
+    n_prompt,
+    num_samples,
+    min_size,
+    downscale,
+    upscale,
+    edm_steps,
+    s_stage1,
+    s_stage2,
+    s_cfg,
+    randomize_seed,
+    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,
+    output_format,
+    allocation
+):
+    return allocation
 
+@spaces.GPU(duration=get_duration)
 def restore_on_gpu(
     noisy_image,
     input_image,
@@ -541,10 +497,10 @@ with gr.Blocks() as interface:
     rotation = gr.Radio([["No rotation", 0], ["⤵ Rotate +90°", 90], ["↩ Return 180°", 180], ["⤴ Rotate -90°", -90]], label="Orientation correction", info="Will apply the following rotation before restoring the image; the AI needs a good orientation to understand the content", value=0, interactive=True, visible=False)
     with gr.Group():
         prompt = gr.Textbox(label="Image description", info="Help the AI understand what the image represents; describe as much as possible, especially the details we can't see on the original image; you can write in any language", value="", placeholder="A 33 years old man, walking, in the street, Santiago, morning, Summer, photorealistic", lines=3)
-        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.")
+        prompt_hint = gr.HTML("You can use a <a href='"'https://huggingface.co/spaces/badayvedat/LLaVA'"'>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], ["x9", 9], ["x10", 10]], label="Upscale factor", info="Resolution x1 to x10", value=2, interactive=True)
         output_format = gr.Radio([["As input", "input"], ["*.png", "png"], ["*.webp", "webp"], ["*.jpeg", "jpeg"], ["*.gif", "gif"], ["*.bmp", "bmp"]], label="Image format for result", info="File extention", value="input", interactive=True)
-        allocation = gr.Radio([["1 min", 1], ["2 min", 2], ["3 min", 3], ["4 min", 4], ["5 min", 5]], label="GPU allocation time", info="lower=May abort run, higher=Quota penalty for next runs", value=3, interactive=True)
+        allocation = gr.Slider(label="GPU allocation time (in seconds)", info="lower=May abort run, higher=Quota penalty for next runs", value=179, minimum=59, maximum=320, step=1)
 
     with gr.Accordion("Pre-denoising (optional)", open=False):
         gamma_correction = gr.Slider(label="Gamma Correction", info = "lower=lighter, higher=darker", minimum=0.1, maximum=2.0, value=1.0, step=0.1)
@@ -557,17 +513,17 @@ with gr.Blocks() as interface:
                               info="Completes the main image description",
                               value='Cinematic, High Contrast, highly detailed, taken using a Canon EOS R '
                                     'camera, hyper detailed photo - realistic maximum detail, 32k, Color '
-                                    'Grading, ultra HD, extreme meticulous detailing, skin pore detailing, '
+                                    'Grading, ultra HD, extreme meticulous detailing, skin pore detailing, clothing fabric detailing, '
                                     'hyper sharpness, perfect without deformations.',
                               lines=3)
         n_prompt = gr.Textbox(label="Negative image description",
                               info="Disambiguate by listing what the image does NOT represent",
                               value='painting, oil painting, illustration, drawing, art, sketch, anime, '
-                                    'cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, unsharp, weird textures, ugly, dirty, messy, '
+                                    'cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, pixel, unsharp, weird textures, ugly, dirty, messy, '
                                     'worst quality, low quality, frames, watermark, signature, jpeg artifacts, '
                                     'deformed, lowres, over-smooth',
                               lines=3)
-        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)
+        edm_steps = gr.Slider(label="Steps", info="lower=faster, higher=more details; too many steps create a checker effect", minimum=1, maximum=200, value=default_setting.edm_steps if torch.cuda.device_count() > 0 else 1, step=1)
         num_samples = gr.Slider(label="Num Samples", info="Number of generated results", minimum=1, maximum=4 if not args.use_image_slider else 1
                                 , value=1, step=1)
         min_size = gr.Slider(label="Minimum size", info="Minimum height, minimum width of the result", minimum=32, maximum=4096, value=1024, step=32)
@@ -610,6 +566,7 @@ with gr.Blocks() as interface:
         diffusion_button = gr.Button(value="🚀 Upscale/Restore", variant = "primary", elem_id = "process_button")
         reset_btn = gr.Button(value="🧹 Reinit page", variant="stop", elem_id="reset_button", visible = False)
 
+        warning = gr.HTML(value = "<center><big>Your computer must <u>not</u> enter into standby mode.</big><br/>On Chrome, you can force to keep a tab alive in <code>chrome://discards/</code></center>", visible = False)
         restore_information = gr.HTML(value = "Restart the process to get another result.", visible = False)
         result_slider = ImageSlider(label = 'Comparator', show_label = False, interactive = False, elem_id = "slider1", show_download_button = False)
         result_gallery = gr.Gallery(label = 'Downloadable results', show_label = True, interactive = False, elem_id = "gallery1")
@@ -622,12 +579,12 @@ with gr.Blocks() as interface:
                     None,
                     "Group of people, walking, happy, in the street, photorealistic, 8k, extremely detailled",
                     "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.",
-                    "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
+                    "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, pixel, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
                     2,
                     1024,
                     1,
                     8,
-                    200,
+                    100,
                     -1,
                     1,
                     7.5,
@@ -645,7 +602,7 @@ with gr.Blocks() as interface:
                     0.,
                     "v0-Q",
                     "input",
-                    5
+                    179
                 ],
                 [
                     "./Examples/Example2.jpeg",
@@ -653,7 +610,7 @@ with gr.Blocks() as interface:
                     None,
                     "La cabeza de un gato atigrado, en una casa, fotorrealista, 8k, extremadamente detallada",
                     "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.",
-                    "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
+                    "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, pixel, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
                     1,
                     1024,
                     1,
@@ -676,7 +633,7 @@ with gr.Blocks() as interface:
                     0.,
                     "v0-Q",
                     "input",
-                    4
+                    179
                 ],
                 [
                     "./Examples/Example3.webp",
@@ -684,7 +641,7 @@ with gr.Blocks() as interface:
                     None,
                     "A red apple",
                     "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.",
-                    "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
+                    "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, pixel, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
                     1,
                     1024,
                     1,
@@ -707,7 +664,7 @@ with gr.Blocks() as interface:
                     0.,
                     "v0-Q",
                     "input",
-                    4
+                    179
                 ],
                 [
                     "./Examples/Example3.webp",
@@ -715,7 +672,7 @@ with gr.Blocks() as interface:
                     None,
                     "A red marble",
                     "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.",
-                    "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
+                    "painting, oil painting, illustration, drawing, art, sketch, anime, cartoon, CG Style, 3D render, unreal engine, blurring, aliasing, pixel, unsharp, weird textures, ugly, dirty, messy, worst quality, low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth",
                     1,
                     1024,
                     1,
@@ -738,7 +695,7 @@ with gr.Blocks() as interface:
                     0.,
                     "v0-Q",
                     "input",
-                    4
+                    179
                 ],
             ],
         run_on_click = True,
@@ -792,9 +749,9 @@ with gr.Blocks() as interface:
         rotation
     ], queue = False, show_progress = False)
 
-    denoise_button.click(fn = check, inputs = [
+    denoise_button.click(fn = check_and_update, inputs = [
         input_image
-    ], outputs = [], queue = False, show_progress = False).success(fn = stage1_process, inputs = [
+    ], outputs = [warning], queue = False, show_progress = False).success(fn = stage1_process, inputs = [
         input_image,
         gamma_correction,
         diff_dtype,
@@ -809,9 +766,9 @@ with gr.Blocks() as interface:
         seed
     ], outputs = [
         seed
-    ], queue = False, show_progress = False).then(fn = check, inputs = [
+    ], queue = False, show_progress = False).then(fn = check_and_update, inputs = [
         input_image
-    ], outputs = [], queue = False, show_progress = False).success(fn=stage2_process, inputs = [
+    ], outputs = [warning], queue = False, show_progress = False).success(fn=stage2_process, inputs = [
         input_image,
         rotation,
         denoise_image,