Are these schedulers supported for Flux-Schnell?

[nitinmukesh]

So I have tested 4 schedulers (ignore euler, it works) with Dev, all 4 works. However when I use these with Schnell and Shuttle-Jaguar, it throws error. Are these supposed to work only with Dev?
Ignore euler as it works.

import torch
import numpy as np
from diffusers import FluxPipeline
from diffusers import FlowMatchEulerDiscreteScheduler, DPMSolverMultistepScheduler, UniPCMultistepScheduler, DPMSolverSinglestepScheduler, DEISMultistepScheduler
from nunchaku import NunchakuFluxTransformer2dModel, NunchakuT5EncoderModel
from nunchaku.utils import get_precision

precision = get_precision()
transformer = NunchakuFluxTransformer2dModel.from_pretrained(
    f"mit-han-lab/nunchaku-flux.1-schnell/svdq-{precision}_r32-flux.1-schnell.safetensors", 
	offload=True
)
text_encoder_2 = NunchakuT5EncoderModel.from_pretrained("mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors")
pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    text_encoder_2=text_encoder_2,
    transformer=transformer,
    torch_dtype=torch.bfloat16,
)
pipeline.enable_sequential_cpu_offload()

schedulers = {
    'dpm': lambda: DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True, prediction_type="flow_prediction"),
    'unipc': lambda: UniPCMultistepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True, prediction_type="flow_prediction"),
    'sdpm': lambda: DPMSolverSinglestepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True, prediction_type="flow_prediction"),
    'deis': lambda: DEISMultistepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True, prediction_type="flow_prediction"),
    # 'euler': lambda: FlowMatchEulerDiscreteScheduler.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="scheduler", revision="refs/pr/3"),
    'euler': lambda: FlowMatchEulerDiscreteScheduler.from_pretrained("black-forest-labs/FLUX.1-schnell", subfolder="scheduler"),
}

num_steps = 50
guidance = 3.5
MAX_SEED = np.iinfo(np.int32).max
def random_seed():
    return torch.randint(0, MAX_SEED, (1,)).item()
generator = torch.Generator(device="cpu").manual_seed(random_seed())


prompts = [
    ("realistic", "A high-resolution photograph of a golden retriever playing in a grassy field at sunset"),
    ("futuristic", "A futuristic city skyline at sunset, glowing neon lights, flying cars, and towering skyscrapers with reflections on water"),
    ("fantasy", "A majestic dragon flying over a magical castle in the clouds, vibrant and colorful fantasy style"),
    ("surreal", "A clock melting over a desert cactus under a starry sky, surrealism style, inspired by Salvador Dali"),
    ("portrait", "A close-up portrait of a young woman with freckles, natural lighting, ultra-detailed, studio photography")
]

for theme, prompt in prompts:
    for name, scheduler_fn in schedulers.items():
        print(f"Running inference with {name} scheduler on {theme} prompt...")
        pipeline.scheduler = scheduler_fn()
        image = pipeline(
            prompt, 
            num_inference_steps=num_steps, 
            guidance_scale=guidance,
            generator=generator
        ).images[0]
        image_path = f"output_flux_scheduler/flux.1-dev-{theme}-{name}.png"
        image.save(image_path)
        print(f"Saved output: {image_path}")

Traceback (most recent call last):
  File "C:\aiOWN\diffuser_webui\FLUX-schedulers.py", line 51, in <module>
    image = pipeline(
            ^^^^^^^^^
  File "C:\Users\nitin\miniconda3\envs\sddw-dev\Lib\site-packages\torch\utils\_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\nitin\miniconda3\envs\sddw-dev\Lib\site-packages\diffusers\pipelines\flux\pipeline_flux.py", line 853, in __call__
    timesteps, num_inference_steps = retrieve_timesteps(
                                     ^^^^^^^^^^^^^^^^^^^
  File "C:\Users\nitin\miniconda3\envs\sddw-dev\Lib\site-packages\diffusers\pipelines\flux\pipeline_flux.py", line 141, in retrieve_timesteps
    scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
  File "C:\Users\nitin\miniconda3\envs\sddw-dev\Lib\site-packages\diffusers\schedulers\scheduling_deis_multistep.py", line 250, in set_timesteps
    assert self.config.use_dynamic_shifting and self.config.time_shift_type == "exponential"

Schnell on left and Dev on Right

I also tried with use_dynamic_shifting=False

schedulers = {
    'dpm': lambda: DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True, use_dynamic_shifting=False, prediction_type="flow_prediction"),
    'unipc': lambda: UniPCMultistepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True,  use_dynamic_shifting=False, prediction_type="flow_prediction"),
    'sdpm': lambda: DPMSolverSinglestepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True,  use_dynamic_shifting=False, prediction_type="flow_prediction"),
    'deis': lambda: DEISMultistepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True,  use_dynamic_shifting=False, prediction_type="flow_prediction"),
    # 'euler': lambda: FlowMatchEulerDiscreteScheduler.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="scheduler", revision="refs/pr/3"),
    'euler': lambda: FlowMatchEulerDiscreteScheduler.from_pretrained("black-forest-labs/FLUX.1-schnell", subfolder="scheduler"),
}