[Community] Implementation of the IADB community pipeline

[tchambon]

This PR is for the implementation of a new community pipeline.
It is for the implementation of the α-(de)Blending: a Minimalist Deterministic Diffusion Model paper.
It is a simple and minimalist diffusion model.

The following code shows how to use the IADB pipeline to generate images using a pretrained celebahq-256 model.

pipeline_iadb = DiffusionPipeline.from_pretrained("thomasc4/iadb-celebahq-256", custom_pipeline='iadb')

pipeline_iadb = pipeline_iadb.to('cuda')

output = pipeline_iadb(batch_size=4,num_inference_steps=128)
for i in range(len(output[0])):
    plt.imshow(output[0][i])
    plt.show()

Sampling with the IADB formulation is easy, and can be done in a few lines (the pipeline already implements it):

def sample_iadb(model, x0, nb_step):
    x_alpha = x0
    for t in range(nb_step):
        alpha = (t/nb_step)
        alpha_next =((t+1)/nb_step)

        d = model(x_alpha, torch.tensor(alpha, device=x_alpha.device))['sample']
        x_alpha = x_alpha + (alpha_next-alpha)*d

    return x_alpha

The training loop is also straightforward:

# Training loop
while True:
    x0 = sample_noise()
    x1 = sample_dataset()

    alpha = torch.rand(batch_size)

    # Blend
    x_alpha = (1-alpha) * x0 + alpha * x1

    # Loss
    loss = torch.sum((D(x_alpha, alpha)- (x1-x0))**2)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[patrickvonplaten]

Very cool!