Refactor code structure for improved readability and maintainability
This commit is contained in:
@@ -1,5 +1,8 @@
|
||||
import { type NextRequest, NextResponse } from 'next/server';
|
||||
|
||||
import { RunpodError } from '../../../../lib/runpod/client';
|
||||
import { rpImg2Img, rpInpaint, rpRemoveBackground, rpUpscale } from '../../../../lib/runpod/endpoints';
|
||||
|
||||
export const runtime = 'nodejs';
|
||||
export const maxDuration = 60; // SD / cold-start models can take a while
|
||||
|
||||
@@ -7,6 +10,12 @@ export const maxDuration = 60; // SD / cold-start models can take a while
|
||||
* Generative image-edit proxy. The BACKEND is pluggable so you are not tied to a
|
||||
* paid model — pick one with env `AI_EDIT_PROVIDER` (default `auto`):
|
||||
*
|
||||
* - `runpod` → your RunPod serverless GPU endpoints (one per model, per
|
||||
* docs/model-api-spec). Maps each op → endpoint: restore/upscale
|
||||
* → Real-ESRGAN (#7), colorize → DDColor (#8), replace-sky/
|
||||
* magic-eraser/generative-fill → SD 3.5 masked inpaint (#9),
|
||||
* prompt → SD 3.5 img2img (#10). Env: RUNPOD_API_KEY + per-model
|
||||
* RUNPOD_*_URL. Key stays server-side. See docs/AI-SETUP.md.
|
||||
* - `local` → your own Stable Diffusion server (Automatic1111 / Forge /
|
||||
* SD.Next img2img API). Env: LOCAL_SD_URL (e.g. http://127.0.0.1:7860).
|
||||
* 100% free + private, needs a GPU box you run. Most reliable.
|
||||
@@ -14,10 +23,10 @@ export const maxDuration = 60; // SD / cold-start models can take a while
|
||||
* Env: HF_API_TOKEN (free), HF_IMAGE_MODEL (default instruct-pix2pix).
|
||||
* - `gemini` → Google Gemini image model (needs a billed key for image output).
|
||||
* Env: GEMINI_API_KEY, GEMINI_IMAGE_MODEL.
|
||||
* - `auto` → first configured of: local → huggingface → gemini.
|
||||
* - `auto` → first configured of: runpod → local → huggingface → gemini.
|
||||
*
|
||||
* NOTE: `remove-background` never reaches here — it runs fully in-browser (@imgly,
|
||||
* free, no key). Object detection / faces / OCR / embeddings are also 100% in-browser.
|
||||
* NOTE: `remove-background` runs in-browser by default (@imgly, free, no key), so
|
||||
* it usually never reaches here. Object detection uses its own route (/api/ai/classify).
|
||||
* See docs/AI-SETUP.md.
|
||||
*/
|
||||
|
||||
@@ -31,19 +40,39 @@ const OP_PROMPTS: Record<string, string> = {
|
||||
|
||||
const MAX_BASE64 = 4_000_000; // ~3 MB decoded — stays under serverless body limits (e.g. Vercel ~4.5MB)
|
||||
|
||||
type Provider = 'local' | 'huggingface' | 'gemini' | 'none';
|
||||
type Provider = 'runpod' | 'local' | 'huggingface' | 'gemini' | 'none';
|
||||
|
||||
function resolveProvider(): Provider {
|
||||
const explicit = (process.env.AI_EDIT_PROVIDER || 'auto').toLowerCase();
|
||||
if (explicit === 'local' || explicit === 'huggingface' || explicit === 'gemini') return explicit;
|
||||
if (
|
||||
explicit === 'runpod' ||
|
||||
explicit === 'local' ||
|
||||
explicit === 'huggingface' ||
|
||||
explicit === 'gemini'
|
||||
)
|
||||
return explicit;
|
||||
if (explicit === 'none') return 'none';
|
||||
// auto: prefer a private local server, then free HF, then Gemini.
|
||||
// auto: prefer RunPod GPU endpoints, then a private local server, then free HF, then Gemini.
|
||||
// Detect RunPod when the key + ANY image endpoint URL is set (an upscale/colorize-only
|
||||
// deployment is valid — not just the SD ones).
|
||||
if (
|
||||
process.env.RUNPOD_API_KEY &&
|
||||
(process.env.RUNPOD_SD_IMG2IMG_URL ||
|
||||
process.env.RUNPOD_SD_INPAINT_URL ||
|
||||
process.env.RUNPOD_UPSCALE_URL ||
|
||||
process.env.RUNPOD_COLORIZE_URL ||
|
||||
process.env.RUNPOD_BG_REMOVE_URL)
|
||||
)
|
||||
return 'runpod';
|
||||
if (process.env.LOCAL_SD_URL) return 'local';
|
||||
if (process.env.HF_API_TOKEN) return 'huggingface';
|
||||
if (process.env.GEMINI_API_KEY) return 'gemini';
|
||||
return 'none';
|
||||
}
|
||||
|
||||
/** Ops that only the RunPod (mask/fixed-function) backend can serve. */
|
||||
const RUNPOD_ONLY_OPS = new Set(['upscale', 'magic-eraser', 'generative-fill']);
|
||||
|
||||
interface EditResult {
|
||||
imageBase64: string;
|
||||
mimeType: string;
|
||||
@@ -55,7 +84,7 @@ export async function POST(req: NextRequest) {
|
||||
return NextResponse.json(
|
||||
{
|
||||
error:
|
||||
'AI image editing is not configured. Set LOCAL_SD_URL (own Stable Diffusion), HF_API_TOKEN (free Hugging Face), or GEMINI_API_KEY. Background removal and all analysis still work with no key. See docs/AI-SETUP.md.',
|
||||
'AI image editing is not configured. Set AI_EDIT_PROVIDER=runpod + RUNPOD_API_KEY + the per-model RUNPOD_*_URL vars (RunPod GPU), or LOCAL_SD_URL (own Stable Diffusion), HF_API_TOKEN (free Hugging Face), or GEMINI_API_KEY. Background removal and all analysis still work with no key. See docs/AI-SETUP.md.',
|
||||
},
|
||||
{ status: 503 },
|
||||
);
|
||||
@@ -68,45 +97,136 @@ export async function POST(req: NextRequest) {
|
||||
return NextResponse.json({ error: 'Invalid request body.' }, { status: 400 });
|
||||
}
|
||||
|
||||
const { imageBase64, mimeType, op } = (body ?? {}) as {
|
||||
const { imageBase64, mimeType, op, maskBase64, params } = (body ?? {}) as {
|
||||
imageBase64?: unknown;
|
||||
mimeType?: unknown;
|
||||
op?: { type?: string; prompt?: string };
|
||||
op?: { type?: string; prompt?: string; factor?: number };
|
||||
maskBase64?: unknown;
|
||||
params?: unknown;
|
||||
};
|
||||
|
||||
if (typeof imageBase64 !== 'string' || imageBase64.length === 0 || imageBase64.length > MAX_BASE64) {
|
||||
return NextResponse.json({ error: 'Invalid or oversized image (try a smaller image).' }, { status: 400 });
|
||||
if (typeof imageBase64 !== 'string' || imageBase64.length === 0) {
|
||||
return NextResponse.json({ error: 'Invalid image.' }, { status: 400 });
|
||||
}
|
||||
const hasMask = typeof maskBase64 === 'string' && maskBase64.length > 0;
|
||||
// Image + mask share one request body — budget them together against the cap.
|
||||
if (imageBase64.length + (hasMask ? (maskBase64 as string).length : 0) > MAX_BASE64) {
|
||||
return NextResponse.json({ error: 'Image (plus mask) is too large — try a smaller image.' }, { status: 400 });
|
||||
}
|
||||
const safeMime =
|
||||
typeof mimeType === 'string' && /^image\/(jpeg|png|webp)$/.test(mimeType) ? mimeType : 'image/jpeg';
|
||||
|
||||
const opType = op?.type ?? '';
|
||||
if (provider !== 'runpod' && RUNPOD_ONLY_OPS.has(opType)) {
|
||||
return NextResponse.json(
|
||||
{ error: 'This edit needs the RunPod backend (set AI_EDIT_PROVIDER=runpod).' },
|
||||
{ status: 400 },
|
||||
);
|
||||
}
|
||||
|
||||
// Build the instruction from an allow-listed op (never trust arbitrary server prompts).
|
||||
let instruction: string;
|
||||
if (op?.type === 'prompt') {
|
||||
const p = typeof op.prompt === 'string' ? op.prompt.trim() : '';
|
||||
let instruction = '';
|
||||
if (opType === 'prompt' || opType === 'generative-fill') {
|
||||
const p = typeof op?.prompt === 'string' ? op.prompt.trim() : '';
|
||||
if (!p) return NextResponse.json({ error: 'Empty prompt.' }, { status: 400 });
|
||||
instruction = p.slice(0, 500);
|
||||
} else if (op?.type === 'replace-sky' && typeof op.prompt === 'string' && op.prompt.trim()) {
|
||||
instruction = `Replace the sky with: ${op.prompt.trim().slice(0, 300)}. Keep the foreground unchanged and photorealistic.`;
|
||||
} else if (op?.type && OP_PROMPTS[op.type]) {
|
||||
instruction = OP_PROMPTS[op.type]!;
|
||||
} else {
|
||||
} else if (opType === 'replace-sky') {
|
||||
instruction =
|
||||
typeof op?.prompt === 'string' && op.prompt.trim()
|
||||
? `Replace the sky with: ${op.prompt.trim().slice(0, 300)}. Keep the foreground unchanged and photorealistic.`
|
||||
: OP_PROMPTS['replace-sky']!;
|
||||
} else if (opType === 'magic-eraser') {
|
||||
instruction = 'Fill the selected region with a clean, seamless, plausible background. Photorealistic.';
|
||||
} else if (OP_PROMPTS[opType]) {
|
||||
instruction = OP_PROMPTS[opType]!;
|
||||
} else if (opType !== 'upscale') {
|
||||
return NextResponse.json({ error: 'Unsupported operation.' }, { status: 400 });
|
||||
}
|
||||
|
||||
try {
|
||||
let result: EditResult;
|
||||
if (provider === 'local') result = await editLocal(instruction, imageBase64);
|
||||
if (provider === 'runpod')
|
||||
result = await editRunPod(op ?? {}, imageBase64, instruction, hasMask ? (maskBase64 as string) : undefined, params);
|
||||
else if (provider === 'local') result = await editLocal(instruction, imageBase64);
|
||||
else if (provider === 'huggingface') result = await editHuggingFace(instruction, imageBase64);
|
||||
else result = await editGemini(instruction, imageBase64, safeMime);
|
||||
return NextResponse.json(result);
|
||||
} catch (err) {
|
||||
const message = err instanceof Error ? err.message : 'AI request failed.';
|
||||
const status = err instanceof AiError ? err.status : 502;
|
||||
const status = err instanceof AiError || err instanceof RunpodError ? err.status : 502;
|
||||
return NextResponse.json({ error: message }, { status });
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Backend: RunPod serverless GPU endpoints (one model per endpoint).
|
||||
// Each op maps to the endpoint from the spec; the API key + URLs stay server-side.
|
||||
// ---------------------------------------------------------------------------
|
||||
interface SdParams {
|
||||
negativePrompt?: string;
|
||||
strength?: number;
|
||||
steps?: number;
|
||||
seed?: number;
|
||||
guidanceScale?: number;
|
||||
}
|
||||
|
||||
function sanitizeParams(raw: unknown): SdParams {
|
||||
const p = (raw ?? {}) as Record<string, unknown>;
|
||||
const out: SdParams = {};
|
||||
if (typeof p.negativePrompt === 'string' && p.negativePrompt.trim())
|
||||
out.negativePrompt = p.negativePrompt.trim().slice(0, 300);
|
||||
const strength = Number(p.strength);
|
||||
if (Number.isFinite(strength)) out.strength = Math.max(0, Math.min(1, strength));
|
||||
const steps = Number(p.steps);
|
||||
if (Number.isFinite(steps)) out.steps = Math.max(1, Math.min(60, Math.round(steps)));
|
||||
const guidance = Number(p.guidanceScale);
|
||||
if (Number.isFinite(guidance)) out.guidanceScale = Math.max(1, Math.min(20, guidance));
|
||||
const seed = Number(p.seed);
|
||||
if (Number.isFinite(seed)) out.seed = Math.max(0, Math.min(2_147_483_647, Math.round(seed)));
|
||||
return out;
|
||||
}
|
||||
|
||||
async function editRunPod(
|
||||
op: { type?: string; prompt?: string; factor?: number },
|
||||
imageBase64: string,
|
||||
instruction: string,
|
||||
maskBase64: string | undefined,
|
||||
rawParams: unknown,
|
||||
): Promise<EditResult> {
|
||||
const params = sanitizeParams(rawParams);
|
||||
switch (op.type) {
|
||||
case 'remove-background':
|
||||
// U²-Net via rembg (#6) — a real endpoint replacing the flaky in-browser remover.
|
||||
return rpRemoveBackground(imageBase64);
|
||||
case 'restore':
|
||||
// Real-ESRGAN (#7) with the GFPGAN face pass = "Restore & Enhance".
|
||||
return rpUpscale(imageBase64, 4, true);
|
||||
case 'upscale':
|
||||
return rpUpscale(imageBase64, op.factor === 4 ? 4 : 2, false);
|
||||
case 'colorize':
|
||||
// The dedicated DDColor endpoint kept hard-crashing (modelscope). Route
|
||||
// colorize through the img2img model as an instruction instead — reliable,
|
||||
// and high quality once img2img is FLUX. (RUNPOD_COLORIZE_URL now unused.)
|
||||
return rpImg2Img({ imageB64: imageBase64, prompt: instruction, ...params });
|
||||
case 'prompt':
|
||||
return rpImg2Img({ imageB64: imageBase64, prompt: instruction, ...params }); // SD 3.5 img2img (#10)
|
||||
case 'replace-sky':
|
||||
// True sky replacement is masked inpaint (#9). Without a mask (no in-app sky
|
||||
// segmentation yet) degrade to a low-strength img2img (#10) so the foreground
|
||||
// is mostly preserved. Add a mask (brush UI / sky-seg) to get real #9.
|
||||
if (maskBase64)
|
||||
return rpInpaint({ imageB64: imageBase64, maskB64: maskBase64, prompt: instruction, ...params });
|
||||
return rpImg2Img({ imageB64: imageBase64, prompt: instruction, ...params, strength: params.strength ?? 0.4 });
|
||||
case 'magic-eraser':
|
||||
case 'generative-fill':
|
||||
// SD 3.5 masked inpaint (#9) — white in the mask = the region to regenerate.
|
||||
if (!maskBase64) throw new AiError('This edit needs a mask/selection.', 400);
|
||||
return rpInpaint({ imageB64: imageBase64, maskB64: maskBase64, prompt: instruction, ...params });
|
||||
default:
|
||||
throw new AiError('Unsupported operation.', 400);
|
||||
}
|
||||
}
|
||||
|
||||
class AiError extends Error {
|
||||
status: number;
|
||||
constructor(message: string, status = 502) {
|
||||
|
||||
Reference in New Issue
Block a user