import { createFileRoute } from "@tanstack/react-router"; import { zipSync, strToU8 } from "fflate"; import { z } from "zod"; import { writeFileSync, unlinkSync } from "fs"; import { execSync } from "child_process"; const PostSchema = z.object({ id: z.string(), fileUrl: z.string().url(), ext: z.string().default("jpg"), previewUrl: z.string().optional(), sampleUrl: z.string().optional(), tags: z.array(z.string()).default([]), source: z.string().optional(), rating: z.string().optional(), score: z.number().optional(), fileSize: z.number().optional(), width: z.number().optional(), height: z.number().optional(), md5: z.string().optional(), uploader: z.union([z.string(), z.number()]).optional(), createdAt: z.string().optional(), rawJson: z.string().optional(), }); const BodySchema = z.object({ siteName: z.string().default("grabber"), posts: z.array(PostSchema).min(1).max(200), excludeTags: z.array(z.string()).default([]), hfToken: z.string().min(1, "HF token is required"), datasetName: z.string().min(1, "Dataset name is required"), }); function sanitize(s: string) { return s.replace(/[^a-zA-Z0-9._-]+/g, "_").slice(0, 60); } function humanSize(n?: number): string { if (n == null || !Number.isFinite(n)) return ""; if (n < 1024) return `${n} B`; if (n < 1024 * 1024) return `${(n / 1024).toFixed(1)} KB`; if (n < 1024 * 1024 * 1024) return `${(n / 1024 / 1024).toFixed(2)} MB`; return `${(n / 1024 / 1024 / 1024).toFixed(2)} GB`; } function buildMarkdown(post: z.infer, excludeSet: Set): string { const dims = post.width && post.height ? `${post.width} x ${post.height}` : ""; const rows: [string, string][] = [ ["Source", post.source ?? ""], ["Rating", post.rating ?? ""], ["Score", post.score != null ? String(post.score) : ""], ["Dimensions", dims], ["File size", humanSize(post.fileSize)], ["MD5", post.md5 ?? ""], ["Uploader", post.uploader != null ? String(post.uploader) : ""], ["Created at", post.createdAt ?? ""], ["Original file", post.fileUrl], ["Sample", post.sampleUrl ?? ""], ["Preview", post.previewUrl ?? ""], ["Ext", post.ext], ]; const meta = rows.map(([k, v]) => `- **${k}:** ${v}`).join("\n"); const filteredTags = post.tags.filter((t) => !excludeSet.has(t.toLowerCase())); const tags = filteredTags.length ? filteredTags.join(" ") : ""; let raw = "{}"; if (post.rawJson) { try { raw = JSON.stringify(JSON.parse(post.rawJson), null, 2); } catch { raw = post.rawJson; } } return `# Post ${post.id} ${meta} ## Tags ${tags} ## Raw metadata \`\`\`json ${raw} \`\`\` `; } export const Route = createFileRoute("/api/upload-to-dataset")({ server: { handlers: { POST: async ({ request }) => { let body: z.infer; try { body = BodySchema.parse(await request.json()); } catch (err) { return new Response( JSON.stringify({ error: err instanceof Error ? err.message : "Invalid body" }), { status: 400, headers: { "Content-Type": "application/json" } }, ); } const files: Record = {}; const errors: string[] = []; const siteName = sanitize(body.siteName); const excludeSet = new Set((body.excludeTags || []).map((t) => t.toLowerCase())); // Generate files await Promise.all( body.posts.map(async (post) => { const idSafe = sanitize(post.id); const extSafe = sanitize(post.ext || "jpg"); const base = `${siteName}_${idSafe}`; const filteredTags = post.tags.filter((t) => !excludeSet.has(t.toLowerCase())); files[`${base}.txt`] = strToU8(filteredTags.join(" ")); files[`${idSafe}.md`] = strToU8(buildMarkdown(post, excludeSet)); try { const res = await fetch(post.fileUrl, { headers: { "User-Agent": "LovableGrabber/1.0" }, }); if (!res.ok) { errors.push(`${post.id}: HTTP ${res.status}`); return; } const buf = new Uint8Array(await res.arrayBuffer()); files[`${base}.${extSafe}`] = buf; } catch (err) { errors.push(`${post.id}: ${err instanceof Error ? err.message : "fetch failed"}`); } }), ); if (errors.length) { files["_errors.txt"] = strToU8(errors.join("\n")); } if (Object.keys(files).length === 0) { return new Response(JSON.stringify({ error: "No files successfully compiled" }), { status: 502, headers: { "Content-Type": "application/json" }, }); } // Build ZIP Sync const zipped = zipSync(files, { level: 0 }); const ts = new Date().toISOString().replace(/[:.]/g, "-"); const filename = `grabber-${ts}.zip`; // Temporary paths const zipPath = `/tmp/${filename}`; const pyScriptPath = `/tmp/hf_dataset_upload_${ts}.py`; try { // Write compiled zip to local temp folder writeFileSync(zipPath, zipped); // Write python uploader helper script const pyCode = ` import sys import os from huggingface_hub import HfApi token = sys.argv[1] repo_id = sys.argv[2] file_path = sys.argv[3] path_in_repo = sys.argv[4] api = HfApi() try: # Attempt to create the dataset repository if not exists api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True, token=token) # Upload the compiled zip to the dataset repository api.upload_file( path_or_fileobj=file_path, path_in_repo=path_in_repo, repo_id=repo_id, repo_type="dataset", token=token ) print("SUCCESS") sys.exit(0) except Exception as e: print(f"ERROR: {str(e)}", file=sys.stderr) sys.exit(1) `; writeFileSync(pyScriptPath, pyCode); // Execute python helper to create/upload the dataset execSync( `python3 "${pyScriptPath}" "${body.hfToken}" "${body.datasetName}" "${zipPath}" "${filename}"`, ); return new Response(JSON.stringify({ success: true, filename }), { status: 200, headers: { "Content-Type": "application/json" }, }); } catch (execErr: unknown) { let errMsg = "Execution error"; if (execErr instanceof Error) { const stderr = (execErr as { stderr?: Buffer }).stderr; errMsg = stderr ? stderr.toString() : execErr.message; } return new Response( JSON.stringify({ error: `Hugging Face dataset upload failed: ${errMsg}` }), { status: 500, headers: { "Content-Type": "application/json" }, }, ); } finally { // Clean up temporary files try { unlinkSync(zipPath); } catch { // ignore } try { unlinkSync(pyScriptPath); } catch { // ignore } } }, }, }, });