Stable Diffusion XL in Rust

I finally got around to porting SDXL to Rust's deep learning framework burn. This time I didn't have an elegant tinygrad reference so I had to go digging through Stability AI's horrendous python repo. Forget spaghetti code, that was a jungle safari. Their diffusion unet model in particular was one of the most unnecessarily convoluted piece of code I've ever seen. It is my hope that people find my implementation cleaner and more comprehensible.

In terms of image quality, SDXL is a huge step up from SD. The square resolution has been increased from 512x512 to 1024x1024. There was a training bug with SD that causes image generations to be cropped. SDXL doesn't have this issue since the crop parameters can be directly specified during generation. Here is an example SDXL output:

Beautiful dwelling on a rocky coast.

Quite beautiful, I'd say. I feel that SDXL outputs are much closer to production quality than those of SD. I haven't yet implemented the refiner model which improves the small details, but SDXL is very good even without it.

Check out the project at https://github.com/Gadersd/stable-diffusion-xl-burn.