inference_runtime.py

"""Progress reporting and scheduler selection shared with the standalone Gradio app."""


def set_scheduler(pipe, shift):
    pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config, shift=shift)


def generate_with_progress(pipe, report_stage, report_step, **kwargs):
    progress_bar = pipe.progress_bar

    def progress(items):
        report_stage("denoise")
        for index, timestep in enumerate(progress_bar(items)):
            yield timestep
            report_step(pipe, index, timestep, {})
        if kwargs.get("output_type", "pil") != "latent":
            report_stage("decode")

    pipe.progress_bar = progress
    try:
        report_stage("encode")
        return pipe(**kwargs)
    finally:
        pipe.progress_bar = progress_bar