Update app.py

app.py

@@ -67,16 +67,70 @@ def run_lora(prompt, selected_state, progress=gr.Progress(track_tqdm=True)):
             logging.error(f"API Error: {response.status_code}")
             raise gr.Error("API Error: Unable to fetch the image.")  # Raise a Gradio error here
 
-def apply_post_processing(image, downscale, limit_colors, grayscale, black_and_white):
+def postprocess(
+    image, 
+    enabled,
+    downscale,
+    need_rescale,
+    enable_color_limit,
+    number_of_colors,
+    quantization_method,
+    dither_method,
+    use_k_means,
+    is_grayscale,
+    number_of_shades,
+    quantization_method_grayscale,
+    dither_method_grayscale,
+    use_k_means_grayscale,
+    is_black_and_white,
+    is_inversed_black_and_white,
+    black_and_white_threshold,
+    use_color_palette,
+    palette_image,
+    palette_colors,
+    dither_method_palette
+):
+    if not enabled:
+        return image
+
     processed_image = image.copy()
+
     if downscale > 1:
         processed_image = downscale_image(processed_image, downscale)
-    if limit_colors:
-        processed_image = limit_colors(processed_image)
-    if grayscale:
+
+    if enable_color_limit:
+        processed_image = limit_colors(
+            image=processed_image,
+            limit=number_of_colors,
+            quantize=QUANTIZATION_METHODS[quantization_method],
+            dither=DITHER_METHODS[dither_method],
+            use_k_means=use_k_means
+        )
+
+    if is_grayscale:
         processed_image = convert_to_grayscale(processed_image)
-    if black_and_white:
-        processed_image = convert_to_black_and_white(processed_image)
+        processed_image = limit_colors(
+            image=processed_image,
+            limit=number_of_shades,
+            quantize=QUANTIZATION_METHODS[quantization_method_grayscale],
+            dither=DITHER_METHODS[dither_method_grayscale],
+            use_k_means=use_k_means_grayscale
+        )
+
+    if is_black_and_white:
+        processed_image = convert_to_black_and_white(processed_image, black_and_white_threshold, is_inversed_black_and_white)
+
+    if use_color_palette:
+        processed_image = limit_colors(
+            image=processed_image,
+            palette=palette_image,
+            palette_colors=palette_colors,
+            dither=DITHER_METHODS[dither_method_palette]
+        )
+
+    if need_rescale:
+        processed_image = resize_image(processed_image, image.size)
+
     return processed_image
 
 with gr.Blocks() as app: