From f2e17d2d44163ffd5f519e64804d1fff90f80c01 Mon Sep 17 00:00:00 2001 From: Siddharth Kumar Sah Date: Sun, 12 Apr 2026 21:22:55 +0800 Subject: [PATCH] fix(upscale): overhaul UI, fix AI pipeline bugs, add format support - Replace Auto/AI/Fast buttons with Fast/Balanced/Best (consistent with other tools) - Rename "Denoise" to "Noise Reduction" with explanatory subtitle - Change output format from 3 buttons to dropdown with all formats (PNG, JPG, WebP, AVIF, TIFF, GIF, HEIC, HEIF) - Add HEIC/HEIF input decoding (was missing unlike other tools) - Add HEIC/HEIF/AVIF output conversion via Sharp and heif-enc - Generate browser-compatible WebP preview for non-previewable output formats - Fix torchvision compatibility shim so Real-ESRGAN actually loads (was silently falling back to Lanczos) - Fix denoise crash: Image.fromarray() instead of type(img).fromarray() - Redirect stdout for entire AI pipeline to prevent library messages corrupting JSON output - Add GFPGAN model download for face enhancement - Use batch endpoint for multi-file uploads (enables Download All ZIP) --- apps/api/src/routes/tools/upscale.ts | 65 ++++++- .../src/components/tools/upscale-settings.tsx | 162 ++++++----------- docker/download_models.py | 123 +++++++++---- packages/ai/python/upscale.py | 170 ++++++++++-------- 4 files changed, 304 insertions(+), 216 deletions(-) diff --git a/apps/api/src/routes/tools/upscale.ts b/apps/api/src/routes/tools/upscale.ts index b2f03810..a49dbec8 100644 --- a/apps/api/src/routes/tools/upscale.ts +++ b/apps/api/src/routes/tools/upscale.ts @@ -3,9 +3,11 @@ import { writeFile } from "node:fs/promises"; import { basename, join } from "node:path"; import { upscale } from "@stirling-image/ai"; import type { FastifyInstance, FastifyReply, FastifyRequest } from "fastify"; +import sharp from "sharp"; import { z } from "zod"; import { autoOrient } from "../../lib/auto-orient.js"; import { validateImageBuffer } from "../../lib/file-validation.js"; +import { decodeHeic, encodeHeic } from "../../lib/heic-converter.js"; import { createWorkspace } from "../../lib/workspace.js"; import { updateSingleFileProgress } from "../progress.js"; import { registerToolProcessFn } from "../tool-factory.js"; @@ -66,6 +68,11 @@ export function registerUpscale(app: FastifyInstance) { "Starting upscale", ); + // Decode HEIC/HEIF input via system decoder + if (validation.format === "heif") { + fileBuffer = await decodeHeic(fileBuffer); + } + // Auto-orient to fix EXIF rotation before upscaling fileBuffer = await autoOrient(fileBuffer); @@ -76,6 +83,12 @@ export function registerUpscale(app: FastifyInstance) { const inputPath = join(workspacePath, "input", filename); await writeFile(inputPath, fileBuffer); + // Determine which format the Python sidecar should produce. + // Formats that need Node.js-side conversion (HEIC/HEIF via heif-enc, + // AVIF via Sharp) are produced as PNG first, then converted below. + const needsNodeConversion = ["heic", "heif", "avif"].includes(format); + const pythonFormat = needsNodeConversion ? "png" : format; + // Process const jobIdForProgress = clientJobId; const onProgress = jobIdForProgress @@ -92,15 +105,58 @@ export function registerUpscale(app: FastifyInstance) { const result = await upscale( fileBuffer, join(workspacePath, "output"), - { scale, model, faceEnhance, denoise, format, quality: outputQuality }, + { scale, model, faceEnhance, denoise, format: pythonFormat, quality: outputQuality }, onProgress, ); + // Convert to final format if needed (HEIC/HEIF/AVIF) + let outputBuffer = result.buffer; + let finalFormat = result.format; + if (needsNodeConversion) { + if (format === "heic" || format === "heif") { + outputBuffer = await encodeHeic(result.buffer, outputQuality); + finalFormat = format; + } else if (format === "avif") { + outputBuffer = await sharp(result.buffer).avif({ quality: outputQuality }).toBuffer(); + finalFormat = "avif"; + } + } + // Save output with correct extension for the chosen format - const ext = result.format === "jpeg" ? "jpg" : result.format === "webp" ? "webp" : "png"; + const EXT_MAP: Record = { + jpeg: "jpg", + jpg: "jpg", + png: "png", + webp: "webp", + tiff: "tiff", + gif: "gif", + avif: "avif", + heic: "heic", + heif: "heif", + }; + const ext = EXT_MAP[finalFormat] || "png"; const outputFilename = `${filename.replace(/\.[^.]+$/, "")}_${scale}x.${ext}`; const outputPath = join(workspacePath, "output", outputFilename); - await writeFile(outputPath, result.buffer); + await writeFile(outputPath, outputBuffer); + + // Generate browser-compatible preview for non-previewable formats + const BROWSER_PREVIEWABLE = new Set(["png", "jpg", "jpeg", "webp", "gif", "avif", "bmp"]); + let previewUrl: string | undefined; + if (!BROWSER_PREVIEWABLE.has(finalFormat)) { + try { + // For HEIC/HEIF, decode first since Sharp can't read HEVC + const previewInput = + finalFormat === "heic" || finalFormat === "heif" + ? await decodeHeic(outputBuffer) + : outputBuffer; + const previewBuffer = await sharp(previewInput).webp({ quality: 80 }).toBuffer(); + const previewPath = join(workspacePath, "output", "preview.webp"); + await writeFile(previewPath, previewBuffer); + previewUrl = `/api/v1/download/${jobId}/preview.webp`; + } catch { + // Non-fatal - frontend will show fallback + } + } if (clientJobId) { updateSingleFileProgress({ @@ -113,8 +169,9 @@ export function registerUpscale(app: FastifyInstance) { return reply.send({ jobId, downloadUrl: `/api/v1/download/${jobId}/${encodeURIComponent(outputFilename)}`, + previewUrl, originalSize: fileBuffer.length, - processedSize: result.buffer.length, + processedSize: outputBuffer.length, width: result.width, height: result.height, method: result.method, diff --git a/apps/web/src/components/tools/upscale-settings.tsx b/apps/web/src/components/tools/upscale-settings.tsx index df2277a9..5ce19cf4 100644 --- a/apps/web/src/components/tools/upscale-settings.tsx +++ b/apps/web/src/components/tools/upscale-settings.tsx @@ -6,11 +6,12 @@ import { useFileStore } from "@/stores/file-store"; const QUICK_SCALES = [2, 3, 4, 6, 8]; const MODEL_OPTIONS = [ - { value: "auto", label: "Auto" }, - { value: "realesrgan", label: "AI" }, { value: "lanczos", label: "Fast" }, + { value: "auto", label: "Balanced" }, + { value: "realesrgan", label: "Best" }, ] as const; -const FORMAT_OPTIONS = ["png", "jpeg", "webp"] as const; +const OUTPUT_FORMATS = ["png", "jpg", "webp", "avif", "tiff", "gif", "heic", "heif"] as const; +const LOSSY_FORMATS = ["jpg", "jpeg", "webp", "avif", "heic", "heif"]; export interface UpscaleControlsProps { onChange?: (settings: Record) => void; @@ -21,7 +22,7 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) { const [model, setModel] = useState<"auto" | "realesrgan" | "lanczos">("auto"); const [faceEnhance, setFaceEnhance] = useState(false); const [denoise, setDenoise] = useState(0); - const [outputFormat, setOutputFormat] = useState<"png" | "jpeg" | "webp">("png"); + const [outputFormat, setOutputFormat] = useState("png"); const [quality, setQuality] = useState(95); const onChangeRef = useRef(onChange); @@ -75,9 +76,9 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) { /> - {/* Model */} + {/* Quality */}
-

