(
config_path: &std::path::Path,
weights_path: &std::path::Path,
device: &Device,
diffusion_steps: Option<usize>,
topology: &crate::cake::Topology,
clus
| 62 | |
| 63 | impl VibeVoiceTTS { |
| 64 | pub async fn load( |
| 65 | config_path: &std::path::Path, |
| 66 | weights_path: &std::path::Path, |
| 67 | device: &Device, |
| 68 | diffusion_steps: Option<usize>, |
| 69 | topology: &crate::cake::Topology, |
| 70 | cluster_key: Option<&str>, |
| 71 | ) -> Result<Self> { |
| 72 | let config = VibeVoiceConfig::from_path(config_path)?; |
| 73 | let common_cfg = config.into_config(); |
| 74 | |
| 75 | info!("Loading VibeVoice-Realtime-0.5B..."); |
| 76 | // Use F32 for exact numerical match with reference implementation. |
| 77 | // This disables Flash Attention but ensures correct conditioning. |
| 78 | let dtype = DType::F32; |
| 79 | |
| 80 | let vb = unsafe { |
| 81 | VarBuilder::from_mmaped_safetensors(&[weights_path.to_path_buf()], dtype, device)? |
| 82 | }; |
| 83 | |
| 84 | // Base LM (4 layers) |
| 85 | info!(" Loading base LM (4 layers)..."); |
| 86 | let base_vb = vb.pp("model").pp("language_model"); |
| 87 | let base_embed_weight = base_vb.pp("embed_tokens").get((common_cfg.vocab_size, common_cfg.hidden_size), "weight")?; |
| 88 | // Base LM norm: not in checkpoint (initialized to ones, matching Qwen2Model default) |
| 89 | let base_norm_weight = Tensor::ones(common_cfg.hidden_size, dtype, device)?; |
| 90 | let base_norm_eps = common_cfg.rms_norm_eps as f32; |
| 91 | let mut base_layers: Vec<Box<dyn crate::cake::Forwarder>> = Vec::new(); |
| 92 | for i in 0..4 { |
| 93 | let layer_name = format!("model.language_model.layers.{i}"); |
| 94 | if let Some((_node_name, node)) = topology.get_node_for_layer(&layer_name) { |
| 95 | info!(" {layer_name} → remote worker at {}", node.host); |
| 96 | base_layers.push(Box::new(crate::cake::Client::new( |
| 97 | device.clone(), &node.host, &layer_name, cluster_key, |
| 98 | ).await?)); |
| 99 | } else { |
| 100 | base_layers.push(Box::new( |
| 101 | crate::models::common::Transformer::load_for_vibevoice(base_vb.pp("layers").pp(i), &common_cfg, crate::backends::create_backend(device))?, |
| 102 | )); |
| 103 | } |
| 104 | } |
| 105 | |
| 106 | // TTS LM (20 layers) |
| 107 | let num_tts = config.tts_backbone_num_hidden_layers; |
| 108 | info!(" Loading TTS LM ({num_tts} layers)..."); |
| 109 | let tts_vb = vb.pp("model").pp("tts_language_model"); |
| 110 | let tts_embed_weight = tts_vb.pp("embed_tokens").get((common_cfg.vocab_size, common_cfg.hidden_size), "weight")?; |
| 111 | let tts_norm_weight = tts_vb.pp("norm").get(common_cfg.hidden_size, "weight")?; |
| 112 | let tts_norm_eps = common_cfg.rms_norm_eps as f32; |
| 113 | let mut tts_layers: Vec<Box<dyn crate::cake::Forwarder>> = Vec::new(); |
| 114 | for i in 0..num_tts { |
| 115 | let layer_name = format!("model.tts_language_model.layers.{i}"); |
| 116 | if let Some((_node_name, node)) = topology.get_node_for_layer(&layer_name) { |
| 117 | info!(" {layer_name} → remote worker at {}", node.host); |
| 118 | tts_layers.push(Box::new(crate::cake::Client::new( |
| 119 | device.clone(), &node.host, &layer_name, cluster_key, |
| 120 | ).await?)); |
| 121 | } else { |
nothing calls this directly
no test coverage detected