TensorStack-AI / OnnxStack

C# Stable Diffusion using ONNX Runtime
Apache License 2.0
221 stars 33 forks source link

Feature Suggestion: Support for Custom Scheduler Implementation #114

Closed ianLeslie2 closed 9 months ago

ianLeslie2 commented 9 months ago

Hi, Firstly, nice project! It's great being able to do AI operations from a C# context.

For a project I'm working on I'd like to have low-level control over how the image is generated. From what I understand about how StableDiffusion works, writing my own Scheduler is probably the right approach. From reading the source code, I see that the Scheduler classes are fairly well abstracted and have a shared IScheduler interface. I'd like to write my own IScheduler implementation and have the system use it.

However it doesn't look like there is a way to insert a custom Scheduler into the system. The scheduler implementation is chosen based on SchedulerType enum with a set of hard-coded values. To resolve this I am currently forking the repo and modifying the source-code directly, but it would be better if I didn't need to fork & maintain my own version of the code.

Could you add a way to inject custom Scheduler implementations? If you like I could make a pull request for this, but I wanted to check whether you are open to the idea first.

saddam213 commented 9 months ago

Yeah, I think I could rearrange the API a bit to allow custom schedulers.

Moving GetScheduler from DiffuserBase to PipelineBase then we can pass down the scheduler to the diffuser Then add some methods to set a custom scheduler on the pipeline similar to how we can customize the Vae models with OverrideVaeEncoder

This should be fairly easy and ill try get it done tonight after work :)

saddam213 commented 9 months ago

Well I thought this would be simple but turns out will be a significant change

It may be simpler to derive a new pipeline/diffuser or create a fork of OnnxStack with your new model/scheduler implementation

ianLeslie2 commented 9 months ago

alright, that's ok by me. As I worked on my project I already found that using a custom scheduler was not enough and that I needed to create custom variants of the other stuff as well.