Spherical VAE Flux
research-only
Modellquelle
- Art
- VAE
- Dateigröße
- 321 MB
- Version
- vdiffusion_pytorch_model
- Aktualisiert
- 9. Okt. 2026
- Downloads
- 17
- Quelle
- Hugging Face
Quellenbeschreibung
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
| Parameter | Value |
|---|---|
| Base model | FLUX.2-dev VAE (AutoencoderKLFlux2) |
| Training steps | 10,000 |
| Dataset | FLUX-Reason-6M (1024x1024) |
| Weights | EMA (decay=0.9999) |
| Sphere radius | sqrt(32) ~ 5.66 |
| Latent channels | 32 |
| Decoder LR | 4e-5 |
| Encoder LR | 4e-6 |
| Loss | L1 + LPIPS + KL + adversarial (after 10k steps) |
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
Comparison at 256x256
512x512
Comparison at 512x512
1024x1024
Comparison at 1024x1024
Files
| File | Description |
|---|---|
diffusion_pytorch_model.safetensors | Model weights in safetensors format |
config.json | Diffusers model config |
spherical_vae_metadata.json | Spherical VAE training metadata (radius, channels, base model) |
inference.py | Self-contained encode/decode inference script |
training_config.yaml | Full training configuration |
comparison_*.png | Qualitative comparison images at 256/512/1024 |
License
This is a research checkpoint for evaluation and collaboration purposes. The base model (FLUX.2-dev) has its own license terms.
Quellen
1 Quelle
Quellen
1 QuelleVerifiziert 9. Okt.
Modellartefakte
1 Artefakt
Modellartefakte
1 Artefaktdiffusion_pytorch_model.safetensors
safetensors · 321 MB · SHA-256 f1750bbc985c…d24e · Hugging Face
HerunterladenQuellenauszüge
3 Auszüge
Quellenauszüge
3 Auszüge--- license: other license_name: research-only tags: - vae - autoencoder - flux - spherical-vae - image-generation - latent-space library_name: diffusers pipeline_tag: image-to-image --- # Spherical VAE (fine-tuned from FLUX.2-dev) A variational autoencoder fine-tuned from [black-forest-labs/FLUX.2-dev](https://huggingface.co/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 | Parameter | Value | |-----------|-------| | Base model | FLUX.2-dev VAE (AutoencoderKLFlux2) | | Training steps | 10,000 | | Dataset | [FLUX-Reason-6M](https://huggingface.co/datasets/LucasFang/FLUX-Reason-6M) (1024x1024) | | Weights | EMA (decay=0.9999) | | Sphere radius | sqrt(32) ~ 5.66 | | Latent channels | 32 | | Decoder LR | 4e-5 | | Encoder LR | 4e-6 | | Loss | L1 + LPIPS + KL + adversarial (after 10k steps) | ## Usage ```python 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} ->...
Source context: 17 downloads · 2 likes · Pipeline image-to-image · Library diffusers · Repo tmeral/spherical-vae-flux
Source context: 4 downloads · 2 likes · Pipeline image-to-image · Library diffusers · Repo tmeral/spherical-vae-flux
