jina-ai / discoart

🪩 Create Disco Diffusion artworks in one line
Other
3.84k stars 249 forks source link

Add progress callback #157

Open entmike opened 2 years ago

entmike commented 2 years ago

Similar to what was added here recently:

https://github.com/jina-ai/discoart/blob/main/discoart/persist.py#L63

              if is_save_step:
                if is_image_output:
                    if cur_t == -1:
                        f_name = os.path.join(output_dir, f'{_nb}-done-{k}.png')
                    else:
                        f_name = os.path.join(output_dir, f'{_nb}-step-{j}-{k}.png')
                    c.save_uri_to_file(f_name)

                    if callable(image_callback):
                        image_callback(f_name)

                da[k].chunks.append(c)

It would be good to also have a more generalized callback for progress updates, as well as it not being bound to the frequency of saves (is_save_step for instance)

Proposed pseudo-code:

              if is_progress_step:
                    if callable(progress_callback):
                        progress_callback( any progress amount, maybe image as bytecode etc passed in an event/dict format )