You can deploy this model and start earning money today!

gabehubner/
vae-256px-8z

Name: vae-256px-8z
Rating: 2.5 (8 reviews)

Views

No views yet

UNet-Style VAE for 256x256 Image Reconstruction

This model is a UNet-style Variational Autoencoder (VAE) trained on the CC3M dataset for high-quality image reconstruction and generation. It integrates adversarial, perceptual, and identity-preserving loss terms to improve semantic and visual fidelity.

Architecture

Encoder/Decoder: Multi-scale UNet architecture
Latent Space: 8-channel latent bottleneck with reparameterization (mu, logvar)
Losses:
- L1 reconstruction loss
- KL divergence with annealing
- LPIPS perceptual loss (VGG backbone)
- Identity loss via MoCo-v2 embeddings
- Adversarial loss via Patch Discriminator w/ Spectral Norm

$$ \mathcal{L}{total} = \mathcal{L}{recon} + \mathcal{L}{PIPS} + 0.5 * \mathcal{L}{GAN} + 0.1 *\mathcal{L}{ID} + 10^{-6} *\mathcal{L}{KL} $$

Reconstructions

Input	Output
input	output

Training Config

Hyperparameter	Value
Dataset	CC3M (850k images)
Image Resolution	256 x 256
Batch Size	16
Optimizer	AdamW
Learning Rate	5e-5
Precision	bf16 (mixed precision)
Total Steps	210,000
GAN Start Step	50,000
KL Annealing	Yes (10% of training)
Augmentations	Crop, flip, jitter, blur, rotation

Trained using a cosine learning rate schedule with gradient clipping and automatic mixed precision (torch.cuda.amp)

gabehubner/vae-256px-8z

UNet-Style VAE for 256x256 Image Reconstruction

Architecture

Reconstructions

Training Config

gabehubner/
vae-256px-8z