Fix SDXL Refiner with Higher Order Schedulers

[Beinsezii]

What does this PR do?

Fixes SDXL refiner w/ higher order schedulers by changing the strange hardcoded order==2 check with a simple tensor stride.

Standalone script using a scheduler with order=15

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "diffusers",
#     "skrample==0.6.*",
#     "torch>=2.11.0",
#     "transformers>=5.5.3",
# ]
#
# [tool.uv.sources]
# diffusers = { git = "https://github.com/Beinsezii/diffusers.git", rev = "beinsezii/fix_denoise_order" }
# ///

import torch
from skrample.diffusers import RKUltraWrapperScheduler

from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import StableDiffusionXLPipeline
from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img import (
    StableDiffusionXLImg2ImgPipeline,
)


with torch.inference_mode():
    DEVICE = "cuda"
    DTYPE = torch.bfloat16
    BASE = "stabilityai/stable-diffusion-xl-base-1.0"
    REFINER = "stabilityai/stable-diffusion-xl-refiner-1.0"

    PROMPT = "bright high resolution dslr photograph of a kitten in a field of lavender flowers"
    CFG: float = 8
    RATIO: float = 0.6
    ORDER: int = 15

    base = StableDiffusionXLPipeline.from_pretrained(BASE).to(device=DEVICE, dtype=DTYPE)
    refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
        REFINER,
        vae=base.vae,
        text_encoder_2=base.text_encoder_2,
    ).to(device=DEVICE, dtype=DTYPE)
    base.scheduler = RKUltraWrapperScheduler.from_diffusers_config(base.scheduler.config, sampler_order=ORDER)
    refiner.scheduler = RKUltraWrapperScheduler.from_diffusers_config(refiner.scheduler.config, sampler_order=ORDER)

    for steps in 5, 10:
        partial = base(
            prompt=PROMPT,
            guidance_scale=CFG,
            num_inference_steps=steps,
            denoising_end=RATIO,
            output_type="latent",
        ).images[0]
        refiner(
            image=partial,
            prompt=PROMPT,
            guidance_scale=CFG,
            num_inference_steps=steps,
            denoising_start=RATIO,
        ).images[0].save(f"order_{refiner.scheduler.order}_{steps}_steps.png")

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@yiyixuxu @sayakpaul

[Beinsezii]

Based on this comment, it was previously hardcoded specifically for Heun's method, and anything else is 100% broken. Thing is, Heun appears to be the only higher order singlestep solver in Diffusers, so I guess we can't add tests for this yet?

[Beinsezii]

technically this might change the current results with Heun but it's necessary because otherwise it'll split wrong on a butcher tableaux with non-sequential coefficients like

                          RKZ.Butcher6
+0.0    | 
+0.2764 | +0.2764
+0.7236 | -0.2236 +0.9472
+0.2764 | +0.0326 +0.309  -0.0652
+0.7236 | +0.0461 +0.0    +0.1667 +0.5109
+0.2764 | +0.1206 +0.0    -0.1817 +0.1667 +0.1708
+1.0    | +0.1667 +0.0    +0.0751 -3.3877 +0.5279 +3.618 
-----------------------------------------------------------------
        | +0.0833 +0.0    +0.0    +0.0    +0.4167 +0.4167 +0.0833

Where it could split on stage 3, but the following stages contain lesser timestep values, and since the refiner is not trained on earlier timesteps this will lead to worse results.