fix the add_noise function for dpm-multi et al

[yiyixuxu]

My PR to refactor schedulers #4986 broke img2img pipelines :(
I updated the add_noise incorrectly and never tested it with img2img and it was not caught by our current tests

this PR updateadd_noise to use sigmas

testing

I compared outputs from this branch vs the output from a branch where I reverted the commit that introduced this bug #4986

I used a single strength value 0.8, and I tested both use_karras_sigma=True and use_karras_sigma=False config. The results are identical.

testing use_karras_sigma=True

import requests
import numpy as np
from PIL import Image
from io import BytesIO
import torch

from diffusers import (
    StableDiffusionImg2ImgPipeline,
    DPMSolverSinglestepScheduler,
    DPMSolverMultistepScheduler,
    UniPCMultistepScheduler,
    DEISMultistepScheduler,
)

url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
response = requests.get(url)
init_image = Image.open(BytesIO(response.content)).convert("RGB")
init_image = init_image.resize((768, 512))

device = "cuda"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
).to(device)

scheduler_config = dict(pipe.scheduler.config)

seed = 33
strength = 0.8
#branch = "current"
branch = "revert-dpm-refactor"

# DPMSolverMultistepScheduler
sched_name = "dpm_miulti"
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    scheduler_config, use_karras_sigmas=True, algorithm_type="dpmsolver++")

generator = torch.Generator(device="cuda").manual_seed(seed)
images = pipe(
    prompt="a fantasy landscape",
    negative_prompt=None,
    image=init_image,
    strength=strength,
    num_inference_steps=15,
    generator=generator,
    guidance_scale=10,
    num_images_per_prompt=1
).images
image = images[0].save(f"test_5_out/[{sched_name}]_{strength:.2f}_[{branch}].png")

# DPMSolverMultistepScheduler(sde) 
sched_name = "dpm_miulti_sde"
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    scheduler_config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")

generator = torch.Generator(device="cuda").manual_seed(seed)
images = pipe(
    prompt="a fantasy landscape",
    negative_prompt=None,
    image=init_image,
    strength=strength,
    num_inference_steps=15,
    guidance_scale=10,
    generator=generator,
    num_images_per_prompt=1
).images
image = images[0].save(f"test_5_out/[{sched_name}]_{strength:.2f}_[{branch}].png")

# DPMSolverSingleScheduler
sched_name = "dpm_single"
pipe.scheduler = DPMSolverSinglestepScheduler.from_config(
    scheduler_config, use_karras_sigmas=True)

generator = torch.Generator(device="cuda").manual_seed(seed)
images = pipe(
    prompt="a fantasy landscape",
    negative_prompt=None,
    image=init_image,
    strength=strength,
    num_inference_steps=15,
    guidance_scale=10,
    generator=generator,
    num_images_per_prompt=1
).images
image = images[0].save(f"test_5_out/[{sched_name}]_{strength:.2f}_[{branch}].png")

# UniPCMultistepScheduler
sched_name = "unipc"
pipe.scheduler = UniPCMultistepScheduler.from_config(
    scheduler_config, use_karras_sigmas=True)

generator = torch.Generator(device="cuda").manual_seed(seed)
images = pipe(
    prompt="a fantasy landscape",
    negative_prompt=None,
    image=init_image,
    strength=strength,
    num_inference_steps=15,
    guidance_scale=10,
    generator=generator,
    num_images_per_prompt=1
).images
image = images[0].save(f"test_5_out/[{sched_name}]_{strength:.2f}_[{branch}].png")

# DEISMultistepScheduler
sched_name = "deis"
pipe.scheduler = DEISMultistepScheduler.from_config(
    scheduler_config, use_karras_sigmas=True)

generator = torch.Generator(device="cuda").manual_seed(seed)
images = pipe(
    prompt="a fantasy landscape",
    negative_prompt=None,
    image=init_image,
    strength=strength,
    num_inference_steps=15,
    generator=generator,
    guidance_scale=10,
    num_images_per_prompt=1
).images
image = images[0].save(f"test_5_out/[{sched_name}]_{strength:.2f}_[{branch}].png")

DEISMultistepScheduler

this PR	revert dpm-refactor PR

DPMSolverMultistepScheduler(sde)

this PR	revert dpm-refactor PR

DPMSolverMultistepScheduler

this PR	revert dpm-refactor PR

DPMSolverSingleScheduler

this PR	revert dpm-refactor PR

UniPCMultistepScheduler

