Update app.py

fixing runtime error

app.py

@@ -83,7 +83,6 @@ def update_seed(is_randomize_seed, seed):
 
 def reset():
     return [
-        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",
@@ -102,28 +101,81 @@ def reset():
         "Wavelet",
         "fp32",
         "fp32",
-        1.0,
         True,
         False,
         default_setting.spt_linear_CFG_Quality if torch.cuda.device_count() > 0 else 1.0,
         0.,
         "v0-Q",
-        6
+        4
     ]
 
 def check(input_image):
     if input_image is None:
         raise gr.Error("Please provide an image to restore.")
 
-def stage2_process(*args, **kwargs):
+def stage2_process(
+    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,
+    allocation
+):
     try:
-        return restore_in_Xmin(*args, **kwargs)
+        return restore_in_Xmin(
+            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,
+            allocation
+        )
     except Exception as e:
         print(f"Exception occurred: {str(e)}")
         raise e
 
 def restore_in_Xmin(
-    noisy_image,
+    input_image,
     prompt,
     a_prompt,
     n_prompt,
@@ -151,7 +203,7 @@ def restore_in_Xmin(
     allocation
 ):
     print("Starting image restoration process...")
-    input_format = re.sub(r"^.*\.([^\.]+)$", r"\1", noisy_image)
+    input_format = re.sub(r"^.*\.([^\.]+)$", r"\1", input_image)
 
     if input_format.lower() not in ['png', 'webp', 'jpg', 'jpeg', 'gif', 'bmp', 'heic']:
         gr.Warning('Invalid image format. Please use a supported image format.')
@@ -174,7 +226,7 @@ def restore_in_Xmin(
         a_prompt = prompt + a_prompt
     print("Final prompt: " + str(a_prompt))
 
-    denoise_image = np.array(Image.open(noisy_image))
+    denoise_image = np.array(Image.open(input_image))
 
     if 1 < downscale:
         input_height, input_width, input_channel = denoise_image.shape
@@ -184,7 +236,7 @@ def restore_in_Xmin(
 
     if torch.cuda.device_count() == 0:
         gr.Warning('Set this space to GPU config to make it work.')
-        return [noisy_image, denoise_image], gr.update(label="Downloadable results", format=output_format, value=[denoise_image]), None, gr.update(visible=True)
+        return [input_image, denoise_image], gr.update(label="Downloadable results", format=output_format, value=[denoise_image]), None, gr.update(visible=True)
 
     if model_select != model.current_model:
         print('Loading model: ' + model_select)
@@ -211,9 +263,9 @@ def restore_in_Xmin(
         10: restore_in_10min,
     }
 
-    restore_function = allocation_functions.get(allocation, restore_in_6min)
+    restore_function = allocation_functions.get(allocation, restore_in_4min)
     return restore_function(
-        noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale,
+        input_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
@@ -260,8 +312,8 @@ def restore_in_10min(*args, **kwargs):
     return restore_on_gpu(*args, **kwargs)
 
 def restore_on_gpu(
-    noisy_image,
-    input_image,
+    input_image_path,
+    denoise_image,
     prompt,
     a_prompt,
     n_prompt,
@@ -295,6 +347,7 @@ def restore_on_gpu(
     torch.cuda.set_device(SUPIR_device)
 
     with torch.no_grad():
+        input_image = Image.open(input_image_path)
         input_image = upscale_image(input_image, upscale, unit_resolution=32, min_size=min_size)
         LQ = np.array(input_image) / 255.0
         LQ = np.power(LQ, gamma_correction)
@@ -331,7 +384,7 @@ def restore_on_gpu(
         "The image has been enhanced successfully."
 
     # Only one image can be shown in the slider
-    return [noisy_image] + [results[0]], gr.update(label="Downloadable results", format=output_format, value=results), gr.update(value=information, visible=True), gr.update(visible=True)
+    return [input_image_path] + [results[0]], gr.update(label="Downloadable results", format=output_format, value=results), gr.update(value=information, visible=True), gr.update(visible=True)
 
 def load_and_reset(param_setting):
     print('Resetting parameters...')
@@ -418,6 +471,8 @@ with gr.Blocks() as interface:
             interactive=True
         )
 
+    gamma_correction = gr.Number(value=1.0, visible=False)  # Hidden component with default value 1.0
+
     with gr.Accordion("Advanced options", open=False):
         a_prompt = gr.Textbox(
             label="Additional image description",
@@ -655,7 +710,7 @@ with gr.Blocks() as interface:
             color_fix_type,
             diff_dtype,
             ae_dtype,
-            1.0,  # gamma_correction
+            gamma_correction,  # Use the hidden gamma_correction component
             linear_CFG,
             linear_s_stage2,
             spt_linear_CFG,
@@ -724,7 +779,7 @@ with gr.Blocks() as interface:
             color_fix_type,
             diff_dtype,
             ae_dtype,
-            1.0,  # gamma_correction
+            gamma_correction,  # Use the hidden gamma_correction component
             linear_CFG,
             linear_s_stage2,
             spt_linear_CFG,