Spaces:
Running
Running
| import { | |
| AutoTokenizer, | |
| AutoProcessor, | |
| AutoModelForCausalLM, | |
| AutoModelForImageTextToText, | |
| TextStreamer, | |
| RawImage, | |
| env, | |
| } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.8.1'; | |
| env.allowLocalModels = false; | |
| env.useBrowserCache = true; | |
| // Mapping of model prefixes to their corresponding model classes | |
| const MODEL_CLASSES = { | |
| 'onnx-community/granite-4.0-micro-ONNX-web': { | |
| model: AutoModelForCausalLM, | |
| processor: AutoTokenizer | |
| }, | |
| 'mistralai/Ministral-3-3B-Instruct-2512-ONNX': { | |
| model: AutoModelForImageTextToText, | |
| processor: AutoProcessor | |
| }, | |
| }; | |
| class TextGenerationPipeline { | |
| static processor = null; // Can be a Processor or a Tokenizer | |
| static model = null; | |
| static model_id = null; // Track loaded model | |
| static model_class = null; // Track loaded model class | |
| static async getInstance(model_id, progress_callback = null) { | |
| if (model_id && model_id !== this.model_id) { | |
| // Force reload if a new model is requested | |
| this.processor = null; | |
| this.model = null; | |
| this.model_class = null; | |
| } | |
| if (this.processor === null || this.model === null) { | |
| if (!model_id) throw new Error("Model not loaded and no model ID provided to load."); | |
| const modelInfo = MODEL_CLASSES[model_id] || { model: AutoModelForCausalLM, processor: AutoTokenizer }; | |
| this.model_class = modelInfo.model; | |
| const ProcessorClass = modelInfo.processor; | |
| this.processor = await ProcessorClass.from_pretrained(model_id, { | |
| progress_callback, | |
| }); | |
| // Set the image processor size for vision models | |
| if (this.model_class === AutoModelForImageTextToText && this.processor.image_processor) { | |
| this.processor.image_processor.size = { longest_edge: 480 }; | |
| } | |
| const modelConfig = { | |
| device: "webgpu", | |
| progress_callback, | |
| }; | |
| if (this.model_class === AutoModelForImageTextToText) { | |
| modelConfig.dtype = { | |
| embed_tokens: "fp16", | |
| vision_encoder: "q4", | |
| decoder_model_merged: "q4f16", | |
| }; | |
| } else { | |
| modelConfig.dtype = "q4f16"; | |
| } | |
| this.model = await this.model_class.from_pretrained(model_id, modelConfig); | |
| this.model_id = model_id; | |
| } | |
| return [this.processor, this.model, this.model_class]; | |
| } | |
| } | |
| self.addEventListener("message", async (e) => { | |
| const { type, data } = e.data; | |
| switch (type) { | |
| case 'load': | |
| await loadModel(data.model_id); | |
| break; | |
| case 'generate': | |
| await generateText(data); | |
| break; | |
| } | |
| }); | |
| async function loadModel(model_id) { | |
| self.postMessage({ | |
| status: "loading", | |
| message: "Loading model...", | |
| }); | |
| try { | |
| const [processor, model, model_class] = await TextGenerationPipeline.getInstance(model_id, (x) => { | |
| self.postMessage({ status: 'progress', data: x }); | |
| }); | |
| self.postMessage({ | |
| status: "loading", | |
| message: "Compiling shaders and warming up model...", | |
| }); | |
| // Perform a dry run to compile shaders | |
| if (model_class === AutoModelForImageTextToText) { | |
| const image = new RawImage(new Uint8ClampedArray(480 * 480 * 4).fill(0), 480, 480, 4); | |
| const messages = [ | |
| { role: "system", content: "You are a helpful AI assistant." }, | |
| { role: "user", content: "[IMG]a" } | |
| ]; | |
| self.postMessage({ status: 'log', message: 'Warming up model... (1/3) Applying chat template' }); | |
| let prompt, inputs; | |
| try { | |
| prompt = processor.apply_chat_template(messages); | |
| } catch (e) { | |
| self.postMessage({ status: 'log', message: 'ERROR during apply_chat_template' }); | |
| throw e; | |
| } | |
| self.postMessage({ status: 'log', message: 'Warming up model... (2/3) Processing inputs' }); | |
| try { | |
| inputs = await processor(image, prompt); | |
| } catch (e) { | |
| self.postMessage({ status: 'log', message: 'ERROR during processor(image, prompt)' }); | |
| throw e; | |
| } | |
| self.postMessage({ status: 'log', message: 'Warming up model... (3/3) Generating' }); | |
| try { | |
| await model.generate({ ...inputs, max_new_tokens: 1 }); | |
| } catch (e) { | |
| self.postMessage({ status: 'log', message: 'ERROR during model.generate' }); | |
| throw e; | |
| } | |
| } else { | |
| const inputs = processor("a"); | |
| await model.generate({ ...inputs, max_new_tokens: 1 }); | |
| } | |
| self.postMessage({ status: "ready", message: `Successfully loaded ${model_id}` }); | |
| } catch (err) { | |
| handleError(err); | |
| } | |
| } | |
| async function generateText({ messages, image }) { | |
| try { | |
| const [processor, model, model_class] = await TextGenerationPipeline.getInstance(); | |
| let inputs; | |
| if (model_class === AutoModelForImageTextToText) { | |
| let final_image; | |
| // Create a shallow copy of the messages to avoid mutation | |
| const processedMessages = JSON.parse(JSON.stringify(messages)); | |
| if (image) { | |
| // Use the user-uploaded image | |
| final_image = new RawImage(image.data, image.width, image.height, 4); | |
| // Prepend [IMG] to the last user message | |
| const lastUserMessage = processedMessages.filter(x => x.role === 'user').pop(); | |
| if (lastUserMessage) { | |
| lastUserMessage.content = `[IMG]${lastUserMessage.content}`; | |
| } | |
| } else { | |
| // Fallback to a dummy black image for text-only queries | |
| final_image = new RawImage(new Uint8ClampedArray(480 * 480 * 4).fill(0), 480, 480, 4); | |
| } | |
| const prompt = processor.apply_chat_template(processedMessages, { add_generation_prompt: true }); | |
| inputs = await processor(final_image, prompt); | |
| } else { | |
| // Standard text-only chat model | |
| inputs = processor.apply_chat_template(messages, { | |
| add_generation_prompt: true, | |
| return_dict: true, | |
| }); | |
| } | |
| const tokenizer = model_class === AutoModelForImageTextToText ? processor.tokenizer : processor; | |
| const streamer = new TextStreamer(tokenizer, { | |
| skip_prompt: true, | |
| skip_special_tokens: true, | |
| callback_function: (text) => { | |
| self.postMessage({ status: 'stream-chunk', text: text }); | |
| }, | |
| }); | |
| self.postMessage({ status: 'log', message: 'Generating response...' }); | |
| const output = await model.generate({ | |
| ...inputs, | |
| max_new_tokens: 512, | |
| temperature: 0.7, | |
| do_sample: true, | |
| top_k: 20, | |
| streamer: streamer, | |
| return_dict_in_generate: true, | |
| }); | |
| const decoded = processor.batch_decode(output.sequences, { skip_special_tokens: true }); | |
| self.postMessage({ status: 'complete', output: decoded[0] }); | |
| } catch (err) { | |
| handleError(err); | |
| } | |
| } | |
| async function handleError(err) { | |
| // Stringify the error to ensure all information is sent | |
| const errorObject = { | |
| message: err.message || 'No error message', | |
| stack: err.stack || 'No stack trace', | |
| name: err.name || 'No error name', | |
| // Include any other properties that might be useful | |
| ...err | |
| }; | |
| const errorMsg = JSON.stringify(errorObject, Object.getOwnPropertyNames(err)); | |
| self.postMessage({ status: 'error', message: errorMsg, stack: err.stack }); | |
| } | |