AI Upscaler por dentro: dos redes neuronales en una extensión de Chrome

Un amigo que trabajaba en un ecommerce tenía que preparar cientos de fotos de artículos al día: ampliarlas y quitarles el fondo. Las herramientas con IA que lo hacen cobran por imagen o por suscripción y procesan la foto en su servidor, y a mano es inviable a ese volumen. Yo quería que lo hiciera su propio ordenador, gratis y sin que la imagen saliera del equipo. De ahí salió AI Upscaler & Background Remover, y este artículo cuenta cómo está hecho por dentro. Si buscas qué hace y cómo se usa, eso está en el artículo para usuarios.

La regla que lo condiciona todo: ninguna imagen sale del dispositivo. No hay backend ni API, y los modelos van empaquetados dentro de la extensión. Si algo no se puede hacer en el cliente, no se hace.

La pieza central: ONNX Runtime Web

Hay dos modelos, los dos en formato ONNX:

• Real-ESRGAN general x4v3 para ampliar. Es la variante compacta de Real-ESRGAN: unos 4,9 MB entre el grafo y sus pesos.
• U²-Netp para quitar fondos. Es la versión ligera de U²-Net, de unos 4,6 MB.

Los ejecuta ONNX Runtime Web, que puede tirar de dos motores: WebGPU, si el equipo tiene una GPU disponible, o WebAssembly con SIMD e hilos si no la tiene. La sesión se crea primero con WebGPU y, si falla, se repite con WASM. Así funciona en cualquier equipo, y en los que tienen GPU va mucho más rápido.

try {
  return await ort.InferenceSession.create(model, { ...options, executionProviders: ['webgpu'] })
}
catch {
  return await ort.InferenceSession.create(model, { ...options, executionProviders: ['wasm'] })
}

Crear una sesión es caro: hay que leer el modelo, compilarlo y subir los pesos a la GPU. Por eso cada worker guarda la promesa de la sesión, no la sesión, y la reutiliza en cada imagen. Si la creación falla, la promesa se borra para que el siguiente intento empiece de cero en vez de devolver el mismo error para siempre.

Los modelos y el runtime van dentro de la extensión y se cargan con chrome.runtime.getURL(). El runtime de WebAssembly pesa más que los dos modelos juntos (unos 24 MB), y es la mayor parte de lo que ocupa la extensión.

Todo fuera del hilo principal

Una inferencia puede tardar varios segundos, y la interfaz no puede congelarse mientras tanto. Cada modelo tiene su propio Web Worker (ai-worker.ts y bg-worker.ts) y la página solo habla con ellos por mensajes:

Algunos detalles que importan:

• La imagen viaja como un ImageBitmap transferido, no copiado: pasa al worker sin duplicar memoria.
• Todos los mensajes llevan un requestId. Si eliges otra imagen a mitad de un proceso, la respuesta del anterior llega con un id viejo y se descarta. Sin esto, el resultado de una foto podría aparecer sobre otra.
• Cancelar es otro mensaje con el mismo id. El worker lo apunta en un Set y lo comprueba entre mosaico y mosaico, así que se detiene en el siguiente sin esperar al final.
• Con el progreso, la página calcula un tiempo estimado (tiempo por mosaico × mosaicos que faltan) y lo pone también en el título de la pestaña.
• El worker de quitar fondos se crea solo la primera vez que se usa. Quien solo amplía imágenes no carga ese modelo.

Ampliar sin quedarse sin memoria

Real-ESRGAN multiplica por 16 los píxeles de la imagen. Pasar una foto grande entera de una vez se come la memoria de la GPU y, en el mejor caso, tarda una eternidad. La solución clásica es trocearla en mosaicos (tiles), ampliar cada uno y recomponer.

El problema de trocear es que el modelo ve el contexto de alrededor de cada píxel. En el borde de un mosaico ese contexto no existe, y el resultado se nota: puede aparecer una cuadrícula de costuras en la imagen final. En las primeras versiones el margen estaba a 0, y activarlo fue de lo primero del roadmap.

Lo que hace ahora:

1. Divide la imagen en mosaicos de 112 px de lado.
2. Recorta cada uno con 8 px de margen extra por cada lado que tenga vecino. Así el modelo recibe una pieza de 128 px.
3. Si el mosaico es del borde y no llega a 128, lo completa repitiendo los píxeles del borde hasta el múltiplo de 128 que espera el modelo.
4. Al recomponer, descarta el margen ampliado y copia solo la parte central. El contexto se usa para calcular, pero nunca acaba en la imagen.

const TILE_SIZE = 112 // TILE_SIZE + 2*TILE_PADDING must equal MODEL_INPUT_MULTIPLE
const TILE_PADDING = 8
const MODEL_INPUT_MULTIPLE = 128

El cálculo de las regiones (getTileRegions) y el relleno de bordes (padImageDataToMultiple) son funciones puras con sus tests: comprobar que cada píxel de destino se cubre exactamente una vez es mucho más fácil sin modelo de por medio.

