spherical-vae-flux

Spherical VAE (fine-tuned from FLUX.2-dev)

A variational autoencoder fine-tuned from black-forest-labs/FLUX.2-dev with a spherical latent space constraint. All encoder outputs are projected onto the hypersphere S^{d-1} with radius sqrt(latent_channels) = sqrt(32) ~ 5.66.

What is the spherical constraint?

Standard VAEs produce latents with unconstrained norms, leading to "holes" in latent space and unstable generation. This model replaces the standard quant_conv with a SphericalQuantConv that L2-normalizes encoder outputs and scales them to lie exactly on a hypersphere of radius sqrt(32). This ensures:

• Uniform latent coverage: all latent vectors have the same L2 norm
• No posterior collapse: the spherical constraint prevents mode collapse
• Compatible with diffusers: exported weights use standard Conv2d (spherical projection can be applied at inference time)

Training Details

Usage

from diffusers import AutoencoderKLFlux2
import torch

# Load the fine-tuned VAE
vae = AutoencoderKLFlux2.from_pretrained("tmeral/spherical-vae-flux")
vae = vae.to("cuda").eval()

# Encode an image (B, 3, H, W) -> latent
image = torch.randn(1, 3, 512, 512, device="cuda")  # replace with real image
with torch.no_grad():
    latent = vae.encode(image).latent_dist.sample()

# Apply spherical projection (optional, for strict sphere constraint)
import math
radius = math.sqrt(32)
latent_flat = latent.flatten(2)  # (B, C, H*W)
norms = latent_flat.norm(dim=1, keepdim=True)  # (B, 1, H*W)
latent_flat = latent_flat / norms * radius
latent = latent_flat.view_as(latent)

# Decode back to pixel space
with torch.no_grad():
    recon = vae.decode(latent).sample

print(f"Input: {image.shape} -> Latent: {latent.shape} -> Recon: {recon.shape}")

Qualitative Comparisons

Side-by-side comparisons of original FLUX.2-dev VAE vs this fine-tuned spherical VAE at different resolutions:

256x256

512x512

1024x1024

Files

License

This is a research checkpoint for evaluation and collaboration purposes. The base model (FLUX.2-dev) has its own license terms.