madebyollin / taesd

Tiny AutoEncoder for Stable Diffusion
MIT License
495 stars 27 forks source link

About the Color augmentation details #21

Closed stardusts-hj closed 3 weeks ago

stardusts-hj commented 3 weeks ago

Hi madebyollin! Thanks for sharing this wonderful work! I find taesd really useful and I'm trying to finetune the pretrained taesd. I also observed the shift of color in reconstructed images without color augmentation. So I would like to ask you for the details of the color augmentation strategy used in taesd training. 20240626-114132

Thanks!

madebyollin commented 3 weeks ago

Nothing particularly exciting - this is the code I used (assuming [0, 1] images, applied to 5-10% of samples)

def color_augment(im):
    scale = 0.5 + th.randn(3, 3)
    blend = th.rand(3, 1, 1)
    return (scale @ im.flatten(1)).view(im.shape).clamp(0, 1) * blend + (1 - blend) * im
stardusts-hj commented 3 weeks ago

Thanks for your kind reply!