Update app.py

app.py

@@ -316,34 +316,11 @@ clip_tok = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
 clip_enc = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=DTYPE)
 print("✓ Text encoders loaded")
 
-# VAE (local weights - Apache 2.0 from Flux)
+# VAE (local diffusers format)
 print("Loading VAE...")
-
-# 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 = AutoencoderKL.from_pretrained("./vae", torch_dtype=DTYPE)
 vae.eval()
-VAE_SCALE = 0.3611
+VAE_SCALE = vae.config.scaling_factor
 print("✓ VAE loaded")