Add progress_update hook to do_run

discoart/persist.py

@@ -32,7 +32,9 @@ def _sample(
  is_save_gif,
  is_image_output,
  is_display_step,
+ is_progress_step,
  image_callback,
+ progress_callback
 ):
  with threading.Lock():
  is_sampling_done.clear()
@@ -65,6 +67,14 @@ def _sample(
 
  da[k].chunks.append(c)
 
+ if is_progress_step:
+ if callable(progress_callback):
+ progress_callback({
+ "sample" : sample['pred_xstart'],
+ "cur_t" : cur_t,
+ "j" : j
+ })
+
  if is_save_gif and is_image_output:
  da_gif[k].chunks.append(c)
 

discoart/resources/default.yml

@@ -37,6 +37,7 @@ cut_icgray_p: "[0.2]*400+[0]*600"
 cut_ic_pow: 1.
 
 save_rate: 20
+progress_rate: 10
 gif_fps: 20
 gif_size_ratio: 0.5
 n_batches: 4

discoart/resources/docstrings.yml

@@ -123,7 +123,8 @@ cut_icgray_p: |
 
 save_rate: |
  [DiscoArt] The number of steps to save intermediate results. It is a replacement to original `display_rate` parameter. Set it to -1 for not saving any intermediate result.
-
+progress_rate: |
+ [DiscoArt] The number of steps to per progress update callback function
 n_batches: |
  This variable sets the number of still images you want DD to create. If you are using an animation mode (see below for details) DD will ignore n_batches and create a single set of animated frames based on the animation settings.
 batch_name: |

discoart/runner.py

@@ -34,7 +34,9 @@
 
 
 def do_run(
- args, models, device, events, image_callback: Optional[Callable[[str], None]] = None
+ args, models, device, events,
+ image_callback: Optional[Callable[[str], None]] = None,
+ progress_callback: Optional[Callable[[str], None]] = None
 ) -> 'DocumentArray':
  skip_event, stop_event = events
 
@@ -422,6 +424,7 @@ def cond_fn(x, t, **kwargs):
  is_save_step = args.save_rate > 0 and j % args.save_rate == 0
  is_complete = cur_t == -1
  is_display_step = args.display_rate > 0 and j % args.display_rate == 0
+ is_progress_step = args.progress_rate > 0 and j % args.progress_rate == 0
 
  threads.append(
  _sample_thread(
@@ -439,7 +442,9 @@ def cond_fn(x, t, **kwargs):
  args.gif_fps > 0,
  args.image_output,
  is_display_step,
+ is_progress_step,
  image_callback,
+ progress_callback,
  )
  )