apssouza22/webgpu-cluster

1import {2  CLUSTER_MODELS,3  DEFAULT_CLUSTER_MODEL_ID,4  getClusterModel,5  isImplementedClusterModel,6} from '../../shared/clusterModels';7import type {ObjectDetector} from '../detection/ObjectDetector';8import type {DetectionResult} from '../detection/workerMessages';9import type {VideoDescriber} from '../videodescription/VideoDescriber';10 11const API_BASE = import.meta.env.VITE_API_BASE ?? '';12 13const OBJECT_DETECTION_MODEL_ID = 'rfdetr-medium';14const VIDEO_DESCRIPTION_MODEL_ID = 'smolvlm-500m';15 16type ClusterTaskKind = 'detect' | 'describe';17 18interface HostTaskMessage {19  id: string;20  kind: ClusterTaskKind;21  image_base64: string;22  mime_type: string;23  threshold?: number;24  instruction?: string;25  max_new_tokens?: number;26}27 28const statusEl = document.getElementById('status')!;29const hostIdInput = document.getElementById('host-id') as HTMLInputElement;30const modelSelect = document.getElementById('model-select') as HTMLSelectElement;31const registerBtn = document.getElementById('register-btn') as HTMLButtonElement;32const curlExample = document.getElementById('curl-example')!;33 34let detector: ObjectDetector | null = null;35let describer: VideoDescriber | null = null;36let hostId = '';37let selectedModelId = DEFAULT_CLUSTER_MODEL_ID;38let processing = false;39let eventSource: EventSource | null = null;40const taskQueue: HostTaskMessage[] = [];41 42function setStatus(message: string): void {43  statusEl.textContent = message;44  console.log('[detection-host]', message);45}46 47function populateModelSelect(): void {48  modelSelect.replaceChildren(49    ...CLUSTER_MODELS.map((model) => {50      const option = document.createElement('option');51      option.value = model.id;52      option.textContent = model.implemented53        ? model.label54        : `${model.label} (coming soon)`;55      option.disabled = !model.implemented;56      option.title = model.description;57      return option;58    }),59  );60  modelSelect.value = DEFAULT_CLUSTER_MODEL_ID;61}62 63function updateCurlExample(): void {64  const origin = window.location.origin;65 66  if (selectedModelId === OBJECT_DETECTION_MODEL_ID) {67    curlExample.textContent = `curl -X POST '${origin}/v1/detect' \\68  -H 'Content-Type: application/json' \\69  -d '{70    "host": "${hostId}",71    "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/coco_sample.png",72    "threshold": 0.573  }'`;74    return;75  }76 77  if (selectedModelId === VIDEO_DESCRIPTION_MODEL_ID) {78    curlExample.textContent = `curl -X POST '${origin}/v1/describe' \\79  -H 'Content-Type: application/json' \\80  -d '{81    "host": "${hostId}",82    "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/coco_sample.png",83    "instruction": "What do you see?",84    "max_new_tokens": 10085  }'`;86    return;87  }88 89  const model = getClusterModel(selectedModelId);90  curlExample.textContent = model91    ? `# ${model.label}\n# API endpoint for this model is not wired yet.`92    : '';93}94 95async function ensureModelLoaded(modelId: string): Promise<void> {96  if (modelId === OBJECT_DETECTION_MODEL_ID) {97    if (!detector) {98      setStatus('Loading RF-DETR (WebGPU)…');99      const {ObjectDetector: Detector} = await import('../detection/ObjectDetector');100      detector = await Detector.create(setStatus);101    }102    return;103  }104 105  if (modelId === VIDEO_DESCRIPTION_MODEL_ID) {106    if (!describer) {107      setStatus('Loading SmolVLM-500M (WebGPU)…');108      const {VideoDescriber: Describer} = await import('../videodescription/VideoDescriber');109      describer = await Describer.create(setStatus);110    }111    return;112  }113 114  const model = getClusterModel(modelId);115  throw new Error(116    model117      ? `${model.label} is not available on this host yet.`118      : `Unknown model "${modelId}".`,119  );120}121 122async function base64ToVideoFrame(base64: string, mimeType: string): Promise<VideoFrame> {123  const binary = atob(base64);124  const bytes = new Uint8Array(binary.length);125  for (let i = 0; i < binary.length; i++) {126    bytes[i] = binary.charCodeAt(i);127  }128 129  const blob = new Blob([bytes], {type: mimeType});130  const bitmap = await createImageBitmap(blob);131  const frame = new VideoFrame(bitmap, {timestamp: 0});132  bitmap.close();133  return frame;134}135 136async function processDetectTask(task: HostTaskMessage): Promise<void> {137  if (!detector) {138    return;139  }140 141  let frame: VideoFrame | null = null;142 143  try {144    frame = await base64ToVideoFrame(task.image_base64, task.mime_type);145    const results = await detector.detect(frame, {threshold: task.threshold ?? 0.5});146 147    await fetch(`${API_BASE}/api/tasks/${task.id}/complete`, {148      method: 'POST',149      headers: {'Content-Type': 'application/json'},150      body: JSON.stringify({151        status: 'done',152        results: serializeDetectionResults(results),153      }),154    });155  } catch (error) {156    const message = error instanceof Error ? error.message : String(error);157    await fetch(`${API_BASE}/api/tasks/${task.id}/complete`, {158      method: 'POST',159      headers: {'Content-Type': 'application/json'},160      body: JSON.stringify({status: 'error', error: message}),161    });162    throw error;163  } finally {164    frame?.close();165  }166}167 168async function processDescribeTask(task: HostTaskMessage): Promise<void> {169  if (!describer) {170    return;171  }172 173  let frame: VideoFrame | null = null;174 175  try {176    frame = await base64ToVideoFrame(task.image_base64, task.mime_type);177    const description = await describer.describe(frame, {178      instruction: task.instruction ?? 'What do you see?',179      maxNewTokens: task.max_new_tokens ?? 100,180    });181 182    await fetch(`${API_BASE}/api/tasks/${task.id}/complete`, {183      method: 'POST',184      headers: {'Content-Type': 'application/json'},185      body: JSON.stringify({186        status: 'done',187        description,188      }),189    });190  } catch (error) {191    const message = error instanceof Error ? error.message : String(error);192    await fetch(`${API_BASE}/api/tasks/${task.id}/complete`, {193      method: 'POST',194      headers: {'Content-Type': 'application/json'},195      body: JSON.stringify({status: 'error', error: message}),196    });197    throw error;198  } finally {199    frame?.close();200  }201}202 203async function processTask(task: HostTaskMessage): Promise<void> {204  processing = true;205  setStatus(`Processing task ${task.id.slice(0, 8)}…`);206 207  try {208    if (task.kind === 'describe') {209      await processDescribeTask(task);210    } else {211      await processDetectTask(task);212    }213    setStatus(`Task ${task.id.slice(0, 8)} complete. Waiting for requests…`);214  } catch (error) {215    const message = error instanceof Error ? error.message : String(error);216    setStatus(`Task failed: ${message}`);217  } finally {218    processing = false;219    void drainTaskQueue();220  }221}222 223async function drainTaskQueue(): Promise<void> {224  if (processing || taskQueue.length === 0) {225    return;226  }227 228  const task = taskQueue.shift();229  if (task) {230    await processTask(task);231  }232}233 234function enqueueTask(task: HostTaskMessage): void {235  taskQueue.push(task);236  void drainTaskQueue();237}238 239function serializeDetectionResults(results: DetectionResult[]): unknown[] {240  return results.map((result) => ({241    label: result.label,242    score: result.score,243    box: {244      xmin: result.box.xmin,245      ymin: result.box.ymin,246      xmax: result.box.xmax,247      ymax: result.box.ymax,248    },249  }));250}251 252function connectStream(): void {253  eventSource?.close();254 255  const url = `${API_BASE}/api/hosts/stream?host_id=${encodeURIComponent(hostId)}`;256  eventSource = new EventSource(url);257 258  eventSource.addEventListener('ready', () => {259    const model = getClusterModel(selectedModelId);260    const modelLabel = model?.label ?? selectedModelId;261    setStatus(`Hosting "${hostId}" (${modelLabel}). Waiting for requests…`);262  });263 264  eventSource.addEventListener('task', (event) => {265    const raw = JSON.parse(event.data) as HostTaskMessage & {kind?: ClusterTaskKind};266    const task: HostTaskMessage = {267      ...raw,268      kind: raw.kind ?? 'detect',269    };270    enqueueTask(task);271  });272 273  eventSource.onerror = () => {274    if (eventSource?.readyState === EventSource.CLOSED) {275      setStatus('Broker connection lost. Reconnecting in 2s…');276      setTimeout(connectStream, 2000);277    }278  };279}280 281async function registerHost(): Promise<void> {282  hostId = hostIdInput.value.trim();283  selectedModelId = modelSelect.value;284 285  if (!hostId) {286    setStatus('Enter a host id (e.g. my-gpu-node).');287    return;288  }289 290  if (!isImplementedClusterModel(selectedModelId)) {291    setStatus('Choose an available model from the list.');292    return;293  }294 295  registerBtn.disabled = true;296  modelSelect.disabled = true;297  hostIdInput.disabled = true;298 299  try {300    await ensureModelLoaded(selectedModelId);301 302    const res = await fetch(`${API_BASE}/api/hosts/register`, {303      method: 'POST',304      headers: {'Content-Type': 'application/json'},305      body: JSON.stringify({id: hostId, model: selectedModelId}),306    });307 308    if (!res.ok) {309      const json = await res.json().catch(() => ({}));310      throw new Error(json.error ?? `Register failed (${res.status})`);311    }312 313    updateCurlExample();314    connectStream();315  } catch (error) {316    const message = error instanceof Error ? error.message : String(error);317    setStatus(`Failed to start host:\n${message}`);318    registerBtn.disabled = false;319    modelSelect.disabled = false;320    hostIdInput.disabled = false;321  }322}323 324registerBtn.addEventListener('click', () => {325  void registerHost();326});327 328modelSelect.addEventListener('change', () => {329  selectedModelId = modelSelect.value;330  updateCurlExample();331});332 333populateModelSelect();334hostIdInput.value = `node-${Math.random().toString(36).slice(2, 8)}`;335selectedModelId = modelSelect.value;336updateCurlExample();337 338setStatus('Pick a model, enter a host id, and click “Start hosting”.');339