this PR	revert dpm-refactor PR

testing use_karras_sigma=False

DEISMultistepScheduler

this PR	revert dpm-refactor PR

DPMSolverMultistepScheduler(sde)

this PR	revert dpm-refactor PR

DPMSolverMultistepScheduler

this PR	revert dpm-refactor PR

DPMSolverSingleScheduler

this PR	revert dpm-refactor PR

UniPCMultistepScheduler

this PR	revert dpm-refactor PR

[yiyixuxu]

More testing here

I also test various strength levels (from 0.1 ~ 0.95 with 0.05 interval) for each of the schedulers we updated, for both use_karras_sigmas=True and use_karras_sigmas=False,

I uploaded outputs here https://huggingface.co/datasets/YiYiXu/pr5158/tree/main

import requests
import numpy as np
from PIL import Image
from io import BytesIO
import torch

from diffusers import (
    StableDiffusionImg2ImgPipeline,
    DPMSolverSinglestepScheduler,
    DPMSolverMultistepScheduler,
    UniPCMultistepScheduler,
    DEISMultistepScheduler,
)

url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
response = requests.get(url)
init_image = Image.open(BytesIO(response.content)).convert("RGB")
init_image = init_image.resize((768, 512))

device = "cuda"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
).to(device)

scheduler_config = dict(pipe.scheduler.config)


# DPMSolverMultistepScheduler
sched_name = "dpm_miulti"
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    scheduler_config, use_karras_sigmas=True, algorithm_type="dpmsolver++")


for strength in torch.linspace(0.1,0.95,18):
    images = pipe(
        prompt="a fantasy landscape",
        negative_prompt=None,
        image=init_image,
        strength=strength,
        num_inference_steps=15,
        guidance_scale=10,
        num_images_per_prompt=1
    ).images
    image = images[0].save(f"test_2_out/[{sched_name}]_{strength:.2f}.png")

# DPMSolverMultistepScheduler(sde) 
sched_name = "dpm_miulti_sde"
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    scheduler_config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")


for strength in torch.linspace(0.1,0.95,18):
    images = pipe(
        prompt="a fantasy landscape",
        negative_prompt=None,
        image=init_image,
        strength=strength,
        num_inference_steps=15,
        guidance_scale=10,
        num_images_per_prompt=1
    ).images
    image = images[0].save(f"test_2_out/[{sched_name}]_{strength:.2f}.png")

# DPMSolverSingleScheduler
sched_name = "dpm_single"
pipe.scheduler = DPMSolverSinglestepScheduler.from_config(
    scheduler_config, use_karras_sigmas=True)


for strength in torch.linspace(0.1,0.95,18):
    images = pipe(
        prompt="a fantasy landscape",
        negative_prompt=None,
        image=init_image,
        strength=strength,
        num_inference_steps=15,
        guidance_scale=10,
        num_images_per_prompt=1
    ).images
    image = images[0].save(f"test_2_out/[{sched_name}]_{strength:.2f}.png")

# UniPCMultistepScheduler
sched_name = "unipc"
pipe.scheduler = UniPCMultistepScheduler.from_config(
    scheduler_config, use_karras_sigmas=True)


for strength in torch.linspace(0.1,0.95,18):
    images = pipe(
        prompt="a fantasy landscape",
        negative_prompt=None,
        image=init_image,
        strength=strength,
        num_inference_steps=15,
        guidance_scale=10,
        num_images_per_prompt=1
    ).images
    image = images[0].save(f"test_2_out/[{sched_name}]_{strength:.2f}.png")

# DEISMultistepScheduler
sched_name = "deis"
pipe.scheduler = DEISMultistepScheduler.from_config(
    scheduler_config, use_karras_sigmas=True)

for strength in torch.linspace(0.1,0.95,18):
    images = pipe(
        prompt="a fantasy landscape",
        negative_prompt=None,
        image=init_image,
        strength=strength,
        num_inference_steps=15,
        guidance_scale=10,
        num_images_per_prompt=1
    ).images
    image = images[0].save(f"test_2_out/[{sched_name}]_{strength:.2f}.png")

[pcuenca]

Works for me (tested DPMSolverMultistepScheduler). The SDXL refiner step was broken, wasn't it? Did we receive any reports?

I'd suggest we merge this soon but then create tests (or fix them) as a followup, to ensure we catch similar problems in the future.

[rvorias]

Thanks for fixing this so fast!

[patrickvonplaten]

Very clean fix - thanks a lot @yiyixuxu