improving app

app.py

@@ -1,7 +1,4 @@
 import os
-import time
-
-os.environ["KERAS_BACKEND"] = "jax"
 
 import gradio as gr
 import jax
@@ -24,35 +21,68 @@ Two parameters that are interesting to control and adjust the amount of dehazing
 """
 
 
-@spaces.GPU
+def initialize_model():
+    config = Config.from_yaml(CONFIG_PATH)
+    diffusion_model = init(config)
+    return config, diffusion_model
+
+
+@spaces.GPU(duration=30)
+
+# Generator function for status updates
 def process_image(input_img, diffusion_steps, omega, omega_vent, omega_sept, eta):
     if input_img is None:
-        raise gr.Error(
-            "No input image was provided. Please select or upload an image before running."
+        yield (
+            gr.update(
+                value='<div style="background:#ffeeba;padding:10px;border-radius:6px;font-weight:bold;font-size:1.1em;color:#856404;">⚠️ No input image was provided. Please select or upload an image before running.</div>'
+            ),
+            None,
+        )
+        return
+    # Show loading message
+    yield (
+        gr.update(
+            value='<div style="background:#ffeeba;padding:10px;border-radius:6px;font-weight:bold;font-size:1.1em;color:#856404;">⏳ Loading model...</div>'
+        ),
+        None,
+    )
+
+    try:
+        config, diffusion_model = initialize_model()
+        params = config.params
+    except Exception as e:
+        yield (
+            gr.update(
+                value=f'<div style="background:#f8d7da;padding:10px;border-radius:6px;font-weight:bold;font-size:1.1em;color:#721c24;">❌ Error initializing model: {e}</div>'
+            ),
+            None,
         )
+        return
 
     def _prepare_image(image):
         resized = False
-
         if image.mode != "L":
             image = image.convert("L")
-
         orig_shape = image.size[::-1]
         h, w = diffusion_model.input_shape[:2]
         if image.size != (w, h):
             image = image.resize((w, h), Image.BILINEAR)
             resized = True
-
         image = np.array(image)
-
         image = image.astype(np.float32)
         image = image[None, ...]
         return image, resized, orig_shape
 
     try:
         image, resized, orig_shape = _prepare_image(input_img)
-    except Exception:
-        raise gr.Error("Something went wrong with preparing the input image.")
+    except Exception as e:
+        yield (
+            gr.update(
+                value=f'<div style="background:#f8d7da;padding:10px;border-radius:6px;font-weight:bold;font-size:1.1em;color:#721c24;">❌ Error preparing input image: {e}</div>'
+            ),
+            None,
+        )
+        return
 
     guidance_kwargs = {
         "omega": omega,
@@ -65,6 +95,12 @@ def process_image(input_img, diffusion_steps, omega, omega_vent, omega_sept, eta
     seed = jax.random.PRNGKey(config.seed)
 
     try:
+        yield (
+            gr.update(
+                value='<div style="background:#ffeeba;padding:10px;border-radius:6px;font-weight:bold;font-size:1.1em;color:#856404;">🌀 Running dehazing algorithm... (First time takes longer...) <span style="font-weight:normal;font-size:0.95em;">(first time takes longer)</span></div>'
+            ),
+            None,
+        )
         _, pred_tissue_images, *_ = run(
             hazy_images=image,
             diffusion_model=diffusion_model,
@@ -75,23 +111,32 @@ def process_image(input_img, diffusion_steps, omega, omega_vent, omega_sept, eta
             skeleton_params=params["skeleton_params"],
             batch_size=1,
             diffusion_steps=diffusion_steps,
-            initial_diffusion_step=params.get("initial_diffusion_step", 0),
             threshold_output_quantile=params.get("threshold_output_quantile", None),
             preserve_bottom_percent=params.get("preserve_bottom_percent", 30.0),
             bottom_transition_width=params.get("bottom_transition_width", 10.0),
             verbose=False,
         )
-    except Exception:
-        raise gr.Error("The algorithm failed to process the image.")
+    except Exception as e:
+        yield (
+            gr.update(
+                value=f'<div style="background:#f8d7da;padding:10px;border-radius:6px;font-weight:bold;font-size:1.1em;color:#721c24;">❌ The algorithm failed to process the image: {e}</div>'
+            ),
+            None,
+        )
+        return
 
     out_img = np.squeeze(pred_tissue_images[0])
     out_img = np.clip(out_img, 0, 255).astype(np.uint8)
     out_pil = Image.fromarray(out_img)
-    # Resize back to original input size if needed
     if resized and out_pil.size != (orig_shape[1], orig_shape[0]):
         out_pil = out_pil.resize((orig_shape[1], orig_shape[0]), Image.BILINEAR)
-    # Return tuple for ImageSlider: (input, output)
-    return (input_img, out_pil)
+    yield gr.update(value="Done!"), (input_img, out_pil)
+    yield (
+        gr.update(
+            value='<div style="background:#d4edda;padding:10px;border-radius:6px;font-weight:bold;font-size:1.1em;color:#155724;">✅ Done!</div>'
+        ),
+        (input_img, out_pil),
+    )
 
 
 slider_params = Config.from_yaml(SLIDER_CONFIG_PATH)
@@ -133,7 +178,7 @@ examples = [[img] for img in example_images]
 
 with gr.Blocks() as demo:
     gr.Markdown(description)
-    status = gr.Markdown("Initializing model, please wait...", visible=True)
+    status = gr.Markdown("", visible=True)
     with gr.Row():
         img1 = gr.Image(label="Input Image", type="pil", webcam_options=False)
         img2 = gr.ImageSlider(label="Dehazed Image", type="pil")
@@ -176,16 +221,6 @@ with gr.Blocks() as demo:
         )
     run_btn = gr.Button("Run")
 
-    def initialize_model():
-        time.sleep(0.5)  # Let UI update
-        config = Config.from_yaml(CONFIG_PATH)
-        diffusion_model = init(config)
-        params = config.params
-        return config, diffusion_model, params
-
-    config, diffusion_model, params = initialize_model()
-    status.visible = False
-
     run_btn.click(
         process_image,
         inputs=[
@@ -196,8 +231,9 @@ with gr.Blocks() as demo:
             omega_sept_slider,
             eta_slider,
         ],
-        outputs=[img2],
+        outputs=[status, img2],
+        queue=True,
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch()

main.py

@@ -1,9 +1,6 @@
 import copy
-import os
 from pathlib import Path
 
-os.environ["KERAS_BACKEND"] = "jax"
-
 import jax
 import keras
 import matplotlib.pyplot as plt
@@ -277,7 +274,6 @@ def run(
     skeleton_params: dict,
     batch_size: int = 4,
     diffusion_steps: int = 100,
-    initial_diffusion_step: int = 0,
     threshold_output_quantile: float = None,
     preserve_bottom_percent: float = 30.0,
     bottom_transition_width: float = 10.0,
@@ -306,7 +302,6 @@ def run(
             batch,
             n_samples=1,
             n_steps=diffusion_steps,
-            initial_step=initial_diffusion_step,
             seed=seed,
             verbose=True,
             per_pixel_omega=masks["per_pixel_omega"],