fix warmup

Author: pdufour

src/Chat.js

@@ -25,11 +25,13 @@ export function Chat() {
   let isGenerating;
   let error;
   let generatedText;
+  let warmup;
+
   if (INTERFACE === 'IMAGE') {
-    ({ generateText, isGenerating, error, generatedText } =
+    ({ generateText, isGenerating, error, generatedText, warmup } =
       useLLMVisionGeneration(LLM_VISION_MODEL_CONFIG));
   } else {
-    ({ generateCode: generateText, isGenerating, error, generatedCode: generatedText } =
+    ({ generateCode: generateText, isGenerating, error, generatedCode: generatedText, warmup } =
       useLLMHtmlGeneration(LLM_HTML_MODEL_CONFIG));
   }
   const [currentMessageId, setCurrentMessageId] = useState(null);
@@ -84,6 +86,16 @@ export function Chat() {
     }
   }, [generatedText, currentMessageId]);
 
+  const onChangeInput = (e) => {
+    setInput(e.target.value);
+    if (showWarning) {
+      const proceed = window.confirm("Warning: Using this chat will download AI models larger than 1GB in size. Do you want to continue?");
+      if (!proceed) return;
+      warmup();
+      setShowWarning(false);
+    }
+  }
+
   const handleImageUpload = (event) => {
     const file = event.target.files[0];
     if (file) {
@@ -100,12 +112,6 @@ export function Chat() {
     e.preventDefault();
     if ((!input.trim() && !selectedImage) || isGenerating) return;
 
-    if (showWarning) {
-      const proceed = window.confirm("Warning: Using this chat will download AI models larger than 1GB in size. Do you want to continue?");
-      if (!proceed) return;
-      setShowWarning(false);
-    }
-
     const userMessage = {
       role: "user",
       content: input,
@@ -128,17 +134,6 @@ export function Chat() {
     h(
       "div",
       { className: "fixed top-4 right-4 z-20 flex items-center gap-2" },
-      hasCache && h(
-        "button",
-        {
-          onClick: clearModelCache,
-
-          className: "bg-gray-900/80 backdrop-blur-sm rounded-full p-2 text-gray-300 hover:text-white hover:bg-gray-800 transition-colors flex items-center gap-2",
-          title: `Clear cached models (${cacheSize}MB)`
-        },
-        h(XCircle, { className: "w-5 h-5" }),
-        `Delete Cache ${cacheSize}MB`
-      ),
       h(
         "a",
         {
@@ -149,7 +144,17 @@ export function Chat() {
         },
         h(Github, { className: "w-5 h-5" }),
         "GitHub"
-      )
+      ), hasCache && h(
+        "button",
+        {
+          onClick: clearModelCache,
+
+          className: "bg-gray-900/80 backdrop-blur-sm rounded-full p-2 text-gray-300 hover:text-white hover:bg-gray-800 transition-colors flex items-center gap-2",
+          title: `Clear cached models (${cacheSize}MB)`
+        },
+        h(XCircle, { className: "w-5 h-5" }),
+        `Delete Cache ${cacheSize}MB`
+      ),
     ),
     !isGenerating && !messages.length
       ? h(
@@ -298,7 +303,7 @@ export function Chat() {
           h("input", {
             type: "text",
             value: input,
-            onChange: (e) => setInput(e.target.value),
+            onChange: (e) => onChangeInput(e),
             placeholder: "Describe what you want to create...",
             className:
               "flex-1 bg-transparent px-4 py-3 focus:outline-none placeholder-gray-500",

src/constants/chat.js

@@ -22,6 +22,9 @@ export const LLM_HTML_MODEL_CONFIG = {
         temperature: 0.3,
         top_p: 0.9,
       },
+      warmup: {
+        max_tokens: 1,
+      }
     },
     huggingface: {
       modelId: "Qwen/Qwen2.5-Coder-1.5B-Instruct",
@@ -35,6 +38,9 @@ export const LLM_HTML_MODEL_CONFIG = {
         top_p: 0.9,
         do_sample: true,
       },
+      warmup: {
+        max_new_tokens: 1,
+      }
     },
   },
   backend: "webllm",

src/hooks/useLLMGeneration.js

@@ -374,19 +374,28 @@ class Qwen2VLBackend {
 
 class WebLLMBackend {
   constructor(modelId, config) {
-    this.modelId = modelId
-    this.config = config
+    this.modelId = modelId;
+    this.config = config;
   }
 
-  async generate(prompt, systemPrompt, callbacks) {
-    const engine = await CreateMLCEngine(this.modelId)
+  async warmup() {
+    if (!this.warmupPromise) {
+      this.warmupPromise = (async () => {
+        console.log('Creating new engine');
+        this.engine = await CreateMLCEngine(this.modelId);
+        return this.engine;
+      })();
+    }
+    return this.warmupPromise;
+  }
 
+  async generate(prompt, systemPrompt, callbacks) {
     const messages = [
       { role: "system", content: systemPrompt },
       { role: "user", content: prompt }
     ]
 
-    const asyncChunkGenerator = await engine.chat.completions.create({
+    const asyncChunkGenerator = await this.engine.chat.completions.create({
       messages,
       stream: true,
       ...this.config
@@ -491,8 +500,17 @@ export function useLLMGeneration(
     }
   }, [backend, modelConfig])
 
+  const warmup = React.useCallback(async () => {
+    const callbacks = {
+      onToken: () => { },
+      onComplete: () => { },
+      onError: () => { }
+    };
+    await backendRef.current.warmup();
+  });
+
   const generate = React.useCallback(
-    async (prompt, extras) => {
+    async (prompt, extras = {}, config = {}) => {
       if (!backendRef.current) {
         throw new Error(`No backend configured for ${backend}`)
       }
@@ -516,7 +534,10 @@ export function useLLMGeneration(
         }
       }
 
+      backendRef.config = { ...backendRef.config, ...config };
+
       try {
+        await backendRef.current.warmup();
         await backendRef.current.generate(prompt, systemPrompt, callbacks, extras)
       } catch (err) {
         callbacks.onError(err)
@@ -528,6 +549,7 @@ export function useLLMGeneration(
 
   return {
     generate,
+    warmup,
     isGenerating,
     error,
     partialText,

src/hooks/useLLMHtmlGeneration.js

@@ -84,7 +84,7 @@ export function useLLMHtmlGeneration({
   const [generatedCode, setGeneratedCode] = React.useState("");
   const lastGeneratedCode = React.useRef("");
 
-  const { generate, isGenerating, error, partialText } = useLLMGeneration(
+  const { generate, warmup, isGenerating, error, partialText } = useLLMGeneration(
     modelConfig,
     systemPrompt,
     backend
@@ -100,12 +100,12 @@ export function useLLMHtmlGeneration({
   }, [partialText]);
 
   const generateCode = React.useCallback(
-    async (prompt) => {
+    async (prompt, extras, config) => {
       const fullPrompt = lastGeneratedCode.current
         ? `Current HTML: \n${lastGeneratedCode.current}\n\nRequest: ${prompt}`
         : `Generate the HTML for: ${prompt}`;
 
-      await generate(fullPrompt);
+      await generate(fullPrompt, extras, config);
     },
     [generate]
   );
@@ -115,5 +115,6 @@ export function useLLMHtmlGeneration({
     isGenerating,
     error,
     generatedCode,
+    warmup,
   };
 }