Update app.py

app.py

@@ -316,10 +316,34 @@ 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
+# VAE (local weights - Apache 2.0 from Flux)
 print("Loading VAE...")
-vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-schnell", subfolder="vae", torch_dtype=DTYPE)
-VAE_SCALE = vae.config.scaling_factor
+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"]
 print("✓ VAE loaded")