Fix decoding batch size bug

app.py

@@ -1,18 +1,4 @@
-"""
-Gradio app to compare multiple inference methods for Monetico model.
-
-This file wires your existing inference functions (infer_pretrained, infer_smc_grad)
-into a single UI with one shared prompt and per-method collapsed setting panels.
-
-Place this file at repository root (next to src/) and run:
-    python app.py
-
-Notes:
-- The code assumes your module that contains infer_pretrained and infer_smc_grad
-  is importable (e.g. package root with src/ on PYTHONPATH). Adjust imports if needed.
-- Defaults provided are reasonable starting points; tweak as you like.
-"""
-
+import traceback
 import gradio as gr
 
 # Import your inference functions and dataclasses
@@ -79,6 +65,7 @@ def run_inference_all(
         pretrained_output = infer_pretrained(pretrained_cfg, device=DEVICE)
         pretrained_images = pretrained_output.images
     except Exception as e:
+        traceback.print_exc()
         pretrained_images = []
         pretrained_output = None
         pretrained_error = f"Pretrained inference error: {e}"
@@ -107,6 +94,7 @@ def run_inference_all(
         smc_grad_output = infer_smc_grad(smc_grad_cfg, device=DEVICE)
         # The above line is defensive; simpler: pass smc_grad_device value used by gradio - will be provided.
     except Exception as e:
+        traceback.print_exc()
         smc_grad_images = []
         smc_grad_output = None
         smc_grad_error = f"SMC inference error: {e}"

src/smc/pipeline.py

@@ -242,7 +242,7 @@ class Pipeline(
                         sample = F.gumbel_softmax(tmp_logits, tau=tau, hard=True)
                         if use_continuous_formulation:
                             sample = gamma * sample + (ONE - MASK) * latents_one_hot
-                        sample = self._decode_one_hot_latents(sample, batch_p, height, width, "pt")
+                        sample = self._decode_one_hot_latents(sample, height, width, "pt")
                         tmp_rewards[:, phi_i] = reward_fn(sample)
                     tmp_rewards = logmeanexp(tmp_rewards * scale_cur, dim=-1) / scale_cur
                     
@@ -332,7 +332,7 @@ class Pipeline(
                     tmp_logp_x0 = self.model._subs_parameterization(tmp_logits, latents_batch)
                     for phi_i in range(phi):
                         sample = F.gumbel_softmax(tmp_logp_x0, tau=tau, hard=True).argmax(dim=-1)
-                        sample = self._decode_latents(sample, batch_p, height, width, "pt")
+                        sample = self._decode_latents(sample, height, width, "pt")
                         tmp_rewards[:, phi_i] = reward_fn(sample)
                     tmp_rewards = logmeanexp(tmp_rewards * scale_cur, dim=-1) / scale_cur
                 
@@ -427,7 +427,7 @@ class Pipeline(
                         tmp_logp_x0 = self._subs_parameterization(tmp_logits, latents_batch)
                     for phi_i in range(phi):
                         sample = F.gumbel_softmax(tmp_logp_x0, tau=tau, hard=True).argmax(dim=-1)
-                        sample = self._decode_latents(sample, batch_p, height, width, "pt")
+                        sample = self._decode_latents(sample, height, width, "pt")
                         tmp_rewards[:, phi_i] = reward_fn(sample)
                     tmp_rewards = logmeanexp(tmp_rewards * scale_cur, dim=-1) / scale_cur
                 
@@ -517,7 +517,7 @@ class Pipeline(
         for j in range(0, total_particles, batch_p):
             latents_batch = latents[j:j+batch_p]
             with torch.no_grad():
-                sample = self._decode_latents(latents_batch, batch_p, height, width, "pt")
+                sample = self._decode_latents(latents_batch, height, width, "pt")
                 tmp_rewards = reward_fn(sample)
                 rewards[j:j+batch_p] = tmp_rewards
                 log_twist[j:j+batch_p] = tmp_rewards * scale_cur
@@ -563,7 +563,7 @@ class Pipeline(
         for j in range(0, batches, batch_p):
             latents_batch = latents[j:j+batch_p]
             outputs.extend(
-                self._decode_latents(latents_batch, batch_p, height, width, output_type) # type: ignore
+                self._decode_latents(latents_batch, height, width, output_type) # type: ignore
             )
         if output_type == "pt":
             outputs = torch.stack(outputs, dim=0)
@@ -602,7 +602,8 @@ class Pipeline(
         tmp_logits = torch.cat([tmp_logits, pad_logits], dim=-1)
         return tmp_logits
 
-    def _decode_latents(self, latents, batch_size, height, width, output_type):
+    def _decode_latents(self, latents, height, width, output_type):
+        batch_size = latents.shape[0]
         if output_type == "latent":
             output = latents
         else:
@@ -624,7 +625,8 @@ class Pipeline(
                 self.vqvae.half()
         return output
             
-    def _decode_one_hot_latents(self, latents_one_hot, batch_size, height, width, output_type):
+    def _decode_one_hot_latents(self, latents_one_hot, height, width, output_type):
+        batch_size = latents_one_hot.shape[0]
         shape = (
             batch_size,
             height // self.vae_scale_factor,