Update app.py

app.py

@@ -318,9 +318,32 @@ print("✓ Text encoders loaded")
 
 # VAE (local weights - Apache 2.0 from Flux)
 print("Loading VAE...")
-vae = AutoencoderKL.from_single_file("ae.safetensors", torch_dtype=DTYPE)
+
+# Flux VAE config (16 latent channels, not standard SD 8)
+vae_config = {
+    "in_channels": 3,
+    "out_channels": 3,
+    "latent_channels": 16,
+    "block_out_channels": [128, 256, 512, 512],
+    "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"],
+    "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
+    "layers_per_block": 2,
+    "norm_num_groups": 32,
+    "act_fn": "silu",
+    "scaling_factor": 0.3611,
+    "shift_factor": 0.1159,
+    "use_quant_conv": False,
+    "use_post_quant_conv": False,
+    "mid_block_add_attention": True,
+}
+
+vae = AutoencoderKL.from_single_file(
+    "ae.safetensors",
+    config=vae_config,
+    torch_dtype=DTYPE,
+)
 vae.eval()
-VAE_SCALE = 0.3611  # Flux VAE scaling factor
+VAE_SCALE = 0.3611
 print("✓ VAE loaded")