Merge pull request #48 from stirling-image/fix/upscale-bugs-and-features

feat: overhaul upscale with bug fixes and advanced features

apps/api/src/routes/tools/upscale.ts

@@ -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
+      // Save output with correct extension for the chosen format
-      const outputFilename = `${filename.replace(/\.[^.]+$/, "")}_${scale}x.png`;
+      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);
 

apps/web/src/components/tools/upscale-settings.tsx

@@ -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 });
+    onChangeRef.current?.({
-  }, [scale]);
+      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,16 +249,21 @@ 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}
         />
       ) : (
+        <div className="space-y-2">
           <button
             type="button"
             data-testid="upscale-submit"
@@ -106,6 +273,17 @@ export function UpscaleSettings() {
           >
             {`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 */}

apps/web/src/hooks/use-tool-processor.ts

@@ -35,18 +35,8 @@ const AI_PYTHON_TOOLS = new Set<string>(PYTHON_SIDECAR_TOOLS);
 const MEDIUM_TOOLS = new Set(["content-aware-resize", "convert"]);
 
 export function useToolProcessor(toolId: string) {
-  const {
+  const { processing, error, processedUrl, originalSize, processedSize, setProcessing, setError } =
-    processing,
+    useFileStore();
-    error,
-    processedUrl,
-    originalSize,
-    processedSize,
-    setProcessing,
-    setError,
-    setProcessedUrl,
-    setSizes,
-    setJobId,
-  } = useFileStore();
 
   const [progress, setProgress] = useState<ToolProgress>(IDLE_PROGRESS);
   const elapsedRef = useRef<ReturnType<typeof setInterval> | null>(null);
@@ -74,8 +64,19 @@ export function useToolProcessor(toolId: string) {
         return;
       }
 
+      // Capture the file index at request time so results are written
+      // to the correct entry even if the user navigates away.
+      const capturedIndex = useFileStore.getState().selectedIndex;
+
       setError(null);
-      setProcessedUrl(null);
+      // Mark the target entry as processing and clear any old result
+      useFileStore.getState().updateEntry(capturedIndex, {
+        processedUrl: null,
+        processedPreviewUrl: null,
+        processedFilename: null,
+        status: "processing",
+        error: null,
+      });
       setProcessing(true);
       setProgress({ phase: "uploading", percent: 0, elapsed: 0 });
 
@@ -106,7 +107,7 @@ export function useToolProcessor(toolId: string) {
                 setProgress((prev) => ({
                   ...prev,
                   phase: "processing",
-                  percent: scaled,
+                  percent: Math.max(prev.percent, scaled),
                   stage: data.stage,
                 }));
               }
@@ -130,7 +131,7 @@ export function useToolProcessor(toolId: string) {
       delete cleanSettings._bgImageFile;
 
       const formData = new FormData();
-      formData.append("file", files[0]);
+      formData.append("file", files[capturedIndex] ?? files[0]);
       formData.append("settings", JSON.stringify(cleanSettings));
       if (bgImageFile) {
         formData.append("backgroundImage", bgImageFile);
@@ -140,9 +141,9 @@ export function useToolProcessor(toolId: string) {
       }
 
       // If this file came from the Files page, include its ID for version tracking
-      const currentEntry = useFileStore.getState().currentEntry;
+      const capturedEntry = useFileStore.getState().entries[capturedIndex];
-      if (currentEntry?.serverFileId) {
+      if (capturedEntry?.serverFileId) {
-        formData.append("fileId", currentEntry.serverFileId);
+        formData.append("fileId", capturedEntry.serverFileId);
       }
 
       // Use XHR for upload progress tracking
@@ -188,6 +189,20 @@ export function useToolProcessor(toolId: string) {
             });
           }, 500);
         }
+
+        // AI tools: asymptotic fill during long processing gaps.
+        // Slowly creeps toward 88% so the bar never stalls visually.
+        // Real SSE events always win via Math.max in the handler.
+        if (isAiTool) {
+          processingTimerRef.current = setInterval(() => {
+            setProgress((prev) => {
+              if (prev.phase !== "processing") return prev;
+              const remaining = 88 - prev.percent;
+              if (remaining <= 0.5) return prev;
+              return { ...prev, percent: prev.percent + remaining * 0.015 };
+            });
+          }, 1000);
+        }
       };
 
       xhr.onload = () => {
@@ -201,16 +216,17 @@ export function useToolProcessor(toolId: string) {
         if (xhr.status >= 200 && xhr.status < 300) {
           try {
             const result: ProcessResult = JSON.parse(xhr.responseText);
-            setJobId(result.jobId);
+            // Write result to the entry that was being processed (captured at
-            setProcessedUrl(result.downloadUrl, result.previewUrl);
+            // request time), not whatever entry happens to be selected now.
-            setSizes(result.originalSize, result.processedSize);
+            useFileStore.getState().updateEntry(capturedIndex, {
-            // Update serverFileId if a new version was saved
+              processedUrl: result.downloadUrl,
-            if (result.savedFileId) {
+              processedPreviewUrl: result.previewUrl ?? null,
-              const state = useFileStore.getState();
+              processedFilename: null,
-              if (state.entries[state.selectedIndex]) {
+              status: "completed",
-                state.updateEntry(state.selectedIndex, { serverFileId: result.savedFileId });
+              originalSize: result.originalSize,
-              }
+              processedSize: result.processedSize,
-            }
+              ...(result.savedFileId ? { serverFileId: result.savedFileId } : {}),
+            });
           } catch {
             setError("Invalid response from server");
           }
@@ -260,7 +276,7 @@ export function useToolProcessor(toolId: string) {
       });
       xhr.send(formData);
     },
-    [toolId, isAiTool, isMediumTool, setProcessing, setError, setProcessedUrl, setSizes, setJobId],
+    [toolId, isAiTool, isMediumTool, setProcessing, setError],
   );
 
   const processAllFiles = useCallback(

packages/ai/python/upscale.py

@@ -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,16 +49,27 @@ 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
+        method = "lanczos"
+        result = None
+
+        # Try Real-ESRGAN if requested
+        if model_choice in ("auto", "realesrgan"):
             try:
+                emit_progress(10, "Loading AI model")
+
                 # Redirect stdout to stderr so basicsr/realesrgan init messages
                 # cannot contaminate our JSON result on stdout.
                 stdout_fd = os.dup(1)
@@ -48,7 +87,9 @@ def main():
                     os.close(stdout_fd)
 
                 if not os.path.exists(REALESRGAN_MODEL_PATH):
-                raise FileNotFoundError(f"RealESRGAN model not found: {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")
@@ -69,31 +110,93 @@ def main():
                     half=use_gpu,
                     device=device,
                 )
-            emit_progress(20, "Model ready")
+                emit_progress(20, "AI model loaded")
+
                 img_array = np.array(img.convert("RGB"))
-            emit_progress(25, "Upscaling image")
+                emit_progress(30, "Enhancing image with AI")
-            output, _ = upsampler.enhance(img_array, outscale=scale)
+                output_array, _ = upsampler.enhance(img_array, outscale=scale)
-            emit_progress(90, "Upscaling complete")
+                emit_progress(80, "AI enhancement complete")
-            result = Image.fromarray(output)
+                result = Image.fromarray(output_array)
-            emit_progress(95, "Saving result")
-            result.save(output_path)
                 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 - fall back to Lanczos
+                # 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)
+            result = img.resize(new_size, Image.LANCZOS)
-            emit_progress(95, "Saving result")
-            img_upscaled.save(output_path)
             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],
+                    "width": actual_w,
-                    "height": new_size[1],
+                    "height": actual_h,
                     "method": method,
+                    "output_path": final_path,
+                    "format": output_format,
                 }
             )
         )

packages/ai/src/upscaling.ts

@@ -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",
   };
 }