Useful is a funny word, but it would certainly be interesting. I think the pymc-experimental package exists for exactly projectsl like this one. PyMC would allow you to easily define the energy function (the unnormalized logp of the model), then you’ll have to do a lot of work to implement the diffusion scheme in pytensor.
Tag me if you end up working on something, it looks super interesting!