Using diffusers for research and application to physics

[aurelio-amerio]

Hello,

I'm a physics researcher exploring the use of diffusion models for my work. My current focus is on simulation-based inference, where we train models on simulations of physical phenomena and condition them on the parameters describing those phenomena. This allows us to reconstruct key parameters from complex observations. Diffusion models are proving incredibly valuable in this area.

I'm very interested in building a new JAX library with your diffusers library as a core component. I've seen other physics-oriented libraries emerging, but they often seem to re-implement existing solutions, particularly for SDE solving and sampling, which diffusers already handles expertly.

While physics researchers typically train models on their specific data, your library's schedulers and embedders offer excellent potential for reuse. I have, however, noticed some discrepancies between the PyTorch and JAX implementations (for example, the VESDE scheduler in JAX seems to be missing the add_noise method). I'm keen to contribute and help fill these gaps as I develop my library.

Could you provide any guidance on contributing, particularly regarding scheduler/embedding implementations and ensuring consistent testing across PyTorch and JAX?

Thank you for your time.

TL;DR: I'm interested in contributing to diffusers by improving JAX support, especially for schedulers/embedders, and ensuring consistency with the PyTorch implementation. Could you provide guidance on contributing and testing?