Check a single directory and return its model status, or None if it's not a model dir.
(dir: &Path)
| 153 | |
| 154 | /// Check a single directory and return its model status, or None if it's not a model dir. |
| 155 | fn check_model_dir(dir: &Path) -> Option<(ModelStatus, u64)> { |
| 156 | // Must have config.json to be considered a model |
| 157 | if !dir.join("config.json").exists() { |
| 158 | return None; |
| 159 | } |
| 160 | |
| 161 | let mut total_size: u64 = 0; |
| 162 | |
| 163 | // Count config + tokenizer sizes |
| 164 | for name in &["config.json", "tokenizer.json"] { |
| 165 | let p = dir.join(name); |
| 166 | if p.exists() { |
| 167 | total_size += std::fs::metadata(&p).map(|m| m.len()).unwrap_or(0); |
| 168 | } |
| 169 | } |
| 170 | |
| 171 | let index_path = dir.join("model.safetensors.index.json"); |
| 172 | if index_path.exists() { |
| 173 | total_size += std::fs::metadata(&index_path) |
| 174 | .map(|m| m.len()) |
| 175 | .unwrap_or(0); |
| 176 | |
| 177 | // Sharded model — check which shards are present |
| 178 | let index_data = std::fs::read_to_string(&index_path).ok()?; |
| 179 | let index_json: serde_json::Value = serde_json::from_str(&index_data).ok()?; |
| 180 | let weight_map = index_json.get("weight_map")?.as_object()?; |
| 181 | |
| 182 | let mut expected_shards: HashSet<String> = HashSet::new(); |
| 183 | for value in weight_map.values() { |
| 184 | if let Some(file) = value.as_str() { |
| 185 | expected_shards.insert(file.to_string()); |
| 186 | } |
| 187 | } |
| 188 | |
| 189 | let total = expected_shards.len(); |
| 190 | let mut have = 0; |
| 191 | for shard in &expected_shards { |
| 192 | let shard_path = dir.join(shard); |
| 193 | if shard_path.exists() { |
| 194 | have += 1; |
| 195 | total_size += std::fs::metadata(&shard_path) |
| 196 | .map(|m| m.len()) |
| 197 | .unwrap_or(0); |
| 198 | } |
| 199 | } |
| 200 | |
| 201 | let status = if have == total { |
| 202 | ModelStatus::Complete |
| 203 | } else { |
| 204 | ModelStatus::Partial { have, total } |
| 205 | }; |
| 206 | |
| 207 | Some((status, total_size)) |
| 208 | } else { |
| 209 | // Single safetensors file |
| 210 | let single = dir.join("model.safetensors"); |
| 211 | if single.exists() { |
| 212 | total_size += std::fs::metadata(&single) |