A problem on the inference using refiner

[TheLovesOfLadyPurple]

Describe the bug

If you follow the official guidance for sdxl and use refiner, like the code in https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0, and use the Karras time schedule, then it will use the weird inference time schedule. For example, for a 20-step inference and ending the first stage inference at denoising_end=0.8, then it will do only 12 steps of inference in the first stage. But if you don't use karras, it will do 16 steps of inference in the first stage.

That means, when we denote denoising_end=0.8, we don't end the inference when 80% of the inference process is finished.

If 0.8 means the inference returns a noisy image and that noisy image can be sampled by a forward diffusion process where 80% of the diffusion process has not been finished, then, in the 20-step inference using Karras, the current method stops on the 285th step rather than the 197th step, which makes whether this explanation is correct not clear enough and the behavior unpredictable.

When I mean predictable, I mean the user knows what they are doing. For example, if the 0.8 means the inference stops when 80% of the inference is complete, then the user knows what that 0.8 means. However, if the 0.8 means the inference stops when 80% of the inference finished, then, when the inference stops, the variance of the noise in the image will be in a vast range if we choose different time schedules. And that may not be friendly enough to those users who don't have enough knowledge on this scope.

The information in the following table provided by @4ktLuffy

| Schedule | Last kept timestep | Trajectory completed |
|:---|---:|----:|
| Leading | 201.0 |79.9% |
| Exponential | 204.6 |79.5% |
| Linspace | 210.3 | 79.0% |
| Trailing | 249.0 | 75.1% |
| Karras | 272.0 | 72.8% |

,which exhibit different inference result if we use current method to do 20 step inference, and denoising_end=0.8 And if 0.8 means the inference stops when 80% of the inference finished, then the table is:

| Schedule | Current stop steps | Current noise variance at handoff | new stop step on Step-count-based handoff | Noise variance at new handoff |
|:---|---:|---:|---:|---:|
| Linspace | 16/20 | 3.2% | 16/20 | 3.2% |
| Leading | 16/20 | 4.2% | 16/20 | 4.2% |
| Trailing | 16/20 | 3.9% | 16/20 | 3.9% |
| Karras | 12/20 | 5.0% | 16/20 | 1.3% |
| Exponential | 11/20 | 3.9% | 16/20 | 0.9% |

So, what kind of the behaviour is willing? Stop when 80% of the inference finished? Or stop when the noisy image can be sampled by a forward diffusion process where 80% of the diffusion process has not been finished? In the given situation that we do the 20 step inference, should the solver prefer to use the timestep that is closest to 200 in the given timesteps?

Reproduction

"""Minimal SDXL inference with StableDiffusionXLPipeline and a Karras scheduler."""

import argparse
import functools
from pathlib import Path

import torch
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionXLImg2ImgPipeline,
    StableDiffusionXLPipeline,
)
from huggingface_hub import hf_hub_download


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--prompt", default="A majestic lion jumping from a big stone at night")
    parser.add_argument("--negative-prompt", default="")
    parser.add_argument("--base-model", default="stabilityai/stable-diffusion-xl-base-1.0")
    parser.add_argument("--refiner-model", default="stabilityai/stable-diffusion-xl-refiner-1.0")
    parser.add_argument("--vae-repo", default="madebyollin/sdxl-vae-fp16-fix")
    parser.add_argument("--vae-filename", default="sdxl_vae.safetensors")
    parser.add_argument("--steps", type=int, default=20)
    parser.add_argument("--high-noise-frac", type=float, default=0.8)
    parser.add_argument("--guidance-scale", type=float, default=5.0)
    parser.add_argument("--height", type=int, default=1024)
    parser.add_argument("--width", type=int, default=1024)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--use-karras", action=argparse.BooleanOptionalAction, default=False)
    parser.add_argument("--skip-refiner", action="store_true")
    parser.add_argument("--output", default="outputs/sdxl-karras/sdxl-karras.png")
    return parser.parse_args()


def build_scheduler(pipe, use_karras):
    pipe.scheduler = EulerDiscreteScheduler.from_config(
        pipe.scheduler.config,
        use_karras_sigmas=use_karras,
    )
    return pipe


def attach_timestep_output_closure(pipe, label):
    original_forward = pipe.unet.forward
    recorded_timesteps = []

    @functools.wraps(original_forward)
    def wrapped_forward(sample, timestep, *args, **kwargs):
        if torch.is_tensor(timestep):
            timestep_value = timestep.detach().cpu().flatten().tolist()
            if len(timestep_value) == 1:
                timestep_value = timestep_value[0]
        else:
            timestep_value = timestep
        recorded_timesteps.append(timestep_value)
        return original_forward(sample, timestep, *args, **kwargs)

    pipe.unet.forward = wrapped_forward
    return label, recorded_timesteps


