stage2 vae mask