La salida no siempre es la misma

Al leer el tensor de salida hay dos sorpresas posibles. Puede venir en float16, y JavaScript no tiene un tipo para eso. El worker lo decodifica a mano (signo, exponente y mantisa). Además, el rango puede ser 0–1 o 0–255, así que el worker mira el valor máximo y decide si multiplicar por 255 antes de convertir a píxeles RGBA.

×2 sale de ×4

Solo hay un modelo empaquetado, el de ×4. Para ×2, la extensión amplía ×4 y luego reduce a la mitad con el suavizado de alta calidad del canvas. Así no hay que llevar dos modelos dentro de la extensión, y el ×2 tiene el detalle reconstruido de la IA en vez de una ampliación bicúbica.

Tiene un coste que tuve que corregir en la propia interfaz: al principio la opción ×2 decía «más rápido», y es justo al revés, porque hace el mismo trabajo que ×4 y uno más. Ahora dice lo que es cierto: archivo más pequeño.

Un límite para no romper el navegador

Una imagen de 8000 × 8000 ampliada ×4 son mil millones de píxeles, y eso ya no cabe en la memoria de una pestaña. Antes de enviar nada al worker, la página lee las dimensiones y rechaza lo que pase de 32 megapíxeles con un mensaje claro, en vez de dejar que el worker se quede sin memoria en silencio.

Quitar fondos: un modelo que no arrancaba

U²-Net funciona de otra manera. No amplía: devuelve una máscara que dice, para cada píxel, cuánto pertenece al sujeto. El proceso es:

1. Reducir la imagen a 320 × 320, que es lo que espera el modelo.
2. Normalizarla con las medias y desviaciones de ImageNet, que son con las que se entrenó.
3. Ejecutar el modelo y normalizar la máscara a 0–1 (mínimo y máximo).
4. Escalar la máscara al tamaño original con suavizado de alta calidad, y usarla como canal alfa de la imagen original.

El resultado se exporta siempre en PNG, porque es el formato que conserva la transparencia. Como la máscara se calcula a 320 px y se escala después, los píxeles del resultado son los de la foto original y no se pierde resolución. Los bordes de detalle muy fino, como el pelo, sí quedan algo suaves.

El modelo, tal como se distribuye, no cargaba en ONNX Runtime Web: varias capas MaxPool usan ceil_mode=1, que el runtime web no soporta. La solución está en un script de Python que se ejecuta al descargar el modelo y cambia esas capas a ceil_mode=0.

for node in model.graph.node:
    if node.op_type == "MaxPool":
        for attr in node.attribute:
            if attr.name == "ceil_mode" and attr.i == 1:
                attr.i = 0

No es un truco a ciegas. ceil_mode solo cambia el resultado cuando la dimensión que se divide es impar. La entrada es siempre de 320 × 320, y al ir dividiendo entre dos (160, 80, 40, 20, 10) todas las dimensiones son pares. Redondear hacia arriba o hacia abajo da exactamente lo mismo, así que el parche no cuesta calidad.

Lo que pone difícil una extensión Manifest V3

Ejecutar IA en una web normal ya tiene lo suyo. Dentro de una extensión de Chrome hay obstáculos extra:

• WebAssembly y CSP. Las páginas de una extensión MV3 tienen una política de seguridad estricta. Para compilar WASM hace falta declarar 'wasm-unsafe-eval' en script-src, y nada más: no se permite ningún script externo.
• Los ficheros del runtime. ONNX Runtime necesita saber dónde están sus .wasm. Van empaquetados en ort-wasm/, pero el código no los importa con una ruta como /ort-wasm/...: la ruta se resuelve en tiempo de ejecución a partir del origen del worker, y así funciona igual en desarrollo y en producción.
• Workers en desarrollo. En modo desarrollo, Vite sirve los workers desde localhost, y Chrome cierra las páginas de una extensión que intentan crear un worker de otro origen. La configuración de WXT compila copias de los workers dentro de la propia extensión y las recompila al guardar. Cuando un rebuild borra esa carpeta, el watcher se reinicia solo.
• El menú contextual. El service worker de fondo crea la opción «Mejorar imagen con IA». Al pulsarla, guarda la URL de la imagen en storage.session y abre la página principal, que la descarga al arrancar (o la decodifica directamente si es una URL data:). Por eso la extensión pide permiso sobre las webs: sin él, muchas bloquearían la descarga.
• Copiar al portapapeles. La API del portapapeles solo acepta imágenes PNG, así que si el resultado es WebP o JPEG se convierte a PNG antes de copiarlo.

La interfaz principal es una página de opciones abierta en pestaña completa, no un popup: un popup se cierra al perder el foco y se llevaría por delante una inferencia a medias.

Lotes: cientos de fotos sin tumbar la pestaña

