mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
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)
This commit is contained in:
Siddharth Kumar Sah
@@ -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<string, string> = {
|
||||
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,
|
||||
|
||||
@@ -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<string, unknown>) => 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<string>("png");
|
||||
const [quality, setQuality] = useState(95);
|
||||
|
||||
const onChangeRef = useRef(onChange);
|
||||
@@ -75,9 +76,9 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) {
|
||||
/>
|
||||
</div>
|
||||
|
||||
{/* Model */}
|
||||
{/* Quality */}
|
||||
<div>
|
||||
<p className="text-sm font-medium text-muted-foreground mb-1.5">Model</p>
|
||||
<p className="text-sm font-medium text-muted-foreground mb-1.5">Quality</p>
|
||||
<div className="flex gap-1">
|
||||
{MODEL_OPTIONS.map(({ value, label }) => (
|
||||
<button
|
||||
@@ -94,11 +95,6 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) {
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
<p className="text-[11px] text-muted-foreground/70 mt-1">
|
||||
{model === "auto" && "AI when available, falls back to fast resize"}
|
||||
{model === "realesrgan" && "Real-ESRGAN neural network upscaling"}
|
||||
{model === "lanczos" && "Fast Lanczos interpolation resize"}
|
||||
</p>
|
||||
</div>
|
||||
|
||||
{/* Face Enhancement */}
|
||||
@@ -114,10 +110,10 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) {
|
||||
</label>
|
||||
)}
|
||||
|
||||
{/* Denoise */}
|
||||
{/* Noise Reduction */}
|
||||
<div>
|
||||
<div className="flex justify-between items-center">
|
||||
<p className="text-sm font-medium text-muted-foreground">Denoise</p>
|
||||
<p className="text-sm font-medium text-muted-foreground">Noise Reduction</p>
|
||||
<span className="text-sm font-mono font-medium">
|
||||
{denoise === 0 ? "Off" : denoise.toFixed(1)}
|
||||
</span>
|
||||
@@ -131,31 +127,32 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) {
|
||||
onChange={(e) => setDenoise(Number(e.target.value))}
|
||||
className="w-full mt-1"
|
||||
/>
|
||||
<p className="text-[11px] text-muted-foreground/70 mt-1">
|
||||
Smooths out grain and noise. Higher values remove more noise but may soften details.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
{/* Output Format */}
|
||||
<div>
|
||||
<p className="text-sm font-medium text-muted-foreground mb-1.5">Output Format</p>
|
||||
<div className="flex gap-1">
|
||||
{FORMAT_OPTIONS.map((fmt) => (
|
||||
<button
|
||||
key={fmt}
|
||||
type="button"
|
||||
onClick={() => setOutputFormat(fmt)}
|
||||
className={`flex-1 text-xs py-1.5 rounded uppercase ${
|
||||
outputFormat === fmt
|
||||
? "bg-primary text-primary-foreground"
|
||||
: "bg-muted text-muted-foreground"
|
||||
}`}
|
||||
>
|
||||
{fmt}
|
||||
</button>
|
||||
<label htmlFor="upscale-format" className="text-sm font-medium text-muted-foreground">
|
||||
Output Format
|
||||
</label>
|
||||
<select
|
||||
id="upscale-format"
|
||||
value={outputFormat}
|
||||
onChange={(e) => setOutputFormat(e.target.value)}
|
||||
className="w-full mt-1 px-2 py-1.5 rounded border border-border bg-background text-sm text-foreground"
|
||||
>
|
||||
{OUTPUT_FORMATS.map((f) => (
|
||||
<option key={f} value={f}>
|
||||
{f.toUpperCase()}
|
||||
</option>
|
||||
))}
|
||||
</div>
|
||||
</select>
|
||||
</div>
|
||||
|
||||
{/* Quality (JPEG/WebP only) */}
|
||||
{outputFormat !== "png" && (
|
||||
{/* Quality (lossy formats only) */}
|
||||
{LOSSY_FORMATS.includes(outputFormat) && (
|
||||
<div>
|
||||
<div className="flex justify-between items-center">
|
||||
<p className="text-sm font-medium text-muted-foreground">Quality</p>
|
||||
@@ -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<Record<string, unknown>>({});
|
||||
|
||||
// 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 (
|
||||
<div className="space-y-4">
|
||||
@@ -234,13 +205,6 @@ export function UpscaleSettings() {
|
||||
{/* Error */}
|
||||
{error && <p className="text-xs text-red-500">{error}</p>}
|
||||
|
||||
{/* Multi-file progress summary */}
|
||||
{hasMultiple && completedCount > 0 && (
|
||||
<div className="text-xs text-muted-foreground bg-muted/50 rounded-lg px-3 py-2">
|
||||
{completedCount} of {entries.length} images upscaled
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Size info */}
|
||||
{originalSize != null && processedSize != null && (
|
||||
<div className="text-xs text-muted-foreground space-y-0.5">
|
||||
@@ -254,40 +218,26 @@ export function UpscaleSettings() {
|
||||
<ProgressCard
|
||||
active={processing}
|
||||
phase={progress.phase === "idle" ? "uploading" : progress.phase}
|
||||
label={
|
||||
isQueueActive
|
||||
? `Upscaling ${completedCount + 1} of ${entries.length}`
|
||||
: "Upscaling image"
|
||||
}
|
||||
label={hasMultiple ? `Upscaling ${files.length} images` : "Upscaling image"}
|
||||
percent={progress.percent}
|
||||
elapsed={progress.elapsed}
|
||||
/>
|
||||
) : (
|
||||
<div className="space-y-2">
|
||||
<button
|
||||
type="button"
|
||||
data-testid="upscale-submit"
|
||||
onClick={handleProcess}
|
||||
disabled={!hasFile || processing}
|
||||
className="w-full py-2.5 rounded-lg bg-primary text-primary-foreground font-medium disabled:opacity-50 disabled:cursor-not-allowed flex items-center justify-center gap-2"
|
||||
>
|
||||
{`Upscale ${(settings.scale as number) ?? 2}x`}
|
||||
</button>
|
||||
{hasMultiple && !allDone && (
|
||||
<button
|
||||
type="button"
|
||||
onClick={handleProcessAll}
|
||||
disabled={processing}
|
||||
className="w-full py-2 rounded-lg border border-primary text-primary font-medium flex items-center justify-center gap-2 hover:bg-primary/5 disabled:opacity-50"
|
||||
>
|
||||
Upscale All ({pendingCount} remaining)
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
<button
|
||||
type="button"
|
||||
data-testid="upscale-submit"
|
||||
onClick={handleProcess}
|
||||
disabled={!hasFile || processing}
|
||||
className="w-full py-2.5 rounded-lg bg-primary text-primary-foreground font-medium disabled:opacity-50 disabled:cursor-not-allowed flex items-center justify-center gap-2"
|
||||
>
|
||||
{hasMultiple
|
||||
? `Upscale ${(settings.scale as number) ?? 2}x (${files.length} files)`
|
||||
: `Upscale ${(settings.scale as number) ?? 2}x`}
|
||||
</button>
|
||||
)}
|
||||
|
||||
{/* Download */}
|
||||
{downloadUrl && (
|
||||
{/* Download (single file - batch uses Download All ZIP in tool-page) */}
|
||||
{!hasMultiple && downloadUrl && (
|
||||
<a
|
||||
href={downloadUrl}
|
||||
download
|
||||
|
||||
Reference in New Issue
Block a user