restore step argument to step_callback

ldm/models/diffusion/ksampler.py

@@ -101,7 +101,7 @@ class KSampler(Sampler):
     ):
         def route_callback(k_callback_values):
             if img_callback is not None:
-                img_callback(k_callback_values['x'])
+                img_callback(k_callback_values['x'],k_callback_values['i'])
 
         # sigmas = self.model.get_sigmas(S)
         # sigmas are now set up in make_schedule - we take the last steps items

ldm/models/diffusion/sampler.py

@@ -272,7 +272,7 @@ class Sampler(object):
             if callback:
                 callback(i)
             if img_callback:
-                img_callback(img)
+                img_callback(img,step)
 
             if index % log_every_t == 0 or index == total_steps - 1:
                 intermediates['x_inter'].append(img)