initial commit

app.py

@@ -52,11 +52,32 @@ def predict_location(image, model_name, cfg=0.0, num_samples=256):
 
         pipe = PIPES[model_name]
 
-        # Get regular predictions
-        predicted_gps = pipe(img, batch_size=num_samples, cfg=cfg, num_steps=16)
+        # Create a progress bar
+        progress_bar = st.progress(0)
+        status_text = st.empty()
+
+        def update_progress(step, total_steps):
+            progress = float(step) / float(total_steps)
+            progress_bar.progress(progress)
+            status_text.text(f"Sampling step {step + 1}/{total_steps}")
+
+        # Get regular predictions with progress updates
+        predicted_gps = pipe(
+            img, 
+            batch_size=num_samples, 
+            cfg=cfg, 
+            num_steps=16,
+            callback=update_progress
+        )
 
         # Get single high-confidence prediction
+        status_text.text("Generating high-confidence prediction...")
         high_conf_gps = pipe(img, batch_size=1, cfg=2.0, num_steps=16)
+        
+        # Clear the status text and progress bar
+        status_text.empty()
+        progress_bar.empty()
+        
         return {
             "lat": predicted_gps[:, 0].astype(float).tolist(),
             "lon": predicted_gps[:, 1].astype(float).tolist(),
@@ -234,14 +255,13 @@ def main():
                 st.markdown("</div>", unsafe_allow_html=True)
 
                 if st.button("🔍 Predict Location", key="predict_upload"):
-                    with st.spinner("🌍 Analyzing image and predicting locations..."):
-                        predictions = predict_location(
-                            original_image,
-                            model_name=model_name,
-                            cfg=cfg_value,
-                            num_samples=num_samples,
-                        )
-                        st.session_state["predictions"] = predictions
+                    predictions = predict_location(
+                        original_image,
+                        model_name=model_name,
+                        cfg=cfg_value,
+                        num_samples=num_samples,
+                    )
+                    st.session_state["predictions"] = predictions
 
         with tab2:
             url = st.text_input("Enter image URL:", key="image_url")
@@ -261,16 +281,13 @@ def main():
                     st.markdown("</div>", unsafe_allow_html=True)
 
                     if st.button("🔍 Predict Location", key="predict_url"):
-                        with st.spinner(
-                            "🌍 Analyzing image and predicting locations..."
-                        ):
-                            predictions = predict_location(
-                                image,
-                                model_name=model_name,
-                                cfg=cfg_value,
-                                num_samples=num_samples,
-                            )
-                            st.session_state["predictions"] = predictions
+                        predictions = predict_location(
+                            image,
+                            model_name=model_name,
+                            cfg=cfg_value,
+                            num_samples=num_samples,
+                        )
+                        st.session_state["predictions"] = predictions
 
         with tab3:
             examples = load_example_images()
@@ -290,17 +307,14 @@ def main():
                         help=f"Click to predict location for {name}",
                         use_container_width=True,
                     ):
-                        with st.spinner(
-                            "🌍 Analyzing image and predicting locations..."
-                        ):
-                            predictions = predict_location(
-                                original_image,
-                                model_name=model_name,
-                                cfg=cfg_value,
-                                num_samples=num_samples,
-                            )
-                            st.session_state["predictions"] = predictions
-                            st.rerun()
+                        predictions = predict_location(
+                            original_image,
+                            model_name=model_name,
+                            cfg=cfg_value,
+                            num_samples=num_samples,
+                        )
+                        st.session_state["predictions"] = predictions
+                        st.rerun()
 
                     st.image(display_image, caption=name, use_container_width=True)
             st.markdown("</div>", unsafe_allow_html=True)

models/samplers/riemannian_flow_sampler.py

@@ -13,6 +13,7 @@ def riemannian_flow_sampler(
     cfg_rate=0,
     generator=None,
     return_trajectories=False,
+    callback=None,
 ):
     if scheduler is None:
         raise ValueError("Scheduler must be provided")
@@ -35,13 +36,13 @@ def riemannian_flow_sampler(
             if cfg_rate > 0 and conditioning_keys is not None:
                 stacked_batch["y"] = torch.cat([x_cur, x_cur], dim=0)
                 stacked_batch["gamma"] = gamma_now.expand(x_cur.shape[0] * 2)
-                denoised_all = net(stacked_batch)
+                denoised_all = net(stacked_batch, current_step=step)
                 denoised_cond, denoised_uncond = denoised_all.chunk(2, dim=0)
                 denoised = denoised_cond * (1 + cfg_rate) - denoised_uncond * cfg_rate
             else:
                 batch["y"] = x_cur
                 batch["gamma"] = gamma_now.expand(x_cur.shape[0])
-                denoised = net(batch)
+                denoised = net(batch, current_step=step)
 
         dt = gamma_next - gamma_now
         x_next = x_cur + dt * denoised  # manifold.expmap(x_cur, dt * denoised)

pipe.py

@@ -216,6 +216,7 @@ class PlonkPipeline:
         scheduler=None,
         cfg=0,
         generator=None,
+        callback=None,
     ):
         """Sample from the model given conditioning.
 
@@ -228,6 +229,7 @@ class PlonkPipeline:
             scheduler: Custom scheduler function (uses default if not provided)
             cfg: Classifier-free guidance scale (default 15)
             generator: Random number generator
+            callback: Optional callback function to report progress (step, total_steps)
 
         Returns:
             Sampled GPS coordinates after postprocessing
@@ -264,26 +266,28 @@ class PlonkPipeline:
         sampler = self.sampler
         if scheduler is None:
             scheduler = self.scheduler
+
         # Sample from model
         if num_steps is None:
-            output = sampler(
-                self.model,
-                batch,
-                conditioning_keys="emb",
-                scheduler=scheduler,
-                cfg_rate=cfg,
-                generator=generator,
-            )
-        else:
-            output = sampler(
-                self.model,
-                batch,
-                conditioning_keys="emb",
-                scheduler=scheduler,
-                num_steps=num_steps,
-                cfg_rate=cfg,
-                generator=generator,
-            )
+            num_steps = 16  # Default number of steps
+
+        # Create a wrapper for the model that updates progress
+        def model_with_progress(*args, **kwargs):
+            step = kwargs.pop('current_step', 0)
+            if callback:
+                callback(step, num_steps)
+            return self.model(*args, **kwargs)
+
+        output = sampler(
+            model_with_progress,
+            batch,
+            conditioning_keys="emb",
+            scheduler=scheduler,
+            num_steps=num_steps,
+            cfg_rate=cfg,
+            generator=generator,
+            callback=callback,
+        )
 
         # Apply postprocessing and return
         output = self.postprocessing(output)