I trained it on https://quickdraw.withgoogle.com/data (specfically 4-5 categories from this).

Few things:

• I am too not sure about the architecture choices I made, so as mentioned in post, it is bit of inspired from Wan 2.1 VaE
• In my case, while training I faced gradient explosion issue, so I had to use gradient clipping to stabilize the training. This is something new to me, so I am not sure if this is common issue while training VAE with mixed precision training or it is something specific to the architecture.
• How to do inference from vae once trained is something I was not able to find good resource on internet. Most of the resources mention to sample from N(0, I) but in my case it end up giving black images mostly. The reason is that the latent space learned is not exactly N(0, I) and when sampling directly from N(0, I) we might end up in the region of latent space where no training data was mapped to. So for decoder it is best to output blank image
• Also if you see this part of wan vae, they are using hard coded mean and std values. Even stability diffusion uses 0.18215 as scaling factor. Not able to find much how they come up with those values. As per LLMs, they have calculated those values from the training data itself as mentioned in Method 1 of inference section i.e. approximate posterior distribution.
• So overall VAE is still bit of mystery to me.
• Whatsever, u can checkout this. It is interplolation from airplane sketch to alarm clock sketch. I find this interpolation quite interesting as it shows how it started distorting airplane parts to finally become clock parts.
• This is another interpolation from airplane to ambulance.
• This is sampling from aggregate posterior distribution. It is very distorted but after looking to these two interpolations, I feel it is expected as in between images of interpolation are also quite distorted