Alternative implementation in Refiners

[limiteinductive]

Hi @Newbeeer !

Thank you for you amazing work and paper.

I've implemented the Restart method for the DDIM Scheduler in the library called Refiners. We are building Refiners, an open source (MIT), PyTorch-based framework made to easily train and run adapters on top of foundational models.

Demo:

Follow these install steps
Run the code snippet below:

import torch

from refiners.foundationals.latent_diffusion import StableDiffusion_1
from refiners.fluxion.utils import manual_seed


device = "cuda"

sd15 = StableDiffusion_1(device="cuda", dtype=torch.float16)
sd15.clip_text_encoder.load_from_safetensors("clip_text.safetensors")
sd15.lda.load_from_safetensors("lda.safetensors")
sd15.unet.load_from_safetensors("unet.safetensors")

with torch.no_grad():
    prompt = "a cute cat, detailed high-quality professional image"
    negative_prompt = "lowres, bad anatomy, bad hands, cropped, worst quality"

    clip_text_embedding = sd15.compute_clip_text_embedding(text=prompt, negative_text=negative_prompt)

    manual_seed(2)
    x = torch.randn(1, 4, 64, 64, device=device, dtype=torch.float16)
    
    restart = Restart(
        ldm=sd15,
        num_steps=10,
        num_iterations=2,
        start_time=0.1,
        end_time=2
    )
    

    for step in sd15.steps:
        x = sd15(
            x,
            step=step,
            clip_text_embedding=clip_text_embedding,
            condition_scale=7.5,
        )
        
        if step == restart.start_step:
            x = restart(
                x,
                clip_text_embedding=clip_text_embedding,
                condition_scale=8,
            )

    predicted_image = sd15.lda.decode_latents(x)

predicted_image.save("output.png")
print("done: see output.png")

With Restart:

Without Restart:

It would be great to have you feedback, since there is currently not many implementation of Restart. It would be great if you had some insights on how to support the DPM Scheduler.

Thank you!