La versión 1.0.8 añadió lo que más pedía el caso de uso original: soltar muchas imágenes, procesarlas todas con los mismos ajustes y descargarlas en un ZIP. Antes de escribir código dejé por escrito qué sería un fracaso aunque funcionara: que el flujo de una sola imagen cambiara, que un lote grande tumbara la pestaña por memoria (con ×4 ya había pasado en desarrollo) o que una imagen fallida parara el resto.

La cola son funciones puras. El estado de cada imagen (pending, processing, done o error) y las transiciones viven en utils/batch-queue.ts, sin Vue: elegir la siguiente, reintentar las fallidas, devolver a la espera lo que estaba a medias al cancelar, decidir qué queda seleccionado al quitar una... Todo se prueba sin montar la interfaz.

Secuencial, con los workers de siempre. El lote no cambia el protocolo de los workers: es un bucle que coge la siguiente imagen en espera, le manda su petición con su requestId y espera el resultado. Como el bucle vuelve a buscar en cada vuelta, las imágenes que se añaden a mitad del lote se procesan solas. No proceso en paralelo a propósito: la GPU es una sola, así que dos inferencias ×4 a la vez ganarían poco y duplicarían la memoria.

Un ZIP que nunca mezcla ajustes. Los ajustes se fijan al pulsar el botón. Si después cambias la escala, el formato o el modo y ya hay resultados, se pide confirmación y se descartan. Así no hace falta un estado de «resultado desactualizado», y el ZIP siempre es coherente.

Un presupuesto de memoria. Antes de cada imagen se estima el peor caso de su resultado, el bitmap RGBA de salida:

export function estimateResultBytes(width: number, height: number, scale: 1 | 2 | 4): number {
  return width * height * scale * scale * 4
}

// A lone image always fits, so the single-image flow keeps working for huge outputs
export function fitsMemoryBudget(heldBytes: number, estimatedBytes: number, budget = BATCH_MEMORY_BUDGET_BYTES): boolean {
  return heldBytes === 0 || heldBytes + estimatedBytes <= budget
}

Si sumada a lo que ya ocupan los resultados pasa de 1,5 GB, esa imagen se queda en rojo con un aviso y la cola sigue con las demás, que quizá sí caben. La excepción del primer caso es deliberada: una imagen sola siempre se intenta, para no romper el flujo normal con fotos grandes.

El ZIP, sin comprimir. Los resultados ya son PNG, WebP o JPEG comprimidos, así que JSZip los guarda con STORE: comprimirlos otra vez solo gastaría tiempo y memoria. Los nombres dentro del ZIP necesitaron una corrección: dos fotos Foto.png y foto.png son archivos distintos en Linux, pero al descomprimir en Windows o macOS una pisa a la otra. Ahora los nombres se aplanan (sin barras) y se desduplican sin distinguir mayúsculas, añadiendo (2), (3)...

Cómo está organizado

Es una extensión hecha con WXT, Vue 3, TypeScript y Tailwind CSS v4. En ejecución solo depende de tres paquetes: Vue, ONNX Runtime Web y JSZip. La lógica que no necesita navegador vive en utils/ como funciones puras: cálculo de mosaicos, relleno, dimensiones, nombres de archivo, normalización de la máscara, ETA, la cola del lote... Cada una tiene sus tests con Vitest, y los componentes con interacción (la zona de soltar, el comparador, el selector de modo, la tira del lote, las acciones del resultado) los tienen con Testing Library. En total son 118 tests en 20 archivos.

Los textos están en inglés, español y chino. Todo lo visible pasa por una función t(clave, texto por defecto), y un test comprueba que los tres catálogos tienen exactamente las mismas claves. Así no se escapa una traducción a medias.

GitHub Actions ejecuta lint, comprobación de tipos y tests en cada push. El trabajo está planificado en un roadmap por fases (correcciones críticas, calidad del resultado, experiencia de uso, calidad del código y publicación), y las funcionalidades principales tienen su plan escrito antes de empezar a programarlas.

Qué haría distinto

El ×2 hace más trabajo que el ×4. Funciona y la calidad es buena, pero es tiempo que se tira. Hoy buscaría un modelo ×2 igual de compacto para empaquetarlo también, aunque eso haga crecer la extensión.

La máscara de fondo se calcula a 320 px aunque la foto tenga 4000. Para un sujeto bien definido sobra, pero en pelo y bordes finos se nota. El siguiente paso sería un modelo de segmentación más grande, o refinar el borde de la máscara usando la imagen original a resolución completa.

Y el worker de ampliar es el archivo más crítico y el que menos tests tiene. Sus funciones puras ya están separadas y probadas, pero todavía falta extraer la conversión de tensores y la recomposición de mosaicos, y un test de integración con una sesión ONNX simulada. Está apuntado en el roadmap, junto con la versión para Firefox.

Y en los lotes, el presupuesto de memoria compara una estimación del bitmap con el tamaño de los archivos ya comprimidos. Es una cota conservadora que funciona, pero medir la memoria real ajustaría mejor cuántas imágenes caben antes de pedir que descargues el ZIP.

Puedes instalar AI Upscaler desde la Chrome Web Store.