fix: oor

src/improved_diffusion/gaussian_diffusion.py

@@ -679,6 +679,7 @@ class GaussianDiffusion:
                 caption[1].to(img.device),
             )  # (caption_state, caption_mask)
         my_bar = st.progress(0, text="Processing")
+        max_pro = len(indices)
         for pro, i in enumerate(indices):
             t = torch.tensor([i] * shape[0], device=device)
             with torch.no_grad():
@@ -694,7 +695,7 @@ class GaussianDiffusion:
                 )
                 yield out
                 img = out["sample"]
-            my_bar.progress(pro, text="Processing")
+            my_bar.progress((pro + 1) / max_pro, text="Processing")
         my_bar.empty()
 
     def p_sample_loop_langevin_progressive(