Update app.py

app.py

@@ -318,32 +318,9 @@ print("✓ Text encoders loaded")
 
 # VAE (local weights - Apache 2.0 from Flux)
 print("Loading VAE...")
-from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
-
-# VAE config from Flux
-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",
-    "sample_size": 1024,
-    "scaling_factor": 0.3611,
-    "shift_factor": 0.1159,
-    "use_quant_conv": False,
-    "use_post_quant_conv": False,
-    "mid_block_add_attention": True,
-}
-
-vae = AutoencoderKL(**vae_config)
-vae_weights = load_file("ae.safetensors")
-vae.load_state_dict(vae_weights)
-vae.to(DTYPE).eval()
-VAE_SCALE = vae_config["scaling_factor"]
+vae = AutoencoderKL.from_single_file("ae.safetensors", torch_dtype=DTYPE)
+vae.eval()
+VAE_SCALE = 0.3611  # Flux VAE scaling factor
 print("✓ VAE loaded")