Fix VAE: switch to madebyollin/sdxl-vae-fp16-fix (open, no auth needed, 4ch latent)

train.py

@@ -6,6 +6,7 @@ Optimized for Colab free tier:
 - No VAE needed during training loop → saves ~1GB VRAM + faster iterations
 - Streaming support for large datasets
 - Multiple small dataset presets
+- Uses madebyollin/sdxl-vae-fp16-fix (fully open, no login, fp16 stable)
 
 Flow Matching training objective (velocity prediction):
 - Forward: x_t = (1 - t) * x_0 + t * ε
@@ -27,7 +28,7 @@ from dataclasses import dataclass, asdict
 
 
 # =============================================================================
-# Dataset Presets (all verified, fast to download)
+# Dataset Presets (all verified, fast to download, no auth needed)
 # =============================================================================
 
 DATASET_PRESETS = {
@@ -88,11 +89,12 @@ class TrainConfig:
     image_size: int = 256              # 256 or 512
     max_images: int = 0                # 0 = use all, >0 = limit (for streaming/testing)
 
-    # VAE (for pre-caching only — NOT loaded during training)
-    vae_id: str = "black-forest-labs/FLUX.1-schnell"
-    vae_subfolder: str = "vae"
-    vae_scaling_factor: float = 0.3611
-    vae_shift_factor: float = 0.1159
+    # VAE — fully open, no login needed
+    # madebyollin/sdxl-vae-fp16-fix: SDXL VAE with fp16 NaN fix
+    # 4 latent channels, 8x spatial compression, scaling_factor=0.13025
+    vae_id: str = "madebyollin/sdxl-vae-fp16-fix"
+    vae_scaling_factor: float = 0.13025
+    latent_channels: int = 4
 
     # Training
     batch_size: int = 32               # Can be large since training on cached tensors!
@@ -147,7 +149,7 @@ def get_model_config(size, num_classes=0, class_drop_prob=0.1):
 
 
 # =============================================================================
-# Latent Pre-Caching (the key optimization for Colab)
+# Latent Pre-Caching
 # =============================================================================
 
 class CachedLatentDataset(Dataset):
@@ -174,13 +176,7 @@ class CachedLatentDataset(Dataset):
 def precache_latents(config, cache_path=None):
     """
     Encode all images to VAE latents once, save to disk.
-
-    After caching:
-    - VAE unloaded → frees ~1GB VRAM
-    - Training loads pure tensors → much faster iterations
-    - Larger batch sizes possible (no VAE memory overhead)
-
-    Returns path to cache file.
+    Uses madebyollin/sdxl-vae-fp16-fix (no auth needed).
     """
     if cache_path is None:
         cache_path = os.path.join(config.output_dir, "cached_latents.pt")
@@ -194,14 +190,15 @@ def precache_latents(config, cache_path=None):
     os.makedirs(os.path.dirname(cache_path) if os.path.dirname(cache_path) else ".", exist_ok=True)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # Load VAE temporarily
-    print("Loading VAE for encoding...")
+    # Load VAE — no subfolder, no auth needed
+    print(f"Loading VAE: {config.vae_id} (open, no login needed)...")
     from diffusers import AutoencoderKL
     vae = AutoencoderKL.from_pretrained(
-        config.vae_id, subfolder=config.vae_subfolder, torch_dtype=torch.float16
+        config.vae_id, torch_dtype=torch.float16
     ).to(device).eval()
     for p in vae.parameters():
         p.requires_grad_(False)
+    print(f"  VAE loaded: {sum(p.numel() for p in vae.parameters())/1e6:.0f}M params")
 
     # Load dataset
     preset = DATASET_PRESETS[config.dataset_preset]
@@ -235,7 +232,7 @@ def precache_latents(config, cache_path=None):
     img_col = preset["image_column"]
     lbl_col = preset["label_column"]
 
-    print(f"Encoding images to latents...")
+    print(f"Encoding images to VAE latents...")
     t0 = time.time()
 
     for item in dataset:
@@ -255,7 +252,7 @@ def precache_latents(config, cache_path=None):
             with torch.no_grad():
                 px = torch.stack(batch_pixels).to(device, dtype=torch.float16) * 2 - 1
                 lat = vae.encode(px).latent_dist.sample()
-                lat = (lat - config.vae_shift_factor) * config.vae_scaling_factor
+                lat = lat * config.vae_scaling_factor  # SDXL: scale only, no shift
                 all_latents.append(lat.cpu().float())
             all_labels.extend(batch_labels)
             batch_pixels, batch_labels = [], []
@@ -266,7 +263,7 @@ def precache_latents(config, cache_path=None):
         with torch.no_grad():
             px = torch.stack(batch_pixels).to(device, dtype=torch.float16) * 2 - 1
             lat = vae.encode(px).latent_dist.sample()
-            lat = (lat - config.vae_shift_factor) * config.vae_scaling_factor
+            lat = lat * config.vae_scaling_factor
             all_latents.append(lat.cpu().float())
         all_labels.extend(batch_labels)
 
@@ -379,8 +376,9 @@ def train(config):
     train_dl = DataLoader(train_ds, batch_size=config.batch_size, shuffle=True,
                           num_workers=config.num_workers, pin_memory=True, drop_last=True)
 
-    # Step 3: Model
+    # Step 3: Model (in_channels=4 for SDXL VAE)
     mcfg = get_model_config(config.model_size, config.num_classes, config.class_drop_prob)
+    mcfg["in_channels"] = config.latent_channels  # 4 for SDXL VAE
     model = LiquidGen(**mcfg).to(device)
     print(f"LiquidGen-{config.model_size}: {model.count_params()/1e6:.1f}M params")
 
@@ -397,8 +395,9 @@ def train(config):
     fm = FlowMatchingScheduler(config.min_timestep, config.max_timestep)
     lat_size = config.image_size // 8
 
-    print(f"\nTotal steps: {total_steps}, Batch: {config.batch_size}×{config.gradient_accumulation_steps}")
-    print(f"No VAE during training → max VRAM for model")
+    print(f"\nTotal steps: {total_steps}, Batch: {config.batch_size}x{config.gradient_accumulation_steps}")
+    print(f"Latent: [{config.batch_size}, {config.latent_channels}, {lat_size}, {lat_size}]")
+    print(f"No VAE during training -> max VRAM for model")
     if torch.cuda.is_available():
         print(f"VRAM: {torch.cuda.memory_allocated()/1024**3:.1f} / "
               f"{torch.cuda.get_device_properties(0).total_mem/1024**3:.1f} GB")
@@ -444,20 +443,23 @@ def train(config):
                         print("💥 Diverged!"); return
 
                 if gs % config.sample_every_n_steps == 0:
+                    # Load VAE lazily (only for decoding samples)
                     if not vae_loaded:
                         from diffusers import AutoencoderKL
                         vae = AutoencoderKL.from_pretrained(
-                            config.vae_id, subfolder=config.vae_subfolder,
-                            torch_dtype=torch.float16).to(device).eval()
+                            config.vae_id, torch_dtype=torch.float16
+                        ).to(device).eval()
                         for p in vae.parameters(): p.requires_grad_(False)
                         vae_loaded = True
                     ema.apply(model); model.eval()
                     sl = torch.randint(0, max(1, config.num_classes), (config.num_samples,),
                                        device=device) if config.num_classes > 0 else None
-                    samp = fm.sample(model, (config.num_samples, 16, lat_size, lat_size),
+                    # 4 channels for SDXL VAE
+                    samp = fm.sample(model, (config.num_samples, config.latent_channels, lat_size, lat_size),
                                      device, config.num_sample_steps, sl, config.cfg_scale)
                     with torch.no_grad():
-                        dec = samp.half() / config.vae_scaling_factor + config.vae_shift_factor
+                        # SDXL VAE: unscale only, no shift
+                        dec = samp.half() / config.vae_scaling_factor
                         imgs = ((vae.decode(dec).sample + 1) / 2).clamp(0, 1).float()
                     from torchvision.utils import save_image
                     sp = f"{config.output_dir}/samples/step_{gs:07d}.png"
@@ -476,7 +478,7 @@ def train(config):
     final = f"{config.output_dir}/checkpoints/final.pt"
     torch.save({"model": model.state_dict(), "ema": ema.shadow,
                 "model_config": mcfg, "step": gs}, final)
-    print(f"\n🎉 Done! {gs} steps, {(time.time()-t_start)/60:.1f}min → {final}")
+    print(f"\n🎉 Done! {gs} steps, {(time.time()-t_start)/60:.1f}min -> {final}")
 
 
 if __name__ == "__main__":