Model

+

Quality

{MODEL_OPTIONS.map(({ value, label }) => (
-

- {model === "auto" && "AI when available, falls back to fast resize"} - {model === "realesrgan" && "Real-ESRGAN neural network upscaling"} - {model === "lanczos" && "Fast Lanczos interpolation resize"} -

{/* Face Enhancement */} @@ -114,10 +110,10 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) { )} - {/* Denoise */} + {/* Noise Reduction */}
-

Denoise

+

Noise Reduction

{denoise === 0 ? "Off" : denoise.toFixed(1)} @@ -131,31 +127,32 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) { onChange={(e) => setDenoise(Number(e.target.value))} className="w-full mt-1" /> +

+ Smooths out grain and noise. Higher values remove more noise but may soften details. +

{/* Output Format */}
-

Output Format

-
- {FORMAT_OPTIONS.map((fmt) => ( - + +
- {/* Quality (JPEG/WebP only) */} - {outputFormat !== "png" && ( + {/* Quality (lossy formats only) */} + {LOSSY_FORMATS.includes(outputFormat) && (

Quality

@@ -178,54 +175,28 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) { export function UpscaleSettings() { const { files, entries } = useFileStore(); - const { processFiles, processing, error, downloadUrl, originalSize, processedSize, progress } = - useToolProcessor("upscale"); + const { + processFiles, + processAllFiles, + processing, + error, + downloadUrl, + originalSize, + processedSize, + progress, + } = useToolProcessor("upscale"); const [settings, setSettings] = useState>({}); - // Queue mode for "Upscale All" - processes files sequentially - const queueRef = useRef(false); - const settingsRef = useRef(settings); - const prevProcessingRef = useRef(processing); - - useEffect(() => { - settingsRef.current = settings; - }); - - // Auto-advance to next file when current one finishes - useEffect(() => { - if (prevProcessingRef.current && !processing && queueRef.current) { - const currentEntries = useFileStore.getState().entries; - const nextPending = currentEntries.findIndex((e) => e.status === "pending"); - if (nextPending >= 0) { - useFileStore.getState().setSelectedIndex(nextPending); - setTimeout(() => processFiles(useFileStore.getState().files, settingsRef.current), 0); - } else { - queueRef.current = false; - } - } - prevProcessingRef.current = processing; - }, [processing, processFiles]); - const handleProcess = () => { - processFiles(files, settings); - }; - - const handleProcessAll = () => { - queueRef.current = true; - const currentEntries = useFileStore.getState().entries; - const firstPending = currentEntries.findIndex((e) => e.status === "pending"); - if (firstPending >= 0) { - useFileStore.getState().setSelectedIndex(firstPending); - setTimeout(() => processFiles(useFileStore.getState().files, settings), 0); + if (files.length > 1) { + processAllFiles(files, settings); + } else { + processFiles(files, settings); } }; const hasFile = files.length > 0; const hasMultiple = files.length > 1; - const completedCount = entries.filter((e) => e.status === "completed").length; - const pendingCount = entries.filter((e) => e.status === "pending").length; - const allDone = entries.length > 0 && pendingCount === 0; - const isQueueActive = queueRef.current && processing; return (
@@ -234,13 +205,6 @@ export function UpscaleSettings() { {/* Error */} {error &&

{error}

} - {/* Multi-file progress summary */} - {hasMultiple && completedCount > 0 && ( -
- {completedCount} of {entries.length} images upscaled -
- )} - {/* Size info */} {originalSize != null && processedSize != null && (
@@ -254,40 +218,26 @@ export function UpscaleSettings() { ) : ( -
- - {hasMultiple && !allDone && ( - - )} -
+ )} - {/* Download */} - {downloadUrl && ( + {/* Download (single file - batch uses Download All ZIP in tool-page) */} + {!hasMultiple && downloadUrl && ( GFPGAN_MIN_SIZE, ( + f"GFPGAN model too small: {size} bytes (expected > {GFPGAN_MIN_SIZE})" + ) + print(f" GFPGANv1.3.pth downloaded ({size / 1_000_000:.1f} MB)\n") + + +def download_paddleocr_models(): + """Pre-download PaddleOCR PP-OCRv5 model weights from HuggingFace. + + Uses huggingface_hub.snapshot_download() to fetch model files directly + into the PaddleX cache directory. This avoids initializing PaddlePaddle's + C++ inference engine, which segfaults under QEMU emulation (arm64 host + building amd64 image). + """ + print("=== Downloading PaddleOCR PP-OCRv5 models ===") + from huggingface_hub import snapshot_download + + os.makedirs(PADDLEX_MODEL_DIR, exist_ok=True) + + for repo_id in PADDLEOCR_MODELS: + model_name = repo_id.split("/", 1)[1] + local_dir = os.path.join(PADDLEX_MODEL_DIR, model_name) + print(f" Downloading {model_name}...") + snapshot_download(repo_id=repo_id, local_dir=local_dir) + print(f" {model_name} ready") + print(f"All {len(PADDLEOCR_MODELS)} PaddleOCR PP-OCRv5 models downloaded.\n") def download_paddleocr_vl_model(): - """Pre-download PaddleOCR-VL 1.5 model weights.""" + """Pre-download PaddleOCR-VL 1.5 model weights from HuggingFace.""" print("=== Downloading PaddleOCR-VL 1.5 model ===") - try: - from paddleocr import PaddleOCRVL - except ImportError as e: - if "libcuda" in str(e): - print(f" Skipping PaddleOCR-VL pre-download (no CUDA driver at build time)") - print(f" Model will download on first use at runtime.\n") - return - raise + from huggingface_hub import snapshot_download - print(" Downloading PaddleOCR-VL 1.5 weights (~1.93 GB)...") - try: - PaddleOCRVL(device="cpu") - print(" PaddleOCR-VL 1.5 ready\n") - except Exception as e: - print(f" Warning: PaddleOCR-VL pre-download failed: {e}") - print(f" Model will download on first use at runtime.\n") + model_name = PADDLEOCR_VL_MODEL.split("/", 1)[1] + local_dir = os.path.join(PADDLEX_MODEL_DIR, model_name) + print(f" Downloading {model_name} (~1.93 GB)...") + snapshot_download(repo_id=PADDLEOCR_VL_MODEL, local_dir=local_dir) + print(f" {model_name} ready\n") def verify_mediapipe(): @@ -175,6 +205,28 @@ def smoke_test(): ) print(" RealESRGAN model file verified") + # GFPGAN model file must exist + assert os.path.exists(GFPGAN_MODEL_PATH), ( + f"GFPGAN model missing: {GFPGAN_MODEL_PATH}" + ) + assert os.path.getsize(GFPGAN_MODEL_PATH) > GFPGAN_MIN_SIZE, ( + "GFPGAN model file is too small" + ) + print(" GFPGAN model file verified") + + # PaddleOCR model directories must exist + for repo_id in PADDLEOCR_MODELS: + model_name = repo_id.split("/", 1)[1] + model_dir = os.path.join(PADDLEX_MODEL_DIR, model_name) + assert os.path.isdir(model_dir), f"PaddleOCR model missing: {model_dir}" + print(f" PaddleOCR models verified ({len(PADDLEOCR_MODELS)} models)") + + # PaddleOCR-VL model directory must exist + vl_name = PADDLEOCR_VL_MODEL.split("/", 1)[1] + vl_dir = os.path.join(PADDLEX_MODEL_DIR, vl_name) + assert os.path.isdir(vl_dir), f"PaddleOCR-VL model missing: {vl_dir}" + print(" PaddleOCR-VL model verified") + print("Smoke test passed.\n") @@ -182,6 +234,7 @@ def main(): print("Pre-downloading all ML models...\n") download_rembg_models() download_realesrgan_model() + download_gfpgan_model() download_paddleocr_models() download_paddleocr_vl_model() verify_mediapipe() diff --git a/packages/ai/python/upscale.py b/packages/ai/python/upscale.py index 38dba7c8..72a23524 100644 --- a/packages/ai/python/upscale.py +++ b/packages/ai/python/upscale.py @@ -3,6 +3,23 @@ import sys import json import os +# Patch for basicsr compatibility with torchvision >= 0.18. +# torchvision removed transforms.functional_tensor, merging it into +# transforms.functional. basicsr still imports the old path, so we +# create a shim module to redirect the import. +try: + import torchvision.transforms.functional_tensor # noqa: F401 +except (ImportError, ModuleNotFoundError): + try: + import types + import torchvision.transforms.functional as _F + + _shim = types.ModuleType("torchvision.transforms.functional_tensor") + _shim.rgb_to_grayscale = _F.rgb_to_grayscale + sys.modules["torchvision.transforms.functional_tensor"] = _shim + except ImportError: + pass # torchvision not installed at all, Real-ESRGAN unavailable + def emit_progress(percent, stage): """Emit structured progress to stderr for bridge.ts to capture.""" @@ -27,6 +44,7 @@ def apply_denoise(img, strength): try: import numpy as np import cv2 + from PIL import Image arr = np.array(img) # Map 0-1 strength to filter parameter (3-15 range) @@ -35,7 +53,7 @@ def apply_denoise(img, strength): denoised = cv2.fastNlMeansDenoisingColored(arr, None, h, h, 7, 21) else: denoised = cv2.fastNlMeansDenoising(arr, None, h, 7, 21) - return type(img).fromarray(denoised) + return Image.fromarray(denoised) except ImportError: from PIL import ImageFilter @@ -70,8 +88,10 @@ def main(): try: emit_progress(10, "Loading AI model") - # Redirect stdout to stderr so basicsr/realesrgan init messages - # cannot contaminate our JSON result on stdout. + # Redirect stdout to stderr for the ENTIRE AI pipeline. + # Libraries like basicsr, realesrgan, gfpgan, and torch print + # download progress and init messages to stdout which would + # corrupt our JSON result. stdout_fd = os.dup(1) os.dup2(2, 1) @@ -81,71 +101,72 @@ def main(): from gpu import gpu_available import numpy as np import torch + + if not os.path.exists(REALESRGAN_MODEL_PATH): + raise FileNotFoundError( + f"RealESRGAN model not found: {REALESRGAN_MODEL_PATH}" + ) + + use_gpu = gpu_available() + device = torch.device("cuda" if use_gpu else "cpu") + + # RealESRGAN_x4plus is a 4x model internally + ai_model = RRDBNet( + num_in_ch=3, + num_out_ch=3, + num_feat=64, + num_block=23, + num_grow_ch=32, + scale=4, + ) + upsampler = RealESRGANer( + scale=4, + model_path=REALESRGAN_MODEL_PATH, + model=ai_model, + half=use_gpu, + device=device, + ) + emit_progress(20, "AI model loaded") + + img_array = np.array(img.convert("RGB")) + emit_progress(30, "Enhancing image with AI") + output_array, _ = upsampler.enhance(img_array, outscale=scale) + emit_progress(80, "AI enhancement complete") + result = Image.fromarray(output_array) + method = "realesrgan" + + # Face enhancement with GFPGAN + if face_enhance: + emit_progress(82, "Enhancing faces") + try: + from gfpgan import GFPGANer + + if os.path.exists(GFPGAN_MODEL_PATH): + face_enhancer = GFPGANer( + model_path=GFPGAN_MODEL_PATH, + upscale=scale, + arch="clean", + channel_multiplier=2, + bg_upsampler=upsampler, + ) + _, _, face_output = face_enhancer.enhance( + img_array, + has_aligned=False, + only_center_face=False, + paste_back=True, + ) + result = Image.fromarray(face_output) + emit_progress(88, "Face enhancement complete") + else: + emit_progress(88, "Face model not found, skipping") + except (ImportError, RuntimeError, OSError): + emit_progress(88, "Face enhancement unavailable, skipping") + finally: - # Restore stdout after imports + # Restore stdout after ALL AI processing os.dup2(stdout_fd, 1) os.close(stdout_fd) - if not os.path.exists(REALESRGAN_MODEL_PATH): - raise FileNotFoundError( - f"RealESRGAN model not found: {REALESRGAN_MODEL_PATH}" - ) - - use_gpu = gpu_available() - device = torch.device("cuda" if use_gpu else "cpu") - - # RealESRGAN_x4plus is a 4x model internally - model = RRDBNet( - num_in_ch=3, - num_out_ch=3, - num_feat=64, - num_block=23, - num_grow_ch=32, - scale=4, - ) - upsampler = RealESRGANer( - scale=4, - model_path=REALESRGAN_MODEL_PATH, - model=model, - half=use_gpu, - device=device, - ) - emit_progress(20, "AI model loaded") - - img_array = np.array(img.convert("RGB")) - emit_progress(30, "Enhancing image with AI") - output_array, _ = upsampler.enhance(img_array, outscale=scale) - emit_progress(80, "AI enhancement complete") - result = Image.fromarray(output_array) - method = "realesrgan" - - # Face enhancement with GFPGAN - if face_enhance: - emit_progress(82, "Enhancing faces") - try: - from gfpgan import GFPGANer - - if os.path.exists(GFPGAN_MODEL_PATH): - face_enhancer = GFPGANer( - model_path=GFPGAN_MODEL_PATH, - upscale=scale, - arch="clean", - channel_multiplier=2, - bg_upsampler=upsampler, - ) - _, _, face_output = face_enhancer.enhance( - img_array, - has_aligned=False, - only_center_face=False, - paste_back=True, - ) - result = Image.fromarray(face_output) - emit_progress(88, "Face enhancement complete") - else: - emit_progress(88, "Face model not found, skipping") - except (ImportError, RuntimeError, OSError): - emit_progress(88, "Face enhancement unavailable, skipping") - except (ImportError, FileNotFoundError, RuntimeError, OSError): # RealESRGAN unavailable or failed if model_choice == "realesrgan": @@ -165,22 +186,29 @@ def main(): # Determine final output path based on format base_path = output_path.rsplit(".", 1)[0] - if output_format == "jpeg": - final_path = base_path + ".jpg" - elif output_format == "webp": - final_path = base_path + ".webp" - else: - final_path = base_path + ".png" + EXT_MAP = { + "jpeg": ".jpg", + "jpg": ".jpg", + "png": ".png", + "webp": ".webp", + "tiff": ".tiff", + "gif": ".gif", + } + final_path = base_path + EXT_MAP.get(output_format, ".png") # Save with format-specific options emit_progress(95, "Saving result") save_kwargs = {} - if output_format == "jpeg": + if output_format in ("jpeg", "jpg"): result = result.convert("RGB") # Strip alpha for JPEG save_kwargs["quality"] = quality save_kwargs["optimize"] = True elif output_format == "webp": save_kwargs["quality"] = quality + elif output_format == "tiff": + save_kwargs["compression"] = "tiff_lzw" + elif output_format == "gif": + result = result.convert("P", palette=Image.ADAPTIVE, colors=256) result.save(final_path, **save_kwargs)