Issue with subclassing a scheduler

[christopher-beckham]

Hi,

How exactly are we meant to subclass a particular scheduler if wanted to do it? Consider the following code block:

from diffusers import AutoPipelineForText2Image
import torch
from diffusers import EulerAncestralDiscreteScheduler
from diffusers.configuration_utils import register_to_config
from dataclasses import dataclass

model_id = "stabilityai/sdxl-turbo"

pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
pipe.to("cuda")

prompt = "A cinematic shot of a baby racoon wearing an intricate italian priest robe."

sampler = EulerAncestralDiscreteScheduler.from_pretrained(
    model_id, subfolder="scheduler"
)
pipe.scheduler = sampler

image = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
image.save("image1.png")

Now maybe I want to add my own scheduler which subclasses another. Let me do the simplest example which is just calling the constructor of the parent class:

@dataclass
class EulerAncestralDiscreteSchedulerChris(EulerAncestralDiscreteScheduler):
    #@register_to_config
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

sampler2 = EulerAncestralDiscreteSchedulerChris.from_pretrained(
    model_id, subfolder="scheduler"
)
pipe.scheduler = sampler2

image = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
image.save("image2.png")

image1 and image2 are the following, respectively, so something is going on:

If you print sampler2, you just get a string like this EulerAncestralDiscreteSchedulerChris(), so for some reason it's not printing the attributes inside the object. If you omit @dataclass from the subclass then you get the error *** TypeError: Object of type type is not JSON serializable.

Is this a bug? Or am I not subclassing the scheduler correctly?

Thanks.