def load_vae(args, dtype):
    vae_path = hf_hub_download(repo_id=args.vae_repo, filename=args.vae_filename, cache_dir=".")
    return AutoencoderKL.from_single_file(vae_path, torch_dtype=dtype)


def main():
    args = parse_args()

    if not torch.cuda.is_available():
        raise RuntimeError("This script requires a CUDA-capable PyTorch installation.")
    if args.height % 8 != 0 or args.width % 8 != 0:
        raise ValueError("--height and --width must both be divisible by 8.")
    if not args.skip_refiner and not 0.0 < args.high_noise_frac < 1.0:
        raise ValueError("--high-noise-frac must be between 0 and 1 when the refiner is enabled.")

    device = torch.device("cuda")
    dtype = torch.float16
    vae = load_vae(args, dtype).to(device)

    base = StableDiffusionXLPipeline.from_pretrained(
        args.base_model,
        vae=vae,
        torch_dtype=dtype,
        variant="fp16",
        use_safetensors=True,
    ).to(device)
    build_scheduler(base, args.use_karras)
    timestep_records = {}
    base_label, base_timesteps = attach_timestep_output_closure(base, "base")
    timestep_records[base_label] = base_timesteps

    generator = torch.Generator(device=device).manual_seed(args.seed)

    with torch.inference_mode():
        if args.skip_refiner:
            image = base(
                prompt=args.prompt,
                negative_prompt=args.negative_prompt,
                num_inference_steps=args.steps,
                guidance_scale=args.guidance_scale,
                height=args.height,
                width=args.width,
                generator=generator,
            ).images[0]
        else:
            refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
                args.refiner_model,
                text_encoder_2=base.text_encoder_2,
                vae=base.vae,
                torch_dtype=dtype,
                variant="fp16",
                use_safetensors=True,
            ).to(device)
            build_scheduler(refiner, args.use_karras)
            refiner_label, refiner_timesteps = attach_timestep_output_closure(refiner, "refiner")
            timestep_records[refiner_label] = refiner_timesteps

            latent_image = base(
                prompt=args.prompt,
                negative_prompt=args.negative_prompt,
                num_inference_steps=args.steps,
                guidance_scale=args.guidance_scale,
                height=args.height,
                width=args.width,
                denoising_end=args.high_noise_frac,
                output_type="latent",
                generator=generator,
            ).images

            image = refiner(
                prompt=args.prompt,
                negative_prompt=args.negative_prompt,
                num_inference_steps=args.steps,
                guidance_scale=args.guidance_scale,
                denoising_start=args.high_noise_frac,
                image=latent_image,
                generator=generator,
            ).images[0]

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    image.save(output_path)
    print(f"Recorded timestep inputs to UNet: {timestep_records}")
    print(f"Saved image to {output_path}")


if __name__ == "__main__":
    main()

Logs

The log of running the code is:
use karras: UNet: {'base': [951.0, 913.7394409179688, 873.1348266601562, 828.625244140625, 779.5436401367188, 725.1243286132812, 664.55517578125, 597.1300659179688, 522.5737915039062, 441.6042785644531, 356.6148986816406, 272.02899169921875], 'refiner': [193.7161407470703, 127.3919906616211, 76.73764038085938, 42.189613342285156, 21.15185546875, 9.592663764953613, 3.763125419616699, 1.0000001192092896]} (19 steps)
Don't use Karras: UNet: {'base': [951.0, 901.0, 851.0, 801.0, 751.0, 701.0, 651.0, 601.0, 551.0, 501.0, 451.0, 401.0, 351.0, 301.0, 251.0, 201.0], 'refiner': [151.0, 101.0, 51.0, 1.0]} (20 steps)

System Info

• 🤗 Diffusers version: 0.36.0
• Platform: Windows-10-10.0.26200-SP0
• Running on Google Colab?: No
• Python version: 3.9.25
• PyTorch version (GPU?): 2.8.0+cu129 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 0.36.0
• Transformers version: 4.57.6
• Accelerate version: 1.10.1
• PEFT version: 0.17.1
• Bitsandbytes version: not installed
• Safetensors version: 0.7.0
• xFormers version: 0.0.32.post2
• Accelerator: NVIDIA GeForce RTX 5060 Ti, 16311 MiB
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: not

Who can help?

@yiyixuxu @sayakpaul @DN6

[Sravan1011]

@TheLovesOfLadyPurple @sayakpaul @4ktLuffy can I work in this or try this out ?

[TheLovesOfLadyPurple]

@TheLovesOfLadyPurple @sayakpaul @4ktLuffy can I work in this or try this out ?

Actually, I already have a fix; that is why I and @4ktLuffy have Table 2, but I am not sure whether this fix is what the maintainer wants. According to my discussion with @4ktLuffy, I think we should first wait for the maintainer's reply to see what his requirements are.