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:
@@ -3,9 +3,11 @@ import { writeFile } from "node:fs/promises";
|
|||||||
import { basename, join } from "node:path";
|
import { basename, join } from "node:path";
|
||||||
import { upscale } from "@stirling-image/ai";
|
import { upscale } from "@stirling-image/ai";
|
||||||
import type { FastifyInstance, FastifyReply, FastifyRequest } from "fastify";
|
import type { FastifyInstance, FastifyReply, FastifyRequest } from "fastify";
|
||||||
|
import sharp from "sharp";
|
||||||
import { z } from "zod";
|
import { z } from "zod";
|
||||||
import { autoOrient } from "../../lib/auto-orient.js";
|
import { autoOrient } from "../../lib/auto-orient.js";
|
||||||
import { validateImageBuffer } from "../../lib/file-validation.js";
|
import { validateImageBuffer } from "../../lib/file-validation.js";
|
||||||
|
import { decodeHeic, encodeHeic } from "../../lib/heic-converter.js";
|
||||||
import { createWorkspace } from "../../lib/workspace.js";
|
import { createWorkspace } from "../../lib/workspace.js";
|
||||||
import { updateSingleFileProgress } from "../progress.js";
|
import { updateSingleFileProgress } from "../progress.js";
|
||||||
import { registerToolProcessFn } from "../tool-factory.js";
|
import { registerToolProcessFn } from "../tool-factory.js";
|
||||||
@@ -66,6 +68,11 @@ export function registerUpscale(app: FastifyInstance) {
|
|||||||
"Starting upscale",
|
"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
|
// Auto-orient to fix EXIF rotation before upscaling
|
||||||
fileBuffer = await autoOrient(fileBuffer);
|
fileBuffer = await autoOrient(fileBuffer);
|
||||||
|
|
||||||
@@ -76,6 +83,12 @@ export function registerUpscale(app: FastifyInstance) {
|
|||||||
const inputPath = join(workspacePath, "input", filename);
|
const inputPath = join(workspacePath, "input", filename);
|
||||||
await writeFile(inputPath, fileBuffer);
|
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
|
// Process
|
||||||
const jobIdForProgress = clientJobId;
|
const jobIdForProgress = clientJobId;
|
||||||
const onProgress = jobIdForProgress
|
const onProgress = jobIdForProgress
|
||||||
@@ -92,15 +105,58 @@ export function registerUpscale(app: FastifyInstance) {
|
|||||||
const result = await upscale(
|
const result = await upscale(
|
||||||
fileBuffer,
|
fileBuffer,
|
||||||
join(workspacePath, "output"),
|
join(workspacePath, "output"),
|
||||||
{ scale, model, faceEnhance, denoise, format, quality: outputQuality },
|
{ scale, model, faceEnhance, denoise, format: pythonFormat, quality: outputQuality },
|
||||||
onProgress,
|
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
|
// 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 outputFilename = `${filename.replace(/\.[^.]+$/, "")}_${scale}x.${ext}`;
|
||||||
const outputPath = join(workspacePath, "output", outputFilename);
|
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) {
|
if (clientJobId) {
|
||||||
updateSingleFileProgress({
|
updateSingleFileProgress({
|
||||||
@@ -113,8 +169,9 @@ export function registerUpscale(app: FastifyInstance) {
|
|||||||
return reply.send({
|
return reply.send({
|
||||||
jobId,
|
jobId,
|
||||||
downloadUrl: `/api/v1/download/${jobId}/${encodeURIComponent(outputFilename)}`,
|
downloadUrl: `/api/v1/download/${jobId}/${encodeURIComponent(outputFilename)}`,
|
||||||
|
previewUrl,
|
||||||
originalSize: fileBuffer.length,
|
originalSize: fileBuffer.length,
|
||||||
processedSize: result.buffer.length,
|
processedSize: outputBuffer.length,
|
||||||
width: result.width,
|
width: result.width,
|
||||||
height: result.height,
|
height: result.height,
|
||||||
method: result.method,
|
method: result.method,
|
||||||
|
|||||||
@@ -6,11 +6,12 @@ import { useFileStore } from "@/stores/file-store";
|
|||||||
|
|
||||||
const QUICK_SCALES = [2, 3, 4, 6, 8];
|
const QUICK_SCALES = [2, 3, 4, 6, 8];
|
||||||
const MODEL_OPTIONS = [
|
const MODEL_OPTIONS = [
|
||||||
{ value: "auto", label: "Auto" },
|
|
||||||
{ value: "realesrgan", label: "AI" },
|
|
||||||
{ value: "lanczos", label: "Fast" },
|
{ value: "lanczos", label: "Fast" },
|
||||||
|
{ value: "auto", label: "Balanced" },
|
||||||
|
{ value: "realesrgan", label: "Best" },
|
||||||
] as const;
|
] 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 {
|
export interface UpscaleControlsProps {
|
||||||
onChange?: (settings: Record<string, unknown>) => void;
|
onChange?: (settings: Record<string, unknown>) => void;
|
||||||
@@ -21,7 +22,7 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) {
|
|||||||
const [model, setModel] = useState<"auto" | "realesrgan" | "lanczos">("auto");
|
const [model, setModel] = useState<"auto" | "realesrgan" | "lanczos">("auto");
|
||||||
const [faceEnhance, setFaceEnhance] = useState(false);
|
const [faceEnhance, setFaceEnhance] = useState(false);
|
||||||
const [denoise, setDenoise] = useState(0);
|
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 [quality, setQuality] = useState(95);
|
||||||
|
|
||||||
const onChangeRef = useRef(onChange);
|
const onChangeRef = useRef(onChange);
|
||||||
@@ -75,9 +76,9 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) {
|
|||||||
/>
|
/>
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
{/* Model */}
|
{/* Quality */}
|
||||||
<div>
|
<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">
|
<div className="flex gap-1">
|
||||||
{MODEL_OPTIONS.map(({ value, label }) => (
|
{MODEL_OPTIONS.map(({ value, label }) => (
|
||||||
<button
|
<button
|
||||||
@@ -94,11 +95,6 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) {
|
|||||||
</button>
|
</button>
|
||||||
))}
|
))}
|
||||||
</div>
|
</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>
|
</div>
|
||||||
|
|
||||||
{/* Face Enhancement */}
|
{/* Face Enhancement */}
|
||||||
@@ -114,10 +110,10 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) {
|
|||||||
</label>
|
</label>
|
||||||
)}
|
)}
|
||||||
|
|
||||||
{/* Denoise */}
|
{/* Noise Reduction */}
|
||||||
<div>
|
<div>
|
||||||
<div className="flex justify-between items-center">
|
<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">
|
<span className="text-sm font-mono font-medium">
|
||||||
{denoise === 0 ? "Off" : denoise.toFixed(1)}
|
{denoise === 0 ? "Off" : denoise.toFixed(1)}
|
||||||
</span>
|
</span>
|
||||||
@@ -131,31 +127,32 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) {
|
|||||||
onChange={(e) => setDenoise(Number(e.target.value))}
|
onChange={(e) => setDenoise(Number(e.target.value))}
|
||||||
className="w-full mt-1"
|
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>
|
</div>
|
||||||
|
|
||||||
{/* Output Format */}
|
{/* Output Format */}
|
||||||
<div>
|
<div>
|
||||||
<p className="text-sm font-medium text-muted-foreground mb-1.5">Output Format</p>
|
<label htmlFor="upscale-format" className="text-sm font-medium text-muted-foreground">
|
||||||
<div className="flex gap-1">
|
Output Format
|
||||||
{FORMAT_OPTIONS.map((fmt) => (
|
</label>
|
||||||
<button
|
<select
|
||||||
key={fmt}
|
id="upscale-format"
|
||||||
type="button"
|
value={outputFormat}
|
||||||
onClick={() => setOutputFormat(fmt)}
|
onChange={(e) => setOutputFormat(e.target.value)}
|
||||||
className={`flex-1 text-xs py-1.5 rounded uppercase ${
|
className="w-full mt-1 px-2 py-1.5 rounded border border-border bg-background text-sm text-foreground"
|
||||||
outputFormat === fmt
|
>
|
||||||
? "bg-primary text-primary-foreground"
|
{OUTPUT_FORMATS.map((f) => (
|
||||||
: "bg-muted text-muted-foreground"
|
<option key={f} value={f}>
|
||||||
}`}
|
{f.toUpperCase()}
|
||||||
>
|
</option>
|
||||||
{fmt}
|
|
||||||
</button>
|
|
||||||
))}
|
))}
|
||||||
</div>
|
</select>
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
{/* Quality (JPEG/WebP only) */}
|
{/* Quality (lossy formats only) */}
|
||||||
{outputFormat !== "png" && (
|
{LOSSY_FORMATS.includes(outputFormat) && (
|
||||||
<div>
|
<div>
|
||||||
<div className="flex justify-between items-center">
|
<div className="flex justify-between items-center">
|
||||||
<p className="text-sm font-medium text-muted-foreground">Quality</p>
|
<p className="text-sm font-medium text-muted-foreground">Quality</p>
|
||||||
@@ -178,54 +175,28 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) {
|
|||||||
|
|
||||||
export function UpscaleSettings() {
|
export function UpscaleSettings() {
|
||||||
const { files, entries } = useFileStore();
|
const { files, entries } = useFileStore();
|
||||||
const { processFiles, processing, error, downloadUrl, originalSize, processedSize, progress } =
|
const {
|
||||||
useToolProcessor("upscale");
|
processFiles,
|
||||||
|
processAllFiles,
|
||||||
|
processing,
|
||||||
|
error,
|
||||||
|
downloadUrl,
|
||||||
|
originalSize,
|
||||||
|
processedSize,
|
||||||
|
progress,
|
||||||
|
} = useToolProcessor("upscale");
|
||||||
const [settings, setSettings] = useState<Record<string, unknown>>({});
|
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 = () => {
|
const handleProcess = () => {
|
||||||
processFiles(files, settings);
|
if (files.length > 1) {
|
||||||
};
|
processAllFiles(files, settings);
|
||||||
|
} else {
|
||||||
const handleProcessAll = () => {
|
processFiles(files, settings);
|
||||||
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);
|
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
const hasFile = files.length > 0;
|
const hasFile = files.length > 0;
|
||||||
const hasMultiple = files.length > 1;
|
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 (
|
return (
|
||||||
<div className="space-y-4">
|
<div className="space-y-4">
|
||||||
@@ -234,13 +205,6 @@ export function UpscaleSettings() {
|
|||||||
{/* Error */}
|
{/* Error */}
|
||||||
{error && <p className="text-xs text-red-500">{error}</p>}
|
{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 */}
|
{/* Size info */}
|
||||||
{originalSize != null && processedSize != null && (
|
{originalSize != null && processedSize != null && (
|
||||||
<div className="text-xs text-muted-foreground space-y-0.5">
|
<div className="text-xs text-muted-foreground space-y-0.5">
|
||||||
@@ -254,40 +218,26 @@ export function UpscaleSettings() {
|
|||||||
<ProgressCard
|
<ProgressCard
|
||||||
active={processing}
|
active={processing}
|
||||||
phase={progress.phase === "idle" ? "uploading" : progress.phase}
|
phase={progress.phase === "idle" ? "uploading" : progress.phase}
|
||||||
label={
|
label={hasMultiple ? `Upscaling ${files.length} images` : "Upscaling image"}
|
||||||
isQueueActive
|
|
||||||
? `Upscaling ${completedCount + 1} of ${entries.length}`
|
|
||||||
: "Upscaling image"
|
|
||||||
}
|
|
||||||
percent={progress.percent}
|
percent={progress.percent}
|
||||||
elapsed={progress.elapsed}
|
elapsed={progress.elapsed}
|
||||||
/>
|
/>
|
||||||
) : (
|
) : (
|
||||||
<div className="space-y-2">
|
<button
|
||||||
<button
|
type="button"
|
||||||
type="button"
|
data-testid="upscale-submit"
|
||||||
data-testid="upscale-submit"
|
onClick={handleProcess}
|
||||||
onClick={handleProcess}
|
disabled={!hasFile || processing}
|
||||||
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"
|
||||||
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`}
|
? `Upscale ${(settings.scale as number) ?? 2}x (${files.length} files)`
|
||||||
</button>
|
: `Upscale ${(settings.scale as number) ?? 2}x`}
|
||||||
{hasMultiple && !allDone && (
|
</button>
|
||||||
<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>
|
|
||||||
)}
|
)}
|
||||||
|
|
||||||
{/* Download */}
|
{/* Download (single file - batch uses Download All ZIP in tool-page) */}
|
||||||
{downloadUrl && (
|
{!hasMultiple && downloadUrl && (
|
||||||
<a
|
<a
|
||||||
href={downloadUrl}
|
href={downloadUrl}
|
||||||
download
|
download
|
||||||
|
|||||||
+88
-35
@@ -20,6 +20,13 @@ REALESRGAN_MODEL_URL = (
|
|||||||
REALESRGAN_MODEL_PATH = os.path.join(REALESRGAN_MODEL_DIR, "RealESRGAN_x4plus.pth")
|
REALESRGAN_MODEL_PATH = os.path.join(REALESRGAN_MODEL_DIR, "RealESRGAN_x4plus.pth")
|
||||||
REALESRGAN_MIN_SIZE = 60_000_000 # ~67 MB
|
REALESRGAN_MIN_SIZE = 60_000_000 # ~67 MB
|
||||||
|
|
||||||
|
GFPGAN_MODEL_DIR = "/opt/models/gfpgan"
|
||||||
|
GFPGAN_MODEL_URL = (
|
||||||
|
"https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth"
|
||||||
|
)
|
||||||
|
GFPGAN_MODEL_PATH = os.path.join(GFPGAN_MODEL_DIR, "GFPGANv1.3.pth")
|
||||||
|
GFPGAN_MIN_SIZE = 300_000_000 # ~332 MB
|
||||||
|
|
||||||
REMBG_MODELS = [
|
REMBG_MODELS = [
|
||||||
"u2net",
|
"u2net",
|
||||||
"isnet-general-use",
|
"isnet-general-use",
|
||||||
@@ -30,9 +37,24 @@ REMBG_MODELS = [
|
|||||||
"birefnet-matting",
|
"birefnet-matting",
|
||||||
]
|
]
|
||||||
|
|
||||||
# PaddleOCR language codes (not ISO). German/French/Spanish use "latin" model.
|
# PaddleOCR PP-OCRv5 HuggingFace model repos to pre-download.
|
||||||
# Valid keys: ch, en, korean, japan, chinese_cht, ta, te, ka, latin, arabic, cyrillic, devanagari
|
# These are the models used by PaddleOCR(ocr_version="PP-OCRv5").
|
||||||
PADDLEOCR_LANGUAGES = ["en", "ch", "japan", "korean", "latin"]
|
# Downloaded via huggingface_hub to avoid initializing the PaddlePaddle
|
||||||
|
# inference engine, which segfaults under QEMU emulation at build time.
|
||||||
|
PADDLEOCR_MODELS = [
|
||||||
|
"PaddlePaddle/PP-OCRv5_server_det",
|
||||||
|
"PaddlePaddle/PP-OCRv5_server_rec",
|
||||||
|
"PaddlePaddle/PP-OCRv5_mobile_det",
|
||||||
|
"PaddlePaddle/PP-OCRv5_mobile_rec",
|
||||||
|
"PaddlePaddle/latin_PP-OCRv5_mobile_rec",
|
||||||
|
"PaddlePaddle/korean_PP-OCRv5_mobile_rec",
|
||||||
|
"PaddlePaddle/PP-LCNet_x1_0_textline_ori",
|
||||||
|
]
|
||||||
|
|
||||||
|
PADDLEOCR_VL_MODEL = "PaddlePaddle/PaddleOCR-VL-1.5"
|
||||||
|
|
||||||
|
# PaddleX stores models here by default
|
||||||
|
PADDLEX_MODEL_DIR = os.path.expanduser("~/.paddlex/official_models")
|
||||||
|
|
||||||
|
|
||||||
def _register_birefnet_matting():
|
def _register_birefnet_matting():
|
||||||
@@ -90,44 +112,52 @@ def download_realesrgan_model():
|
|||||||
print(f" RealESRGAN_x4plus.pth downloaded ({size / 1_000_000:.1f} MB)\n")
|
print(f" RealESRGAN_x4plus.pth downloaded ({size / 1_000_000:.1f} MB)\n")
|
||||||
|
|
||||||
|
|
||||||
def download_paddleocr_models():
|
def download_gfpgan_model():
|
||||||
"""Pre-download PaddleOCR PP-OCRv5 models for all supported languages."""
|
"""Download GFPGANv1.3.pth pretrained weights for face enhancement."""
|
||||||
print("=== Downloading PaddleOCR PP-OCRv5 models ===")
|
print("=== Downloading GFPGAN model ===")
|
||||||
try:
|
os.makedirs(GFPGAN_MODEL_DIR, exist_ok=True)
|
||||||
from paddleocr import PaddleOCR
|
print(f" Downloading from {GFPGAN_MODEL_URL}...")
|
||||||
except ImportError as e:
|
urllib.request.urlretrieve(GFPGAN_MODEL_URL, GFPGAN_MODEL_PATH)
|
||||||
if "libcuda" in str(e):
|
|
||||||
print(f" Skipping PaddleOCR model pre-download (no CUDA driver at build time)")
|
|
||||||
print(f" Models will download on first use at runtime.\n")
|
|
||||||
return
|
|
||||||
raise
|
|
||||||
|
|
||||||
for lang in PADDLEOCR_LANGUAGES:
|
size = os.path.getsize(GFPGAN_MODEL_PATH)
|
||||||
print(f" Downloading PP-OCRv5 models for lang={lang}...")
|
assert size > GFPGAN_MIN_SIZE, (
|
||||||
PaddleOCR(lang=lang, use_gpu=False, show_log=False, ocr_version="PP-OCRv5")
|
f"GFPGAN model too small: {size} bytes (expected > {GFPGAN_MIN_SIZE})"
|
||||||
print(f" {lang} ready")
|
)
|
||||||
print(f"All {len(PADDLEOCR_LANGUAGES)} PaddleOCR PP-OCRv5 languages downloaded.\n")
|
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():
|
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 ===")
|
print("=== Downloading PaddleOCR-VL 1.5 model ===")
|
||||||
try:
|
from huggingface_hub import snapshot_download
|
||||||
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
|
|
||||||
|
|
||||||
print(" Downloading PaddleOCR-VL 1.5 weights (~1.93 GB)...")
|
model_name = PADDLEOCR_VL_MODEL.split("/", 1)[1]
|
||||||
try:
|
local_dir = os.path.join(PADDLEX_MODEL_DIR, model_name)
|
||||||
PaddleOCRVL(device="cpu")
|
print(f" Downloading {model_name} (~1.93 GB)...")
|
||||||
print(" PaddleOCR-VL 1.5 ready\n")
|
snapshot_download(repo_id=PADDLEOCR_VL_MODEL, local_dir=local_dir)
|
||||||
except Exception as e:
|
print(f" {model_name} ready\n")
|
||||||
print(f" Warning: PaddleOCR-VL pre-download failed: {e}")
|
|
||||||
print(f" Model will download on first use at runtime.\n")
|
|
||||||
|
|
||||||
|
|
||||||
def verify_mediapipe():
|
def verify_mediapipe():
|
||||||
@@ -175,6 +205,28 @@ def smoke_test():
|
|||||||
)
|
)
|
||||||
print(" RealESRGAN model file verified")
|
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")
|
print("Smoke test passed.\n")
|
||||||
|
|
||||||
|
|
||||||
@@ -182,6 +234,7 @@ def main():
|
|||||||
print("Pre-downloading all ML models...\n")
|
print("Pre-downloading all ML models...\n")
|
||||||
download_rembg_models()
|
download_rembg_models()
|
||||||
download_realesrgan_model()
|
download_realesrgan_model()
|
||||||
|
download_gfpgan_model()
|
||||||
download_paddleocr_models()
|
download_paddleocr_models()
|
||||||
download_paddleocr_vl_model()
|
download_paddleocr_vl_model()
|
||||||
verify_mediapipe()
|
verify_mediapipe()
|
||||||
|
|||||||
@@ -3,6 +3,23 @@ import sys
|
|||||||
import json
|
import json
|
||||||
import os
|
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):
|
def emit_progress(percent, stage):
|
||||||
"""Emit structured progress to stderr for bridge.ts to capture."""
|
"""Emit structured progress to stderr for bridge.ts to capture."""
|
||||||
@@ -27,6 +44,7 @@ def apply_denoise(img, strength):
|
|||||||
try:
|
try:
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import cv2
|
import cv2
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
arr = np.array(img)
|
arr = np.array(img)
|
||||||
# Map 0-1 strength to filter parameter (3-15 range)
|
# 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)
|
denoised = cv2.fastNlMeansDenoisingColored(arr, None, h, h, 7, 21)
|
||||||
else:
|
else:
|
||||||
denoised = cv2.fastNlMeansDenoising(arr, None, h, 7, 21)
|
denoised = cv2.fastNlMeansDenoising(arr, None, h, 7, 21)
|
||||||
return type(img).fromarray(denoised)
|
return Image.fromarray(denoised)
|
||||||
except ImportError:
|
except ImportError:
|
||||||
from PIL import ImageFilter
|
from PIL import ImageFilter
|
||||||
|
|
||||||
@@ -70,8 +88,10 @@ def main():
|
|||||||
try:
|
try:
|
||||||
emit_progress(10, "Loading AI model")
|
emit_progress(10, "Loading AI model")
|
||||||
|
|
||||||
# Redirect stdout to stderr so basicsr/realesrgan init messages
|
# Redirect stdout to stderr for the ENTIRE AI pipeline.
|
||||||
# cannot contaminate our JSON result on stdout.
|
# 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)
|
stdout_fd = os.dup(1)
|
||||||
os.dup2(2, 1)
|
os.dup2(2, 1)
|
||||||
|
|
||||||
@@ -81,71 +101,72 @@ def main():
|
|||||||
from gpu import gpu_available
|
from gpu import gpu_available
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import torch
|
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:
|
finally:
|
||||||
# Restore stdout after imports
|
# Restore stdout after ALL AI processing
|
||||||
os.dup2(stdout_fd, 1)
|
os.dup2(stdout_fd, 1)
|
||||||
os.close(stdout_fd)
|
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):
|
except (ImportError, FileNotFoundError, RuntimeError, OSError):
|
||||||
# RealESRGAN unavailable or failed
|
# RealESRGAN unavailable or failed
|
||||||
if model_choice == "realesrgan":
|
if model_choice == "realesrgan":
|
||||||
@@ -165,22 +186,29 @@ def main():
|
|||||||
|
|
||||||
# Determine final output path based on format
|
# Determine final output path based on format
|
||||||
base_path = output_path.rsplit(".", 1)[0]
|
base_path = output_path.rsplit(".", 1)[0]
|
||||||
if output_format == "jpeg":
|
EXT_MAP = {
|
||||||
final_path = base_path + ".jpg"
|
"jpeg": ".jpg",
|
||||||
elif output_format == "webp":
|
"jpg": ".jpg",
|
||||||
final_path = base_path + ".webp"
|
"png": ".png",
|
||||||
else:
|
"webp": ".webp",
|
||||||
final_path = base_path + ".png"
|
"tiff": ".tiff",
|
||||||
|
"gif": ".gif",
|
||||||
|
}
|
||||||
|
final_path = base_path + EXT_MAP.get(output_format, ".png")
|
||||||
|
|
||||||
# Save with format-specific options
|
# Save with format-specific options
|
||||||
emit_progress(95, "Saving result")
|
emit_progress(95, "Saving result")
|
||||||
save_kwargs = {}
|
save_kwargs = {}
|
||||||
if output_format == "jpeg":
|
if output_format in ("jpeg", "jpg"):
|
||||||
result = result.convert("RGB") # Strip alpha for JPEG
|
result = result.convert("RGB") # Strip alpha for JPEG
|
||||||
save_kwargs["quality"] = quality
|
save_kwargs["quality"] = quality
|
||||||
save_kwargs["optimize"] = True
|
save_kwargs["optimize"] = True
|
||||||
elif output_format == "webp":
|
elif output_format == "webp":
|
||||||
save_kwargs["quality"] = quality
|
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)
|
result.save(final_path, **save_kwargs)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user