SIAF / src /lib /bucket.ts
GUI-STUDIO
Document bucket files, improve upload & parsing
541b32b
Raw
History Blame Contribute Delete
7.09 kB
import path from "node:path";
import { AccessToken, RepoDesignation } from "@huggingface/hub";
import * as hub from "@huggingface/hub";
import { NS_routes_c } from "../../../../../packages/lib/helpers/client/route";
import { NS_misc } from "./misc";
import { NS_promise } from "./promise";
import { NS_ffmpeg_v2 } from "./ffmpeg-v2";
const HF_BUCKET_ACCESS_TOKEN = process.env.HF_BUCKET_ACCESS_TOKEN;
const HF_BUCKET_REPO = process.env.HF_BUCKET_REPO;
export const STORAGE_ROOT =
process.env.STORAGE_ROOT || path.join(process.cwd(), "S3");
export const repo: RepoDesignation = {
type: "bucket",
name: HF_BUCKET_REPO || ""
};
export const credentials: { accessToken: AccessToken } = {
accessToken: HF_BUCKET_ACCESS_TOKEN!
};
// await hub.createRepo({
// repo,
// accessToken: credentials.accessToken,
// private: false
// });
type T_UploadJob = {
status: "in-progress" | "done" | "failed";
tmp_path?: string | null;
created_at: number;
error?: string;
percent?: number;
transcoded_size?: string;
progress_endpoint?: string;
};
const upload_jobs = new Map<string, T_UploadJob>();
export const NS_bucket = {
track_upload_process: {
v1: trackProgress,
v2: trackProgress_v2
},
get_hf_bucket_file_info,
get_hf_bucket_list_of_files,
get_latest_hf_bucket_file_by_name,
upload_to_bucket
};
function trackProgress(
blob: Blob,
onProgress: (pct: number) => void
): ReadableStream<Uint8Array> {
const total = blob.size;
let loaded = 0;
let lastReported = -1;
return new ReadableStream({
async start(controller) {
const reader = blob.stream().getReader();
while (true) {
const { done, value } = await reader.read();
if (done) break;
loaded += value.byteLength;
// only fire when % actually changes
const pct = Math.round((loaded / total) * 100);
if (pct !== lastReported) {
lastReported = pct;
onProgress(pct);
}
controller.enqueue(value);
}
controller.close();
}
});
}
function trackProgress_v2(
blob: Blob,
onProgress: (pct: number) => void
): ReadableStream {
const total = blob.size;
let loaded = 0;
let lastReported = -1;
return new ReadableStream({
async start(controller) {
const reader = blob.stream().getReader();
while (true) {
const { done, value } = await reader.read();
if (done) break;
loaded += value.byteLength;
const pct = Math.round((loaded / total) * 100);
if (pct % 5 === 0 && pct !== lastReported) {
lastReported = pct;
onProgress(pct);
}
controller.enqueue(value);
}
controller.close();
}
});
}
async function get_hf_bucket_file_info(arg0: { file_path: string }) {
const { file_path } = arg0;
if (
!credentials.accessToken ||
typeof file_path !== "string" ||
NS_routes_c.F_is_url(file_path)
) {
return;
}
try {
const info = await hub.fileDownloadInfo({
repo,
path: file_path,
accessToken: credentials.accessToken
});
return info;
} catch {
// file doesn't exist
}
}
async function get_hf_bucket_list_of_files() {
if (!credentials.accessToken) return;
const available_files = await Array.fromAsync(
hub.listFiles({
repo,
recursive: true,
accessToken: credentials.accessToken
})
);
return available_files;
}
async function get_latest_hf_bucket_file_by_name(arg0: {
folder_path: string;
human_filename: string;
sbe_id?: string;
}) {
const { folder_path, human_filename, sbe_id } = arg0;
if (
!credentials.accessToken ||
typeof folder_path !== "string" ||
typeof human_filename !== "string"
) {
return;
}
try {
const files = await Array.fromAsync(
hub.listFiles({
repo,
accessToken: credentials.accessToken,
path: folder_path,
recursive: false
})
);
// match files whose name ends with -{human_filename}
const matching = files.filter((f) => {
const name = f.path.split("/").pop() ?? "";
const withoutPrefix = name.slice(name.indexOf("-") + 1);
return withoutPrefix.toLowerCase() === human_filename.toLowerCase();
});
if (matching.length === 0) return;
// pool 1 — sbe_id prefixed
const sbe_id_pool = sbe_id
? matching.filter((f) => {
const name = f.path.split("/").pop() ?? "";
const prefix = name.slice(0, name.indexOf("-"));
return prefix === sbe_id;
})
: [];
// pool 2 — timestamp prefixed, from what's left after removing sbe_id matches
const remaining = matching.filter((f) => !sbe_id_pool.includes(f));
const timestamp_pool = remaining
.map((f) => {
const name = f.path.split("/").pop() ?? "";
const prefix = Number(name.slice(0, name.indexOf("-")));
return { f, ts: NS_misc.isValidUnixTimestamp(prefix) ? prefix : null };
})
.filter((x) => x.ts !== null)
.sort((a, b) => b.ts! - a.ts!)
.map((x) => x.f);
// prefer timestamp pool, fallback to sbe_id pool
let pool =
timestamp_pool.length > 0
? timestamp_pool
: sbe_id_pool.length > 0
? sbe_id_pool
: [];
if (true) {
const _ = matching.sort((a, b) => {
const da = new Date(a.uploadedAt ?? 0).getTime();
const db = new Date(b.uploadedAt ?? 0).getTime();
return db - da;
});
if (_.length) {
pool = [..._];
}
}
const best_file = await get_hf_bucket_file_info({
file_path: pool[0].path
});
return best_file;
} catch {}
}
async function upload_to_bucket(arg0: {
job_id: string;
bucket_path: string;
file_with_path: string;
skip?: {
probe?: boolean;
};
REQ0?: URL;
}) {
const { job_id, bucket_path, skip, file_with_path, REQ0 } = arg0;
upload_jobs.set(job_id, {
status: "in-progress",
created_at: Date.now()
});
console.log("\n", "🧲 UPLOADING TO BUCKET", "\n");
if (!skip?.probe) {
const [E_, R_] = await NS_promise.F_catch_promise({
promise: NS_ffmpeg_v2.F_probe_t1({
playback: {
url: file_with_path,
type: "video",
protocol: "local"
}
})
});
if (E_ || !R_) {
upload_jobs.delete(job_id);
console.log("\n", "❌ UPLOAD BUCKET: FFPROBE FAILED", "\n");
return;
}
}
const blob: any = Bun.file(file_with_path);
const [E__, R__] = await NS_promise.F_catch_promise({
promise: hub.uploadFiles({
repo,
accessToken: credentials.accessToken,
files: [
{
path: bucket_path,
content: blob
}
]
})
});
upload_jobs.delete(job_id);
if (E__) {
console.log("\n", "❌ UPLOAD BUCKET FAILED", E__, "\n");
return { success: false };
}
console.log("\n", "🔥 UPLOAD BUCKET DONE", "\n");
if (REQ0) {
REQ0.searchParams.set("ops", "just_check");
fetch(REQ0.toString());
}
return { success: true };
}