Make the sampling loop of pipelines modular, with designated inputs that can be handled iteratively with functions and a default sampler step function

[AmericanPresidentJimmyCarter]

Model/Pipeline/Scheduler description

Related to #7761 .

This is an effective replacement for the existing sampling loop function and the many, many kwargs that were made to allow the user to control it or inject callbacks into it.

class SamplingInput:
    def __init__(self, img, text_embedding, unet, timestep=None, **kwargs):
        self.img = img
        self.text_embedding = text_embedding
        self.unet = unet
        self.timestep = timestep

# ... lots of other code ...

        inp = SamplingInput(img, text_embedding, unet)
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                inp.timestep = t
                for sampling_function in self.sampling_functions:
                    inp = sampling_function(inp)

        output_img = inp.img

This will give the end user complete control of the sampling loop, allow the repo to add "official inline sampling functions" like report an image to an endpoint so that the user can view intermediate steps, etc.

We can add an argument sampling_functions: list[Callable]=[default_sampling_function] into the __call__ as a new, backwards compatible kwarg.

This requires a rewrite of all of the pipelines, but as it is a backwards compatible change it can be introduced to any of the more popular pipelines first.

[Beinsezii]

Should the default functionality be a single function or a list of smaller functions do you think?

If it was a list such as [f_a, f_b, f_c] you could easily hook into any stage by simply funcs.insert(f_a1, 1) or funcs[1] = f_b_custom, but that makes the MVP more complex.

[AmericanPresidentJimmyCarter]

You could always decompose the default function to have multiple other functions that are called inside it, which the user can also import and use in their functions if needed.

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.