| 111 | |
| 112 | impl FluxVAE { |
| 113 | pub fn load_model( |
| 114 | device: &Device, |
| 115 | _dtype: DType, |
| 116 | model_repo: &str, |
| 117 | backend: Arc<dyn ComputeBackend>, |
| 118 | ) -> anyhow::Result<Box<Self>> { |
| 119 | // VAE always runs in F32 for numerical stability |
| 120 | let dtype = DType::F32; |
| 121 | |
| 122 | let cache_dir = dirs::cache_dir() |
| 123 | .unwrap_or_else(std::env::temp_dir) |
| 124 | .to_string_lossy() |
| 125 | .to_string(); |
| 126 | |
| 127 | let weights_path = FluxModelFile::Vae.get(model_repo, &cache_dir)?; |
| 128 | info!("loading FLUX VAE from {}", weights_path.display()); |
| 129 | |
| 130 | let vb = unsafe { |
| 131 | VarBuilder::from_mmaped_safetensors(&[weights_path], dtype, device)? |
| 132 | }; |
| 133 | |
| 134 | // Load batch norm running statistics |
| 135 | let bn_running_mean = vb.get(128, "bn.running_mean")?.to_dtype(DType::F32)?; |
| 136 | let bn_running_var = vb.get(128, "bn.running_var")?.to_dtype(DType::F32)?; |
| 137 | |
| 138 | let model = Flux2VAE::load(vb, backend)?; |
| 139 | info!("FLUX VAE loaded (custom decoder)"); |
| 140 | |
| 141 | Ok(Box::new(Self { |
| 142 | model, |
| 143 | bn_running_mean, |
| 144 | bn_running_var, |
| 145 | })) |
| 146 | } |
| 147 | |
| 148 | /// Decode latents to image (legacy interface, used by SD-style callers). |
| 149 | #[allow(dead_code)] |