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,
                     )
                 )