shreyask/Gemma-4-WebGPU
3
1import { useState, useEffect, useRef, useCallback } from "react";2import { useModel } from "./hooks/useModel.js";3import InputBar from "./components/InputBar.jsx";4import MessageList from "./components/MessageList.jsx";5import LoadingBar from "./components/LoadingBar.jsx";6import OrbitalHero from "./components/OrbitalHero.jsx";7 8const SYSTEM_PROMPT = "You are a helpful assistant. When given images, describe and analyze them. When given audio, transcribe or describe it. Be concise and helpful.";9 10const EXAMPLE_IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/thumb/1/16/Artemis_II_patch.svg/500px-Artemis_II_patch.svg.png";11const EXAMPLE_AUDIO_URL = "/neil-armstrong.oga";12const EXAMPLE_VIDEO_URL = "http://images-assets.nasa.gov/video/One_Small_Step_Comparison_720p/One_Small_Step_Comparison_720p~small.mp4";13 14const STARTER_PROMPTS = [15 { label: "Describe this patch", text: "What do you see in this image? Describe it in detail.", icon: "📷", imageUrl: EXAMPLE_IMAGE_URL },16 { label: "Transcribe audio", text: "Transcribe this audio recording.", icon: "🎤", audioUrl: EXAMPLE_AUDIO_URL },17 { label: "Analyze video*", text: "Describe what is happening in this video.", icon: "🎬", videoUrl: EXAMPLE_VIDEO_URL },18 { label: "Explain a concept", text: "Explain quantum entanglement in simple terms.", icon: "💡" },19];20 21const HF_MODEL_URL = "https://huggingface.co/onnx-community/gemma-4-E2B-it-ONNX";22const GEMMA_GRADIENT = "bg-gradient-to-br from-[#3186FF] to-[#4FA0FF]";23 24function calcTokPerSec(text, startTime) {25 const tokens = text.split(/\s+/).length;26 const elapsed = (performance.now() - startTime) / 1000;27 return elapsed > 0.5 ? Math.round(tokens / elapsed * 10) / 10 : null;28}29 30function StatusScreen({ children }) {31 return (32 <div className="min-h-screen flex items-center justify-center px-4">33 <div className="text-center max-w-md">34 <h1 className="text-2xl font-bold mb-3">Gemma 4 WebGPU</h1>35 {children}36 </div>37 </div>38 );39}40 41export default function App() {42 const { status, loadProgress, error, checkWebGPU, loadModel, generate } = useModel();43 44 const [messages, setMessages] = useState([]);45 const [streamingText, setStreamingText] = useState("");46 const [isStreaming, setIsStreaming] = useState(false);47 const [processingStep, setProcessingStep] = useState(null); // null | "extracting frames" | "decoding audio" | "generating"48 const [tokPerSec, setTokPerSec] = useState(null);49 const [isCached, setIsCached] = useState(() => localStorage.getItem("gemma4-cached") === "true");50 const [enableThinking, setEnableThinking] = useState(false);51 const [theme, setTheme] = useState(() => localStorage.getItem("gemma4-theme") || "system");52 53 useEffect(() => {54 const root = document.documentElement;55 root.classList.remove("dark", "light");56 if (theme !== "system") root.classList.add(theme);57 localStorage.setItem("gemma4-theme", theme);58 }, [theme]);59 60 const messagesEndRef = useRef(null);61 const genStartRef = useRef(0);62 const scrollRafRef = useRef(0);63 64 useEffect(() => {65 checkWebGPU();66 }, [checkWebGPU]);67 68 // Throttled scroll-to-bottom via rAF69 useEffect(() => {70 cancelAnimationFrame(scrollRafRef.current);71 scrollRafRef.current = requestAnimationFrame(() => {72 messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });73 });74 }, [messages, streamingText]);75 76 useEffect(() => {77 if (status === "ready") {78 setIsCached(true);79 localStorage.setItem("gemma4-cached", "true");80 }81 }, [status]);82 83 const handleSubmit = useCallback(async ({ imageUrl, audioUrl, videoUrl, text }) => {84 const userContent = [];85 if (imageUrl) userContent.push({ type: "image" });86 // Each video frame needs its own image token in the template87 if (videoUrl) for (let i = 0; i < 4; i++) userContent.push({ type: "image" });88 if (audioUrl) userContent.push({ type: "audio" });89 userContent.push({ type: "text", text: text || "Describe this." });90 91 const userMsg = { role: "user", content: userContent, imageUrl, audioUrl, videoUrl };92 const newMessages = [...messages, userMsg];93 setMessages(newMessages);94 95 const apiMessages = [96 { role: "system", content: SYSTEM_PROMPT },97 ...newMessages.map((m) => ({ role: m.role, content: m.content })),98 ];99 100 setStreamingText("");101 setIsStreaming(true);102 setTokPerSec(null);103 setProcessingStep(videoUrl ? "extracting frames" : audioUrl ? "decoding audio" : "generating");104 genStartRef.current = performance.now();105 106 generate({107 messages: apiMessages,108 imageUrl,109 videoUrl,110 audioUrl,111 enableThinking,112 onUpdate: (text) => {113 setProcessingStep(null);114 const tps = calcTokPerSec(text, genStartRef.current);115 if (tps !== null) setTokPerSec(tps);116 setStreamingText(text);117 },118 onComplete: (text, err) => {119 setProcessingStep(null);120 if (!err && text) {121 setTokPerSec(calcTokPerSec(text, genStartRef.current));122 setMessages((prev) => [...prev, { role: "assistant", content: text }]);123 }124 setStreamingText("");125 setIsStreaming(false);126 },127 });128 }, [messages, generate, enableThinking]);129 130 if (status === "webgpu-unavailable") {131 return (132 <StatusScreen>133 <p className="text-[var(--color-text-secondary)]">WebGPU is required. Use Chrome 113+ or Edge 113+.</p>134 </StatusScreen>135 );136 }137 138 if (error) {139 return (140 <StatusScreen>141 <p className="text-[var(--color-red)] text-sm font-mono">Error: {error}</p>142 </StatusScreen>143 );144 }145 146 const isLoading = status === "idle" || status === "webgpu-available" || status === "loading";147 148 return (149 <div className="min-h-screen flex flex-col max-w-3xl mx-auto">150 <header className="flex items-center justify-between px-4 py-3 border-b border-[var(--color-outline)]">151 <div className="flex items-center gap-2.5">152 <div className={`w-7 h-7 rounded-lg ${GEMMA_GRADIENT} flex items-center justify-center text-white text-xs font-bold`}>G</div>153 <a href={HF_MODEL_URL} target="_blank" rel="noopener" className="text-base font-medium hover:text-[var(--color-blue)] transition-colors">Gemma 4</a>154 {messages.length > 0 && !isStreaming && (155 <button156 onClick={() => { setMessages([]); setStreamingText(""); setTokPerSec(null); }}157 className="ml-2 px-2 py-0.5 text-[10px] text-[var(--color-text-secondary)] hover:text-[var(--color-text)] border border-[var(--color-outline)] rounded-lg transition-colors cursor-pointer"158 >159 New chat160 </button>161 )}162 </div>163 <div className="flex items-center gap-3 text-xs text-[var(--color-text-secondary)]">164 {!isLoading && (165 <button166 onClick={() => setEnableThinking((v) => !v)}167 className={`flex items-center gap-1.5 px-2 py-0.5 rounded-full border transition-colors cursor-pointer ${168 enableThinking169 ? "border-[var(--color-blue)]/50 bg-[var(--color-blue)]/10 text-[var(--color-blue)]"170 : "border-[var(--color-outline)] text-[var(--color-text-secondary)] hover:border-[var(--color-blue)]/30"171 }`}172 title={enableThinking ? "Thinking mode on" : "Thinking mode off"}173 >174 <span className="text-[10px] font-medium">💭 Think</span>175 </button>176 )}177 {tokPerSec != null && (178 <span className="font-mono text-[var(--color-green)]">{tokPerSec} tok/s</span>179 )}180 {isCached && !isLoading && (181 <span className="px-2 py-0.5 rounded-full bg-[var(--color-green)]/10 text-[var(--color-green)] text-[10px] font-medium">Cached</span>182 )}183 <span className="hidden sm:inline">In-Browser · WebGPU</span>184 <button185 onClick={() => setTheme((t) => t === "dark" ? "light" : t === "light" ? "system" : "dark")}186 className="p-1 rounded-full hover:bg-[var(--color-surface)] transition-colors cursor-pointer"187 title={`Theme: ${theme}`}188 >189 {theme === "dark" ? "🌙" : theme === "light" ? "☀️" : "💻"}190 </button>191 </div>192 </header>193 194 {isLoading ? (195 <div className="flex-1 flex flex-col items-center justify-center gap-6 px-4">196 <OrbitalHero />197 198 <div className="text-center -mt-2">199 <h2 className="text-4xl font-bold mb-2 tracking-tight"><a href={HF_MODEL_URL} target="_blank" rel="noopener" className="hover:text-[var(--color-blue)] transition-colors">Gemma 4 E2B</a></h2>200 <p className="text-[var(--color-text-secondary)] text-sm">Multimodal AI running entirely in your browser via WebGPU</p>201 </div>202 203 {status === "loading" ? (204 <LoadingBar loadProgress={loadProgress} isCached={isCached} />205 ) : (206 <button207 onClick={loadModel}208 className="px-8 py-3 bg-[var(--color-blue)] hover:bg-[var(--color-blue)]/90 text-white text-sm font-medium rounded-xl transition-colors cursor-pointer"209 >210 {isCached ? "Load Model (cached)" : "Load Model"}211 </button>212 )}213 214 <footer className="mt-4 text-[10px] text-[var(--color-text-secondary)]/50">215 Powered by <a href="https://huggingface.co/docs/transformers.js" target="_blank" rel="noopener" className="underline hover:text-[var(--color-text-secondary)]">Transformers.js</a>216 </footer>217 </div>218 ) : (219 <>220 <div className="flex-1 overflow-y-auto">221 {messages.length === 0 && !isStreaming ? (222 <div className="flex flex-col items-center justify-center h-full gap-6 px-4 py-12">223 <div className="text-center">224 <div className={`w-12 h-12 rounded-2xl ${GEMMA_GRADIENT} flex items-center justify-center text-white text-xl font-bold mx-auto mb-4`}>G</div>225 <h2 className="text-xl font-medium mb-1">How can I help?</h2>226 <p className="text-sm text-[var(--color-text-secondary)]">Send text, images, audio, or video — all processed locally.</p>227 </div>228 <div className="grid grid-cols-2 gap-2 max-w-md w-full">229 {STARTER_PROMPTS.map((p) => (230 <button231 key={p.label}232 onClick={() => handleSubmit({ imageUrl: p.imageUrl || null, audioUrl: p.audioUrl || null, videoUrl: p.videoUrl || null, text: p.text })}233 className="text-left p-3 bg-[var(--color-surface)] hover:bg-[var(--color-surface-high)] border border-[var(--color-outline)] rounded-xl text-sm transition-colors cursor-pointer"234 >235 {p.imageUrl && (236 <img src={p.imageUrl} alt="" className="w-full h-20 object-contain rounded-lg mb-2 bg-black/20" />237 )}238 {p.audioUrl && (239 <div className="mb-2 text-[10px] text-[var(--color-text-secondary)]/60 truncate">🎤 Neil Armstrong — Apollo 11</div>240 )}241 {p.videoUrl && (242 <div className="mb-2 text-[10px] text-[var(--color-text-secondary)]/60 truncate">🎬 One Small Step — NASA (0:56)</div>243 )}244 <span className="mr-1.5">{p.icon}</span>245 <span className="text-[var(--color-text-secondary)]">{p.label}</span>246 </button>247 ))}248 </div>249 </div>250 ) : (251 <MessageList messages={messages} streamingText={streamingText} isStreaming={isStreaming} processingStep={processingStep} />252 )}253 <div ref={messagesEndRef} />254 </div>255 256 <InputBar onSubmit={handleSubmit} disabled={isStreaming} />257 258 <div className="text-center py-2 text-[10px] text-[var(--color-text-secondary)]/40">259 {isCached && <span>Cached · </span>}260 Powered by <a href="https://huggingface.co/docs/transformers.js" target="_blank" rel="noopener" className="underline hover:text-[var(--color-text-secondary)]">Transformers.js</a>261 <span className="block mt-0.5">*Video analyzes 4 sampled frames — a tradeoff between memory and processing speed</span>262 </div>263 </>264 )}265 </div>266 );267}268