main
rs 659 lines 19.5 KB
Raw
1 //! Upload fine-tuned GGUF models to a HuggingFace repository.
2 //!
3 //! Uses the HuggingFace Hub LFS upload flow:
4 //! 1. `POST /api/repos/create` — create the repo (409 = already exists, OK)
5 //! 2. `POST /api/models/{repo}/preupload/main` — check upload mode
6 //! 3. Compute SHA-256 of the local file
7 //! 4. `POST /{repo}.git/info/lfs/objects/batch` — get the S3 upload URL
8 //! 5. `PUT {s3_url}` — stream the file to S3
9 //! 6. `POST /{repo}.git/info/lfs/objects/verify` — confirm the upload
10 //! 7. `POST /api/models/{repo}/commit/main` — create the commit
11 //!
12 //! Progress is reported through a `tokio::sync::mpsc::UnboundedSender` so the
13 //! TUI can drive a progress bar without blocking.
14
15 use std::path::PathBuf;
16
17 use anyhow::{Context, Result};
18 use sha2::{Digest, Sha256};
19
20 // ---------------------------------------------------------------------------
21 // Public types
22 // ---------------------------------------------------------------------------
23
24 /// Everything needed to upload a single file to HuggingFace.
25 pub struct UploadConfig {
26 /// Path to the local `.gguf` file.
27 pub file_path: PathBuf,
28 /// Full repo identifier, e.g. `ondeinference/qwen2.5-0.5b-finetuned`.
29 pub repo_id: String,
30 /// Destination filename inside the repo, e.g. `model-finetuned-q8_0.gguf`.
31 pub path_in_repo: String,
32 /// HuggingFace Bearer token.
33 pub hf_token: String,
34 /// Commit message shown on HuggingFace.
35 pub commit_message: String,
36 }
37
38 /// Progress updates emitted during the upload.
39 pub enum UploadProgress {
40 /// About to create (or confirm) the HuggingFace repo.
41 CreatingRepo,
42 /// Computing SHA-256 hash of the local file.
43 Hashing { bytes_done: u64, bytes_total: u64 },
44 /// Streaming file bytes to the server.
45 Uploading { bytes_sent: u64, bytes_total: u64 },
46 /// Creating the commit on HuggingFace.
47 Committing,
48 /// Upload finished successfully.
49 Done { url: String },
50 /// Something went wrong.
51 Failed(String),
52 }
53
54 // ---------------------------------------------------------------------------
55 // Public entry point
56 // ---------------------------------------------------------------------------
57
58 /// Spawn-friendly wrapper that never returns an error — failures are sent
59 /// through `tx` as [`UploadProgress::Failed`] instead.
60 pub async fn start_upload(
61 config: UploadConfig,
62 tx: tokio::sync::mpsc::UnboundedSender<UploadProgress>,
63 ) {
64 if let Err(error) = run_upload(&config, &tx).await {
65 eprintln!("[upload] error: {error:#}");
66 let _ = tx.send(UploadProgress::Failed(format!("{error:#}")));
67 }
68 }
69
70 // ---------------------------------------------------------------------------
71 // Implementation
72 // ---------------------------------------------------------------------------
73
74 /// Size of each chunk streamed during upload (1 MiB).
75 const CHUNK_SIZE: usize = 1_048_576;
76
77 /// Size of each chunk read during SHA-256 hashing (4 MiB).
78 const HASH_CHUNK_SIZE: usize = 4 * 1_048_576;
79
80 async fn run_upload(
81 config: &UploadConfig,
82 tx: &tokio::sync::mpsc::UnboundedSender<UploadProgress>,
83 ) -> Result<()> {
84 let client = reqwest::Client::builder()
85 .user_agent("onde-cli/0.2")
86 .redirect(reqwest::redirect::Policy::limited(10))
87 .timeout(std::time::Duration::from_secs(7200))
88 .build()
89 .context("failed to build HTTP client")?;
90
91 let file_size = std::fs::metadata(&config.file_path)
92 .with_context(|| format!("failed to stat {}", config.file_path.display()))?
93 .len();
94
95 // -- Step 1: Create the repo (or confirm it already exists) ---------------
96
97 let _ = tx.send(UploadProgress::CreatingRepo);
98 create_repo(&client, &config.hf_token, &config.repo_id).await?;
99
100 // -- Step 2: Preupload check ----------------------------------------------
101
102 let upload_mode = preupload_check(
103 &client,
104 &config.hf_token,
105 &config.repo_id,
106 &config.path_in_repo,
107 file_size,
108 )
109 .await?;
110
111 eprintln!("[upload] upload mode: {upload_mode}");
112
113 // -- Step 3: Compute SHA-256 of the file ----------------------------------
114
115 let _ = tx.send(UploadProgress::Hashing {
116 bytes_done: 0,
117 bytes_total: file_size,
118 });
119
120 let sha256 = compute_sha256(&config.file_path, file_size, tx).await?;
121 eprintln!("[upload] sha256: {sha256}");
122
123 if upload_mode == "lfs" {
124 // -- Step 4: LFS batch — get the S3 upload URL ------------------------
125
126 let lfs_info = lfs_batch(
127 &client,
128 &config.hf_token,
129 &config.repo_id,
130 &sha256,
131 file_size,
132 )
133 .await?;
134
135 match lfs_info {
136 LfsBatchResult::AlreadyExists => {
137 eprintln!("[upload] LFS object already exists, skipping upload");
138 }
139 LfsBatchResult::NeedsUpload {
140 upload_url,
141 verify_url,
142 verify_token,
143 } => {
144 // -- Step 5: PUT to S3 ----------------------------------------
145
146 let _ = tx.send(UploadProgress::Uploading {
147 bytes_sent: 0,
148 bytes_total: file_size,
149 });
150
151 upload_to_s3(&client, &upload_url, &config.file_path, file_size, tx).await?;
152
153 // -- Step 6: LFS verify ---------------------------------------
154
155 lfs_verify(&client, &verify_url, &verify_token, &sha256, file_size).await?;
156 }
157 }
158
159 // -- Step 7: Commit with LFS pointer ----------------------------------
160
161 let _ = tx.send(UploadProgress::Committing);
162
163 lfs_commit(
164 &client,
165 &config.hf_token,
166 &config.repo_id,
167 &config.path_in_repo,
168 &config.commit_message,
169 &sha256,
170 file_size,
171 )
172 .await?;
173 } else {
174 // Regular (non-LFS) commit for small files.
175 let _ = tx.send(UploadProgress::Uploading {
176 bytes_sent: 0,
177 bytes_total: file_size,
178 });
179
180 regular_commit(
181 &client,
182 &config.hf_token,
183 &config.repo_id,
184 &config.path_in_repo,
185 &config.commit_message,
186 &config.file_path,
187 file_size,
188 tx,
189 )
190 .await?;
191 }
192
193 let model_url = format!("https://huggingface.co/{}", config.repo_id);
194 let _ = tx.send(UploadProgress::Done { url: model_url });
195
196 Ok(())
197 }
198
199 // ---------------------------------------------------------------------------
200 // Step 1: Create repo
201 // ---------------------------------------------------------------------------
202
203 async fn create_repo(client: &reqwest::Client, token: &str, repo_id: &str) -> Result<()> {
204 // The HF API expects "name" (repo only) and "organization" as separate fields.
205 let create_body = if let Some((org, name)) = repo_id.split_once('/') {
206 serde_json::json!({
207 "type": "model",
208 "name": name,
209 "organization": org,
210 "private": false,
211 "sdk": "onde-cli"
212 })
213 } else {
214 serde_json::json!({
215 "type": "model",
216 "name": repo_id,
217 "private": false,
218 "sdk": "onde-cli"
219 })
220 };
221
222 let response = client
223 .post("https://huggingface.co/api/repos/create")
224 .bearer_auth(token)
225 .json(&create_body)
226 .send()
227 .await
228 .context("failed to reach HuggingFace API")?;
229
230 let status = response.status();
231 let body = response.text().await.unwrap_or_default();
232 eprintln!("[upload] create repo: {status} {body}");
233
234 // 200 = created, 409 = already exists — both are fine.
235 if !status.is_success() && status.as_u16() != 409 {
236 anyhow::bail!("failed to create repo {repo_id}: {status} {body}");
237 }
238
239 Ok(())
240 }
241
242 // ---------------------------------------------------------------------------
243 // Step 2: Preupload check
244 // ---------------------------------------------------------------------------
245
246 async fn preupload_check(
247 client: &reqwest::Client,
248 token: &str,
249 repo_id: &str,
250 path_in_repo: &str,
251 file_size: u64,
252 ) -> Result<String> {
253 let body = serde_json::json!({
254 "files": [{
255 "path": path_in_repo,
256 "size": file_size,
257 "sample": ""
258 }]
259 });
260
261 let url = format!("https://huggingface.co/api/models/{repo_id}/preupload/main");
262
263 let response = client
264 .post(&url)
265 .bearer_auth(token)
266 .json(&body)
267 .send()
268 .await
269 .context("preupload check failed")?;
270
271 let status = response.status();
272 let text = response.text().await.unwrap_or_default();
273 eprintln!("[upload] preupload: {status} {text}");
274
275 if !status.is_success() {
276 // Default to LFS for large files if preupload endpoint fails.
277 if file_size > 10_000_000 {
278 return Ok("lfs".to_string());
279 }
280 anyhow::bail!("preupload check failed: {status} {text}");
281 }
282
283 // Parse: {"files":[{"path":"...","uploadMode":"lfs"|"regular",...}],...}
284 let parsed: serde_json::Value =
285 serde_json::from_str(&text).context("failed to parse preupload response")?;
286
287 let upload_mode = parsed
288 .get("files")
289 .and_then(|f| f.as_array())
290 .and_then(|arr| arr.first())
291 .and_then(|entry| entry.get("uploadMode"))
292 .and_then(|v| v.as_str())
293 .unwrap_or("lfs")
294 .to_string();
295
296 Ok(upload_mode)
297 }
298
299 // ---------------------------------------------------------------------------
300 // Step 3: Compute SHA-256
301 // ---------------------------------------------------------------------------
302
303 async fn compute_sha256(
304 path: &std::path::Path,
305 file_size: u64,
306 tx: &tokio::sync::mpsc::UnboundedSender<UploadProgress>,
307 ) -> Result<String> {
308 use std::io::Read;
309
310 let path = path.to_path_buf();
311 let tx = tx.clone();
312
313 tokio::task::spawn_blocking(move || {
314 let mut file = std::fs::File::open(&path)
315 .with_context(|| format!("failed to open {}", path.display()))?;
316
317 let mut hasher = Sha256::new();
318 let mut buf = vec![0u8; HASH_CHUNK_SIZE];
319 let mut done: u64 = 0;
320
321 loop {
322 let n = file
323 .read(&mut buf)
324 .with_context(|| format!("failed to read {}", path.display()))?;
325 if n == 0 {
326 break;
327 }
328 hasher.update(&buf[..n]);
329 done += n as u64;
330
331 let _ = tx.send(UploadProgress::Hashing {
332 bytes_done: done,
333 bytes_total: file_size,
334 });
335 }
336
337 let hash = hasher.finalize();
338 let hex: String = hash.iter().map(|b| format!("{b:02x}")).collect();
339 Ok(hex)
340 })
341 .await
342 .context("SHA-256 task panicked")?
343 }
344
345 // ---------------------------------------------------------------------------
346 // Step 4: LFS batch
347 // ---------------------------------------------------------------------------
348
349 enum LfsBatchResult {
350 /// The object already exists on the server; no upload needed.
351 AlreadyExists,
352 /// The object must be uploaded to the given URL.
353 NeedsUpload {
354 upload_url: String,
355 verify_url: String,
356 verify_token: String,
357 },
358 }
359
360 async fn lfs_batch(
361 client: &reqwest::Client,
362 token: &str,
363 repo_id: &str,
364 sha256: &str,
365 file_size: u64,
366 ) -> Result<LfsBatchResult> {
367 let url = format!("https://huggingface.co/{repo_id}.git/info/lfs/objects/batch");
368
369 let body = serde_json::json!({
370 "operation": "upload",
371 "transfers": ["basic"],
372 "objects": [{
373 "oid": sha256,
374 "size": file_size
375 }],
376 "hash_algo": "sha256"
377 });
378
379 let response = client
380 .post(&url)
381 .bearer_auth(token)
382 .header("Accept", "application/vnd.git-lfs+json")
383 .header("Content-Type", "application/vnd.git-lfs+json")
384 .json(&body)
385 .send()
386 .await
387 .context("LFS batch request failed")?;
388
389 let status = response.status();
390 let text = response.text().await.unwrap_or_default();
391 eprintln!(
392 "[upload] lfs batch: {status} (response len: {})",
393 text.len()
394 );
395
396 if !status.is_success() {
397 anyhow::bail!("LFS batch failed: {status} {text}");
398 }
399
400 let parsed: serde_json::Value =
401 serde_json::from_str(&text).context("failed to parse LFS batch response")?;
402
403 let obj = parsed
404 .get("objects")
405 .and_then(|o| o.as_array())
406 .and_then(|a| a.first())
407 .context("no objects in LFS batch response")?;
408
409 // If there are no "actions", the object already exists on the server.
410 let actions = match obj.get("actions") {
411 Some(a) => a,
412 None => return Ok(LfsBatchResult::AlreadyExists),
413 };
414
415 let upload_url = actions
416 .get("upload")
417 .and_then(|u| u.get("href"))
418 .and_then(|v| v.as_str())
419 .context("no upload href in LFS batch response")?
420 .to_string();
421
422 let verify_url = actions
423 .get("verify")
424 .and_then(|v| v.get("href"))
425 .and_then(|v| v.as_str())
426 .unwrap_or("")
427 .to_string();
428
429 let verify_token = actions
430 .get("verify")
431 .and_then(|v| v.get("header"))
432 .and_then(|h| h.get("Authorization"))
433 .and_then(|v| v.as_str())
434 .unwrap_or("")
435 .to_string();
436
437 Ok(LfsBatchResult::NeedsUpload {
438 upload_url,
439 verify_url,
440 verify_token,
441 })
442 }
443
444 // ---------------------------------------------------------------------------
445 // Step 5: Upload to S3
446 // ---------------------------------------------------------------------------
447
448 async fn upload_to_s3(
449 client: &reqwest::Client,
450 upload_url: &str,
451 file_path: &std::path::Path,
452 file_size: u64,
453 tx: &tokio::sync::mpsc::UnboundedSender<UploadProgress>,
454 ) -> Result<()> {
455 let file = tokio::fs::File::open(file_path)
456 .await
457 .with_context(|| format!("failed to open {}", file_path.display()))?;
458
459 let reader = tokio::io::BufReader::new(file);
460 let progress_tx = tx.clone();
461
462 let stream = futures::stream::unfold(
463 (reader, 0u64, file_size, progress_tx),
464 |(mut reader, mut sent, total, progress)| async move {
465 use tokio::io::AsyncReadExt;
466 let mut buf = vec![0u8; CHUNK_SIZE];
467 match reader.read(&mut buf).await {
468 Ok(0) => None,
469 Ok(n) => {
470 sent += n as u64;
471 let _ = progress.send(UploadProgress::Uploading {
472 bytes_sent: sent,
473 bytes_total: total,
474 });
475 buf.truncate(n);
476 Some((
477 Ok::<_, std::io::Error>(buf),
478 (reader, sent, total, progress),
479 ))
480 }
481 Err(error) => Some((Err(error), (reader, sent, total, progress))),
482 }
483 },
484 );
485
486 let body = reqwest::Body::wrap_stream(stream);
487
488 let response = client
489 .put(upload_url)
490 .header("Content-Length", file_size)
491 .header("Content-Type", "application/octet-stream")
492 .body(body)
493 .send()
494 .await
495 .context("S3 upload request failed")?;
496
497 let status = response.status();
498 if !status.is_success() {
499 let resp_body = response.text().await.unwrap_or_default();
500 anyhow::bail!("S3 upload failed: {status} {resp_body}");
501 }
502
503 eprintln!("[upload] S3 upload complete: {status}");
504 Ok(())
505 }
506
507 // ---------------------------------------------------------------------------
508 // Step 6: LFS verify
509 // ---------------------------------------------------------------------------
510
511 async fn lfs_verify(
512 client: &reqwest::Client,
513 verify_url: &str,
514 verify_token: &str,
515 sha256: &str,
516 file_size: u64,
517 ) -> Result<()> {
518 if verify_url.is_empty() {
519 eprintln!("[upload] no verify URL, skipping");
520 return Ok(());
521 }
522
523 let body = serde_json::json!({
524 "oid": sha256,
525 "size": file_size
526 });
527
528 let mut request = client
529 .post(verify_url)
530 .header("Accept", "application/vnd.git-lfs+json")
531 .header("Content-Type", "application/vnd.git-lfs+json")
532 .json(&body);
533
534 if !verify_token.is_empty() {
535 request = request.header("Authorization", verify_token);
536 }
537
538 let response = request.send().await.context("LFS verify request failed")?;
539
540 let status = response.status();
541 let text = response.text().await.unwrap_or_default();
542 eprintln!("[upload] lfs verify: {status} {text}");
543
544 if !status.is_success() {
545 anyhow::bail!("LFS verify failed: {status} {text}");
546 }
547
548 Ok(())
549 }
550
551 // ---------------------------------------------------------------------------
552 // Step 7: LFS commit
553 // ---------------------------------------------------------------------------
554
555 async fn lfs_commit(
556 client: &reqwest::Client,
557 token: &str,
558 repo_id: &str,
559 path_in_repo: &str,
560 commit_message: &str,
561 sha256: &str,
562 file_size: u64,
563 ) -> Result<()> {
564 let url = format!("https://huggingface.co/api/models/{repo_id}/commit/main");
565
566 // Build the LFS pointer content that HuggingFace expects.
567 let lfs_pointer = format!(
568 "version https://git-lfs.github.com/spec/v1\noid sha256:{sha256}\nsize {file_size}\n"
569 );
570
571 let body = serde_json::json!({
572 "summary": commit_message,
573 "files": [{
574 "path": path_in_repo,
575 "content": lfs_pointer,
576 "lfs": {
577 "oid": sha256,
578 "size": file_size
579 }
580 }]
581 });
582
583 let response = client
584 .post(&url)
585 .bearer_auth(token)
586 .json(&body)
587 .send()
588 .await
589 .context("commit request failed")?;
590
591 let status = response.status();
592 let text = response.text().await.unwrap_or_default();
593 eprintln!("[upload] commit: {status} {text}");
594
595 if !status.is_success() {
596 anyhow::bail!("commit failed: {status} {text}");
597 }
598
599 Ok(())
600 }
601
602 // ---------------------------------------------------------------------------
603 // Regular (non-LFS) commit for small files
604 // ---------------------------------------------------------------------------
605
606 #[allow(clippy::too_many_arguments)]
607 async fn regular_commit(
608 client: &reqwest::Client,
609 token: &str,
610 repo_id: &str,
611 path_in_repo: &str,
612 commit_message: &str,
613 file_path: &std::path::Path,
614 file_size: u64,
615 tx: &tokio::sync::mpsc::UnboundedSender<UploadProgress>,
616 ) -> Result<()> {
617 let content = tokio::fs::read(file_path)
618 .await
619 .with_context(|| format!("failed to read {}", file_path.display()))?;
620
621 let _ = tx.send(UploadProgress::Uploading {
622 bytes_sent: file_size,
623 bytes_total: file_size,
624 });
625
626 // For small regular files, send content as a UTF-8 string.
627 // Binary files should always go through the LFS path instead.
628 let content_str = String::from_utf8_lossy(&content).into_owned();
629
630 let url = format!("https://huggingface.co/api/models/{repo_id}/commit/main");
631
632 let body = serde_json::json!({
633 "summary": commit_message,
634 "files": [{
635 "path": path_in_repo,
636 "content": content_str
637 }]
638 });
639
640 let _ = tx.send(UploadProgress::Committing);
641
642 let response = client
643 .post(&url)
644 .bearer_auth(token)
645 .json(&body)
646 .send()
647 .await
648 .context("commit request failed")?;
649
650 let status = response.status();
651 let text = response.text().await.unwrap_or_default();
652 eprintln!("[upload] regular commit: {status} {text}");
653
654 if !status.is_success() {
655 anyhow::bail!("commit failed: {status} {text}");
656 }
657
658 Ok(())
659 }