log current step

cosmos_transfer1/diffusion/diffusion/modules/res_sampler.py

@@ -30,6 +30,7 @@ import torch
 
 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
+from cosmos_transfer1.utils import log
 from cosmos_transfer1.utils.ddp_config import make_freezable
 
 COMMON_SOLVER_OPTIONS = Literal["2ab", "2mid", "1euler"]
@@ -250,6 +251,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}]")
             input_x_B_StateShape, x0_preds = state
             sigma_cur_0, sigma_next_0 = sigmas_L[i_th], sigmas_L[i_th + 1]