woolball-client
Version:
Client-side library for Woolball enabling secure browser resource sharing for distributed AI task processing
65 lines (64 loc) • 2.58 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.imageTextToText = imageTextToText;
const media_1 = require("../../media");
const SUPPORTED_MODELS = [
'HuggingFaceTB/SmolVLM-256M-Instruct',
];
async function imageTextToText(data) {
const { input, model = 'HuggingFaceTB/SmolVLM-256M-Instruct', dtype, max_new_tokens = 64, provider = 'transformers', } = data;
const { image, text } = typeof input === 'string' ? JSON.parse(input) : input;
if (!image || !text) {
throw new Error('Input must contain image (base64) and text');
}
if (provider === 'prompt-api') {
return await processPromptAPI(image, text);
}
return await processVision(image, text, model, max_new_tokens, dtype);
}
async function processPromptAPI(image, text) {
if (typeof window.LanguageModel === 'undefined') {
throw new Error('Prompt API is not available in this browser');
}
const session = await window.LanguageModel.create({
expectedInputs: [{ type: "image" }]
});
const imageBlob = await fetch(image).then(res => res.blob());
const result = await session.prompt([{
role: "user",
content: [
{ type: "text", value: text },
{ type: "image", value: imageBlob }
]
}]);
return { generatedText: result };
}
async function processVision(image, text, model, max_new_tokens, dtype) {
const { AutoProcessor, AutoModelForVision2Seq, RawImage, env } = await import('@huggingface/transformers');
env.allowLocalModels = false;
const { getTransformersDevice } = await import('../../../utils/environment.js');
const processor = await AutoProcessor.from_pretrained(model);
const visionModel = await AutoModelForVision2Seq.from_pretrained(model, {
dtype: (dtype || 'q4'),
device: getTransformersDevice('webgpu'),
});
const imgBlob = (0, media_1.base64ToBlob)(image);
const img = await RawImage.fromBlob(imgBlob);
const messages = [
{
role: 'user',
content: [
{ type: 'image', source: img },
{ type: 'text', text },
]
}
];
const prompt = processor.apply_chat_template(messages, { add_generation_prompt: true });
const inputs = await processor(prompt, [img]);
const output = await visionModel.generate({
...inputs,
max_new_tokens,
});
const answer = processor.batch_decode(output, { skip_special_tokens: true });
return { generatedText: answer[0] };
}