use tqdm to track the current step

cosmos_transfer1/diffusion/diffusion/modules/res_sampler.py

@@ -27,6 +27,7 @@ from typing import Any, Callable, List, Literal, Optional, Tuple, Union
 
 import attrs
 import torch
+from tqdm import tqdm
 
 from cosmos_transfer1.diffusion.diffusion.functional.multi_step import get_multi_step_fn, is_multi_step_fn_supported
 from cosmos_transfer1.diffusion.diffusion.functional.runge_kutta import get_runge_kutta_fn, is_runge_kutta_fn_supported
@@ -204,7 +205,7 @@ def fori_loop(lower: int, upper: int, body_fun: Callable[[int, Any], Any], init_
         The final result after all iterations.
     """
     val = init_val
-    for i in range(lower, upper):
+    for i in tqdm(range(lower, upper)):
         val = body_fun(i, val)
     return val
 
@@ -251,7 +252,7 @@ def differential_equation_solver(
         def step_fn(
             i_th: int, state: Tuple[torch.Tensor, Optional[List[torch.Tensor]]]
         ) -> Tuple[torch.Tensor, Optional[List[torch.Tensor]]]:
-            log.info(f"Step [{i_th}/{num_step}]")
+            log.debug(f"Step [{i_th}/{num_step}]")
             input_x_B_StateShape, x0_preds = state
             sigma_cur_0, sigma_next_0 = sigmas_L[i_th], sigmas_L[i_th + 1]