From 28e16e09a41f31be2a6fed60a84b3d2700d7ab42 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Kr=C3=BCger?= <valknar@pivoine.art>
Date: Tue, 17 Feb 2026 09:01:35 +0100
Subject: [PATCH] feat: add upscale option to script

---
 img_api_generate.sh | 93 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 92 insertions(+), 1 deletion(-)

diff --git a/img_api_generate.sh b/img_api_generate.sh
index 66bc171..f004db0 100755
--- a/img_api_generate.sh
+++ b/img_api_generate.sh
@@ -76,6 +76,7 @@ INPUT_IMAGE=""
 ASYNC_MODE=false
 DRY_RUN=false
 VERBOSE=false
+UPSCALE=false
 
 # Temp files for cleanup
 _TMP_FILES=()
@@ -776,6 +777,59 @@ swap_face() {
     fi
 }
 
+# ============================================================================
+# Real-ESRGAN UPSCALE
+# ============================================================================
+realesrgan_upscale() {
+    local input_image="$1"
+    local -n _up_output_path="$2"
+    local url="${BASE_URL}/realesrgan/upscale"
+
+    if [[ "$DRY_RUN" == true ]]; then
+        print_section "Dry Run - Real-ESRGAN Upscale curl command"
+        echo ""
+        echo -e "${DIM}curl -s -X POST \\\"\n"
+        echo "  -H 'X-Api-Key: \\${TOKEN}' \\\"\n"
+        echo "  -F 'image=@${input_image}' \\\"\n"
+        echo "  -F 'model=RealESRGAN_x4plus' \\\"\n"
+        echo -e "  '${url}'${RESET}"
+        echo ""
+        return 0
+    fi
+
+    print_section "Upscaling (Real-ESRGAN)"
+    print_field "$PALETTE" "Model" "RealESRGAN_x4plus"
+    print_field "$CAMERA" "Input" "$input_image"
+    echo ""
+
+    local tmp_resp
+    tmp_resp=$(mktemp)
+    _TMP_FILES+=("$tmp_resp")
+
+    (api_curl_binary POST "$url" "${_up_output_path}" \
+        -F "image=@${input_image}" \
+        -F "model=RealESRGAN_x4plus" > "$tmp_resp" 2>&1) &
+    local pid=$!
+    spinner "$pid" "Upscaling image"
+    wait "$pid" || {
+        cat "$tmp_resp" >&2
+        print_error "Upscale failed"
+        return 1
+    }
+
+    if [[ -f "${_up_output_path}" && -s "${_up_output_path}" ]]; then
+        fix_extension "${_up_output_path}" _up_output_path
+        local size
+        size=$(du -h "${_up_output_path}" | cut -f1)
+        print_success "Upscaled saved: ${BOLD_WHITE}${_up_output_path}${RESET} (${size})"
+        return 0
+    else
+        print_error "Upscale result is empty or missing"
+        rm -f "${_up_output_path}" 2>/dev/null || true
+        return 1
+    fi
+}
+
 # ============================================================================
 # HELP
 # ============================================================================
@@ -809,6 +863,7 @@ FACE SWAP:
 
 OUTPUT:
   -o, --output FILE           Output file path (default: auto-generated)
+    -u, --upscale               Upscale final image using Real-ESRGAN (RealESRGAN_x4plus)
 
 MODES:
       --async                 Use async mode with polling instead of sync
@@ -895,6 +950,10 @@ parse_args() {
                 DRY_RUN=true
                 shift
                 ;;
+            -u|--upscale)
+                UPSCALE=true
+                shift
+                ;;
             --verbose)
                 VERBOSE=true
                 shift
@@ -957,9 +1016,18 @@ main() {
 
     if [[ "$DRY_RUN" == true ]]; then
         # Show face swap dry-run if applicable
+        local dry_input="/tmp/generated.png"
         if [[ -n "$FACE_IMAGE" ]]; then
-            swap_face "/tmp/generated.png" "$FACE_IMAGE" "/tmp/faceswap.png"
+            tmp_faceswap="/tmp/faceswap.png"
+            swap_face "/tmp/generated.png" "$FACE_IMAGE" tmp_faceswap
+            dry_input="$tmp_faceswap"
         fi
+
+        if [[ "$UPSCALE" == true ]]; then
+            tmp_upscaled="/tmp/upscaled.png"
+            realesrgan_upscale "$dry_input" tmp_upscaled
+        fi
+
         print_banner "${SPARKLES} Dry Run Complete ${SPARKLES}"
         exit 0
     fi
@@ -990,6 +1058,29 @@ main() {
         swap_face "$output_path" "$FACE_IMAGE" faceswap_path || true
     fi
 
+    # Determine final input for optional upscale (prefer faceswap if present)
+    local final_input="$output_path"
+    if [[ -n "${faceswap_path:-}" && -f "${faceswap_path}" ]]; then
+        final_input="${faceswap_path}"
+    fi
+
+    # Upscale with Real-ESRGAN if requested
+    if [[ "$UPSCALE" == true ]]; then
+        local upscale_path
+        if [[ -n "$OUTPUT_FILE" ]]; then
+            local ext="${OUTPUT_FILE##*.}"
+            local base="${OUTPUT_FILE%.*}"
+            upscale_path="${base}_upscaled.${ext}"
+        else
+            upscale_path="$(generate_output_filename "upscaled")"
+        fi
+
+        realesrgan_upscale "$final_input" upscale_path || exit 1
+
+        # Use the upscaled image for summary/output
+        output_path="$upscale_path"
+    fi
+
     # Summary
     print_section "Summary"
     print_field "$PALETTE" "Model" "$MODEL"