| 312 | |
| 313 | impl FluxTextEncoder { |
| 314 | pub fn load_model( |
| 315 | device: &Device, |
| 316 | dtype: DType, |
| 317 | model_repo: &str, |
| 318 | backend: Arc<dyn ComputeBackend>, |
| 319 | ) -> anyhow::Result<Box<Self>> { |
| 320 | let cfg = EncoderConfig::flux2_klein(); |
| 321 | |
| 322 | let cache_dir = dirs::cache_dir() |
| 323 | .unwrap_or_else(std::env::temp_dir) |
| 324 | .to_string_lossy() |
| 325 | .to_string(); |
| 326 | |
| 327 | let all_paths = FluxModelFile::TextEncoder.get_all(model_repo, &cache_dir)?; |
| 328 | // Filter to only safetensor files (exclude config.json, index.json) |
| 329 | let weight_paths: Vec<_> = all_paths |
| 330 | .into_iter() |
| 331 | .filter(|p| { |
| 332 | p.extension() |
| 333 | .map(|e| e == "safetensors") |
| 334 | .unwrap_or(false) |
| 335 | }) |
| 336 | .collect(); |
| 337 | info!( |
| 338 | "loading FLUX text encoder from {} shard(s)", |
| 339 | weight_paths.len() |
| 340 | ); |
| 341 | |
| 342 | let vb = unsafe { |
| 343 | VarBuilder::from_mmaped_safetensors(&weight_paths, dtype, device)? |
| 344 | }; |
| 345 | |
| 346 | // Weights are prefixed with "model." in the text_encoder safetensors |
| 347 | let vb_model = vb.pp("model"); |
| 348 | |
| 349 | let embeddings_weight = vb_model.pp("embed_tokens").get((cfg.vocab_size, cfg.hidden_size), "weight")?; |
| 350 | |
| 351 | let mut blocks = Vec::with_capacity(cfg.num_hidden_layers); |
| 352 | let vb_layers = vb_model.pp("layers"); |
| 353 | for i in 0..cfg.num_hidden_layers { |
| 354 | let block = EncoderBlock::load(vb_layers.pp(i), &cfg, backend.clone())?; |
| 355 | blocks.push(block); |
| 356 | } |
| 357 | |
| 358 | let final_norm_weight = vb_model.pp("norm").get(cfg.hidden_size, "weight")?; |
| 359 | let final_norm_eps = cfg.rms_norm_eps as f32; |
| 360 | |
| 361 | info!("FLUX text encoder loaded ({} layers)", cfg.num_hidden_layers); |
| 362 | |
| 363 | Ok(Box::new(Self { |
| 364 | embeddings_weight, |
| 365 | blocks, |
| 366 | final_norm_weight, |
| 367 | final_norm_eps, |
| 368 | cfg, |
| 369 | backend, |
| 370 | })) |
| 371 | } |