behind_glass/src/depth/aiDepth.worker.ts

import { pipeline, type DepthEstimationOutput as DepthEstimationPipelineOutput } from '@huggingface/transformers';
import { AI_DEPTH, AI_DEPTH_MODELS, type AiDepthModel } from '../constants/depth';

type Device = 'webgpu' | 'wasm';
type RequestMessage = { type: 'infer'; id: number; image: ImageBitmap; model: AiDepthModel };
type WorkerReply =
  | { type: 'status'; id: number; message: string }
  | { type: 'result'; id: number; width: number; height: number; values: ArrayBuffer; normalizedValues: ArrayBuffer }
  | { type: 'error'; id: number; message: string };
type Estimator = {
  (image: OffscreenCanvas): Promise<DepthEstimationPipelineOutput | DepthEstimationPipelineOutput[]>;
  processor: { image_processor: { size: { width: number; height: number } } };
  dispose(): Promise<void>;
};
type Progress = { status: string; progress?: number };
type PipelineFactory = (task: string, model: string, options: {
  device: Device;
  progress_callback: (progress: Progress) => void;
}) => Promise<Estimator>;
const createPipeline = pipeline as unknown as PipelineFactory;
const workerScope = self as unknown as {
  postMessage(message: WorkerReply, transfer?: Transferable[]): void;
  addEventListener(type: 'message', listener: (event: MessageEvent<RequestMessage>) => void): void;
};

let estimator: Estimator | null = null;
let activeDevice: Device | null = null;
let activeModel: AiDepthModel | null = null;

function reply(message: WorkerReply, transfer: Transferable[] = []): void {
  workerScope.postMessage(message, transfer);
}

function setStatus(id: number, message: string): void {
  reply({ type: 'status', id, message });
}

async function loadEstimator(device: Device, id: number, model: AiDepthModel): Promise<Estimator> {
  setStatus(id, 'AI 모델 준비 중…');
  const loaded = await createPipeline('depth-estimation', AI_DEPTH_MODELS[model].id, {
    device,
    progress_callback: (progress) => {
      if (progress.status === 'progress') {
        setStatus(id, `AI 모델 다운로드 중 · ${Math.round(progress.progress ?? 0)}%`);
      } else if (progress.status === 'initiate' || progress.status === 'download') {
        setStatus(id, 'AI 모델 다운로드 중…');
      }
    },
  });
  loaded.processor.image_processor.size = { width: AI_DEPTH.inputSize, height: AI_DEPTH.inputSize };
  return loaded;
}

async function ensureEstimator(id: number, model: AiDepthModel): Promise<Estimator> {
  if (estimator && activeModel !== model) {
    await estimator.dispose();
    estimator = null;
  }
  if (estimator) return estimator;
  activeModel = model;
  if ('gpu' in navigator) {
    try {
      estimator = await loadEstimator('webgpu', id, model);
      activeDevice = 'webgpu';
      return estimator;
    } catch {
      if (estimator) await estimator.dispose();
      estimator = null;
      setStatus(id, 'AI 모델을 준비하는 중입니다…');
    }
  }
  estimator = await loadEstimator('wasm', id, model);
  activeDevice = 'wasm';
  return estimator;
}

function normalizeDepth(output: DepthEstimationPipelineOutput): { values: Uint8Array; normalizedValues: Float32Array; width: number; height: number } {
  const { predicted_depth: prediction } = output;
  const [height, width] = prediction.dims.slice(-2);
  const data = prediction.data;
  if (!width || !height || data.length !== width * height) throw new Error('모델이 유효한 depth 크기를 반환하지 않았습니다.');

  let min = Infinity;
  let max = -Infinity;
  for (let i = 0; i < data.length; i += 1) {
    const value = Number(data[i]);
    if (!Number.isFinite(value)) continue;
    min = Math.min(min, value);
    max = Math.max(max, value);
  }
  if (!Number.isFinite(min) || !Number.isFinite(max)) throw new Error('모델의 depth 결과가 비어 있습니다.');

  const range = max - min;
  const values = new Uint8Array(width * height);
  const normalizedValues = new Float32Array(width * height);
  if (range <= 1e-8) {
    values.fill(128);
    normalizedValues.fill(0.5);
  } else {
    for (let i = 0; i < values.length; i += 1) {
      const value = Number(data[i]);
      // Depth Anything V2 predicts relative inverse depth: larger values are nearer.
      const normalized = Number.isFinite(value) ? (value - min) / range : 0.5;
      normalizedValues[i] = normalized;
      values[i] = Math.round(normalized * 255);
    }
  }
  return { values, normalizedValues, width, height };
}

async function infer({ id, image, model: modelKey }: RequestMessage): Promise<void> {
  try {
    const canvas = new OffscreenCanvas(image.width, image.height);
    const context = canvas.getContext('2d');
    if (!context) throw new Error('사진을 처리할 수 없습니다.');
    context.drawImage(image, 0, 0);
    image.close();

    let model = await ensureEstimator(id, modelKey);
    setStatus(id, 'AI 깊이를 계산하는 중입니다…');
    let output: DepthEstimationPipelineOutput;
    try {
      const predicted = await model(canvas);
      if (Array.isArray(predicted)) throw new Error('단일 사진 결과를 읽을 수 없습니다.');
      output = predicted;
    } catch (error) {
      if (activeDevice !== 'webgpu') throw error;
      await model.dispose();
      estimator = null;
      activeDevice = null;
      setStatus(id, 'AI 추론을 다시 시도하는 중입니다…');
      model = await loadEstimator('wasm', id, modelKey);
      estimator = model;
      activeDevice = 'wasm';
      const predicted = await model(canvas);
      if (Array.isArray(predicted)) throw new Error('단일 사진 결과를 읽을 수 없습니다.');
      output = predicted;
    }

    try {
      const normalized = normalizeDepth(output);
      const values = normalized.values.buffer as ArrayBuffer;
      const normalizedValues = normalized.normalizedValues.buffer as ArrayBuffer;
      reply({ type: 'result', id, width: normalized.width, height: normalized.height, values, normalizedValues }, [values, normalizedValues]);
    } finally {
      output.predicted_depth.dispose();
    }
  } catch (error) {
    reply({ type: 'error', id, message: error instanceof Error ? error.message : 'AI depth 생성에 실패했습니다.' });
  }
}

workerScope.addEventListener('message', (event: MessageEvent<RequestMessage>) => {
  if (event.data.type === 'infer') void infer(event.data);
});