import { json, type RequestEvent } from '@sveltejs/kit'; import puppeteer from 'puppeteer'; import { promises } from 'fs'; import { randomUUID } from 'crypto'; import prisma from '$lib/prisma'; import { tokenIsAvailable } from '$lib/utils'; /** @type {import('./$types').RequestHandler} */ export async function GET({ url, params } : RequestEvent) { const id = params.id?.replace("@", "/") const full = Boolean(url.searchParams.get('full')) ?? false const metadata = Boolean(url.searchParams.get('metadata')) ?? false const model = await prisma.model.findFirst({ where: { id, }, select: full ? { id: true, likes: true, downloads: true, image: true, instance_prompt: true, metadata: true, gallery: { select: { id: true, prompt: true, image: true, createdAt: true, }, where: { isPublic: true }, orderBy: { createdAt: 'desc' }, take: 10 }, comments: { select: { id: true, createdAt: true, text: true, user: { select: { id: true, name: true, sub: true, picture: true, preferred_username: true, } } } } } : { instance_prompt: true, image: true, id: true, metadata: true, } }) if (!model) { return json({ error: { token: "Model params is required" } }, { status: 401 }) } let infos: Record = {}; if (full) { const hf_model_request = await fetch(`https://huggingface.co/api/models/${id}`, { headers: { "Authorization": `Bearer ${process.env.SECRET_HF_TOKEN}` } }) const hf_model_response = await hf_model_request.json(); infos = { base_model: hf_model_response?.cardData?.base_model, license: hf_model_response?.cardData?.license, tags: hf_model_response?.cardData?.tags, } } // get puppeteer screenshot if (!model.metadata && metadata) { const browser = await puppeteer.launch({ headless: true, defaultViewport: { width: 1920, height: 1080, }, args: [ '--no-sandbox', '--disable-setuid-sandbox', '--disable-dev-shm-usage', ], ignoreDefaultArgs: ['--disable-extensions'] }); const page = await browser.newPage(); await page.goto(process.env.SPACE_HOST + '/metadata/models/' + model.id, { waitUntil: 'networkidle0', }); const buffer = await page.screenshot({ type: 'png', }); const dir = await promises.opendir(process?.env?.PUBLIC_FILE_UPLOAD_DIR as string).catch(() => null) if (!dir) await promises.mkdir(process?.env?.PUBLIC_FILE_UPLOAD_DIR as string) const file_name_formatted = randomUUID() + "_" + "metadata" + model.id.replace("/", "-") + ".png" await promises.writeFile(`${process.env.PUBLIC_FILE_UPLOAD_DIR}/${file_name_formatted}`, buffer) await prisma.model.update({ where: { id, }, data: { metadata: file_name_formatted } }) await browser.close(); } return json({ model: { ...model, infos } }) } export async function POST({ params, cookies, fetch } : RequestEvent) { const token = cookies.get('hf_access_token') if (!token) { return json({ error: { token: "Token is required" } }, { status: 401 }) } const user = await tokenIsAvailable(token) if (!user || !process.env.SECRET_HF_ADMIN?.includes(user.sub)) { return json({ error: { token: "Wrong castle fam :^)" } }, { status: 401 }) } const id = params.id?.replace("@", "/") const model = await prisma.model.findFirst({ where: { id, } }) if (!model) { return json({ error: { token: "Model not found" } }, { status: 404 }) } const data = await fetch(`https://huggingface.co/api/models/${model.id}`) const hf_model = await data.json() await prisma.model.update({ where: { id, }, data: { isPublic: true, instance_prompt: hf_model?.cardData?.instance_prompt } }) return json({ success: true }) } export async function DELETE({ params, cookies } : RequestEvent) { const token = cookies.get('hf_access_token') if (!token) { return json({ error: { token: "Token is required" } }, { status: 401 }) } const user = await tokenIsAvailable(token) if (!user || !process.env.SECRET_HF_ADMIN?.includes(user.sub)) { return json({ error: { token: "Wrong castle fam :^)" } }, { status: 401 }) } const id = params.id?.replace("@", "/") const model = await prisma.model.findFirst({ where: { id, } }) if (!model) { return json({ error: { token: "Model not found" } }, { status: 404 }) } await prisma.model.update({ where: { id, }, data: { isPublic: false, } }) return json({ success: true }) } export async function PATCH({ request, params } : RequestEvent) { const headers = Object.fromEntries(request.headers.entries()); if (headers["x-hf-token"] !== process.env.SECRET_HF_TOKEN) { return Response.json({ message: "Wrong castle fam :^)" }, { status: 401 }); } const id = params.id?.replace("@", "/") const model = await prisma.model.findFirst({ where: { id, } }) if (!model) { return json({ error: { token: "Model not found" } }, { status: 404 }) } const body = await request.json() await prisma.model.update({ where: { id, }, data: { ...body } }) return json({ success: true }) }