jina-ai / discoart

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

Add progress_update hook to do_run #159

Open entmike opened 1 year ago

entmike commented 1 year ago

Proposed addition of progress_update hook mentioned here https://github.com/jina-ai/discoart/issues/157.

Sample test script:

import sys, os, multiprocessing
from types import SimpleNamespace

def pcb(progress):
    print(progress)

def mutate(**kwargs):
    from discoart import runner
    from discoart.config import (
        load_config,
        print_args_table
    )
    from discoart.helper import (
        load_diffusion_model,
        load_clip_models,
        load_secondary_model,
        get_device,
        free_memory,
        show_result_summary,
        get_output_dir,
    )
    events = tuple(
        kwargs.pop(v, multiprocessing.Event()) or multiprocessing.Event()
        for v in ('skip_event', 'stop_event')
    )
    _clip_models_cache = {}
    _args = load_config(user_config=kwargs)
    print_args_table(_args)
    _args = SimpleNamespace(**_args)

    device = get_device()
    model, diffusion = load_diffusion_model(_args, device=device)

    clip_models = load_clip_models(
        device,
        enabled=_args.clip_models,
        clip_models=_clip_models_cache,
        text_clip_on_cpu=_args.text_clip_on_cpu,
    )
    secondary_model = load_secondary_model(_args, device=device)

    free_memory()
    is_exit0 = False
    runner.do_run(
        _args,
        (model, diffusion, clip_models, secondary_model),
        device=device,
        events=events,
        progress_callback = pcb
    )

mutate()