feat: overhaul upscale with bug fixes and advanced features

- Fix multi-image: process selected file, not always first
- Fix progress bar: asymptotic fill prevents visual stalling
- Fix slider: write results to captured index, not current selection
- Add model selection (Auto/AI/Fast), face enhancement, denoise
- Add output format (PNG/JPEG/WebP) with quality control
- Add Upscale All for sequential batch processing with queue
- More granular Python progress stages for smoother UX
This commit is contained in:
Siddharth Kumar Sah
2026-04-12 19:04:23 +08:00
parent 7995d13c81
commit fe376aebd2
4 changed files with 367 additions and 71 deletions
+10 -4
View File
@@ -56,8 +56,13 @@ export function registerUpscale(app: FastifyInstance) {
try {
const settings = settingsRaw ? JSON.parse(settingsRaw) : {};
const scale = Number(settings.scale) || 2;
const model = settings.model || "auto";
const faceEnhance = Boolean(settings.faceEnhance);
const denoise = Number(settings.denoise) || 0;
const format = settings.format || "png";
const outputQuality = Number(settings.quality) || 95;
request.log.info(
{ toolId: "upscale", imageSize: fileBuffer.length, scale },
{ toolId: "upscale", imageSize: fileBuffer.length, scale, model, format },
"Starting upscale",
);
@@ -87,12 +92,13 @@ export function registerUpscale(app: FastifyInstance) {
const result = await upscale(
fileBuffer,
join(workspacePath, "output"),
{ scale },
{ scale, model, faceEnhance, denoise, format, quality: outputQuality },
onProgress,
);
// Save output
const outputFilename = `${filename.replace(/\.[^.]+$/, "")}_${scale}x.png`;
// Save output with correct extension for the chosen format
const ext = result.format === "jpeg" ? "jpg" : result.format === "webp" ? "webp" : "png";
const outputFilename = `${filename.replace(/\.[^.]+$/, "")}_${scale}x.${ext}`;
const outputPath = join(workspacePath, "output", outputFilename);
await writeFile(outputPath, result.buffer);
@@ -5,6 +5,12 @@ import { useToolProcessor } from "@/hooks/use-tool-processor";
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" },
] as const;
const FORMAT_OPTIONS = ["png", "jpeg", "webp"] as const;
export interface UpscaleControlsProps {
onChange?: (settings: Record<string, unknown>) => void;
@@ -12,6 +18,11 @@ export interface UpscaleControlsProps {
export function UpscaleControls({ onChange }: UpscaleControlsProps) {
const [scale, setScale] = useState(2);
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 [quality, setQuality] = useState(95);
const onChangeRef = useRef(onChange);
useEffect(() => {
@@ -19,8 +30,15 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) {
});
useEffect(() => {
onChangeRef.current?.({ scale });
}, [scale]);
onChangeRef.current?.({
scale,
model,
faceEnhance,
denoise,
format: outputFormat,
quality,
});
}, [scale, model, faceEnhance, denoise, outputFormat, quality]);
return (
<div className="space-y-4">
@@ -56,21 +74,158 @@ export function UpscaleControls({ onChange }: UpscaleControlsProps) {
className="w-full mt-2"
/>
</div>
{/* Model */}
<div>
<p className="text-sm font-medium text-muted-foreground mb-1.5">Model</p>
<div className="flex gap-1">
{MODEL_OPTIONS.map(({ value, label }) => (
<button
key={value}
type="button"
onClick={() => setModel(value)}
className={`flex-1 text-xs py-1.5 rounded ${
model === value
? "bg-primary text-primary-foreground"
: "bg-muted text-muted-foreground"
}`}
>
{label}
</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 */}
{model !== "lanczos" && (
<label className="flex items-center gap-2 cursor-pointer">
<input
type="checkbox"
checked={faceEnhance}
onChange={(e) => setFaceEnhance(e.target.checked)}
className="rounded border-border"
/>
<span className="text-sm text-foreground">Enhance faces</span>
</label>
)}
{/* Denoise */}
<div>
<div className="flex justify-between items-center">
<p className="text-sm font-medium text-muted-foreground">Denoise</p>
<span className="text-sm font-mono font-medium">
{denoise === 0 ? "Off" : denoise.toFixed(1)}
</span>
</div>
<input
type="range"
min={0}
max={1}
step={0.1}
value={denoise}
onChange={(e) => setDenoise(Number(e.target.value))}
className="w-full mt-1"
/>
</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>
))}
</div>
</div>
{/* Quality (JPEG/WebP only) */}
{outputFormat !== "png" && (
<div>
<div className="flex justify-between items-center">
<p className="text-sm font-medium text-muted-foreground">Quality</p>
<span className="text-sm font-mono font-medium">{quality}</span>
</div>
<input
type="range"
min={1}
max={100}
step={1}
value={quality}
onChange={(e) => setQuality(Number(e.target.value))}
className="w-full mt-1"
/>
</div>
)}
</div>
);
}
export function UpscaleSettings() {
const { files } = useFileStore();
const { files, entries } = useFileStore();
const { processFiles, 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);
}
};
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">
@@ -79,6 +234,13 @@ 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">
@@ -87,25 +249,41 @@ export function UpscaleSettings() {
</div>
)}
{/* Process button */}
{/* Process buttons / progress */}
{processing ? (
<ProgressCard
active={processing}
phase={progress.phase === "idle" ? "uploading" : progress.phase}
label="Upscaling image"
label={
isQueueActive
? `Upscaling ${completedCount + 1} of ${entries.length}`
: "Upscaling image"
}
percent={progress.percent}
elapsed={progress.elapsed}
/>
) : (
<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>
<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>
)}
{/* Download */}
+155 -52
View File
@@ -14,6 +14,34 @@ REALESRGAN_MODEL_PATH = os.environ.get(
"/opt/models/realesrgan/RealESRGAN_x4plus.pth",
)
GFPGAN_MODEL_PATH = os.environ.get(
"GFPGAN_MODEL_PATH",
"/opt/models/gfpgan/GFPGANv1.3.pth",
)
def apply_denoise(img, strength):
"""Apply denoising to a PIL image. Uses OpenCV when available, falls back to PIL."""
if strength <= 0:
return img
try:
import numpy as np
import cv2
arr = np.array(img)
# Map 0-1 strength to filter parameter (3-15 range)
h = int(3 + strength * 12)
if len(arr.shape) == 3 and arr.shape[2] >= 3:
denoised = cv2.fastNlMeansDenoisingColored(arr, None, h, h, 7, 21)
else:
denoised = cv2.fastNlMeansDenoising(arr, None, h, 7, 21)
return type(img).fromarray(denoised)
except ImportError:
from PIL import ImageFilter
radius = max(0.5, strength * 1.5)
return img.filter(ImageFilter.GaussianBlur(radius=radius))
def main():
input_path = sys.argv[1]
@@ -21,79 +49,154 @@ def main():
settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
scale = settings.get("scale", 2)
model_choice = settings.get("model", "auto")
face_enhance = settings.get("faceEnhance", False)
denoise_strength = float(settings.get("denoise", 0))
output_format = settings.get("format", "png")
quality = int(settings.get("quality", 95))
try:
emit_progress(10, "Loading upscale model")
emit_progress(5, "Opening image")
from PIL import Image
img = Image.open(input_path)
new_size = (img.width * scale, img.height * scale)
# Try Real-ESRGAN first
try:
# Redirect stdout to stderr so basicsr/realesrgan init messages
# cannot contaminate our JSON result on stdout.
stdout_fd = os.dup(1)
os.dup2(2, 1)
method = "lanczos"
result = None
# Try Real-ESRGAN if requested
if model_choice in ("auto", "realesrgan"):
try:
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from gpu import gpu_available
import numpy as np
import torch
finally:
# Restore stdout after imports
os.dup2(stdout_fd, 1)
os.close(stdout_fd)
emit_progress(10, "Loading AI model")
if not os.path.exists(REALESRGAN_MODEL_PATH):
raise FileNotFoundError(f"RealESRGAN model not found: {REALESRGAN_MODEL_PATH}")
# Redirect stdout to stderr so basicsr/realesrgan init messages
# cannot contaminate our JSON result on stdout.
stdout_fd = os.dup(1)
os.dup2(2, 1)
use_gpu = gpu_available()
device = torch.device("cuda" if use_gpu else "cpu")
try:
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from gpu import gpu_available
import numpy as np
import torch
finally:
# Restore stdout after imports
os.dup2(stdout_fd, 1)
os.close(stdout_fd)
# 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, "Model ready")
img_array = np.array(img.convert("RGB"))
emit_progress(25, "Upscaling image")
output, _ = upsampler.enhance(img_array, outscale=scale)
emit_progress(90, "Upscaling complete")
result = Image.fromarray(output)
emit_progress(95, "Saving result")
result.save(output_path)
method = "realesrgan"
except (ImportError, FileNotFoundError, RuntimeError, OSError):
# RealESRGAN unavailable or failed - fall back to Lanczos
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":
emit_progress(15, "AI model not available, using fast resize")
result = None
# Fall back to Lanczos
if result is None:
emit_progress(50, "Upscaling with Lanczos")
img_upscaled = img.resize(new_size, Image.LANCZOS)
emit_progress(95, "Saving result")
img_upscaled.save(output_path)
result = img.resize(new_size, Image.LANCZOS)
method = "lanczos"
# Denoise
if denoise_strength > 0:
emit_progress(90, "Reducing noise")
result = apply_denoise(result, denoise_strength)
# 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"
# Save with format-specific options
emit_progress(95, "Saving result")
save_kwargs = {}
if output_format == "jpeg":
result = result.convert("RGB") # Strip alpha for JPEG
save_kwargs["quality"] = quality
save_kwargs["optimize"] = True
elif output_format == "webp":
save_kwargs["quality"] = quality
result.save(final_path, **save_kwargs)
# Get actual dimensions of the saved result
actual_w, actual_h = result.size
print(
json.dumps(
{
"success": True,
"scale": scale,
"width": new_size[0],
"height": new_size[1],
"width": actual_w,
"height": actual_h,
"method": method,
"output_path": final_path,
"format": output_format,
}
)
)
+10 -1
View File
@@ -4,6 +4,11 @@ import { type ProgressCallback, runPythonWithProgress } from "./bridge.js";
export interface UpscaleOptions {
scale?: number;
model?: string;
faceEnhance?: boolean;
denoise?: number;
format?: string;
quality?: number;
}
export interface UpscaleResult {
@@ -11,6 +16,7 @@ export interface UpscaleResult {
width: number;
height: number;
method: string;
format: string;
}
export async function upscale(
@@ -34,11 +40,14 @@ export async function upscale(
throw new Error(result.error || "Upscaling failed");
}
const buffer = await readFile(outputPath);
// Python may write to a different path when the output format changes
const actualOutputPath = result.output_path || outputPath;
const buffer = await readFile(actualOutputPath);
return {
buffer,
width: result.width,
height: result.height,
method: result.method ?? "unknown",
format: result.format ?? "png",
};
}