Refactor code structure for improved readability and maintainability

This commit is contained in:
2026-07-22 22:58:28 +05:30
parent d1308bf149
commit bc8bf2007d
99 changed files with 4784 additions and 111 deletions
+141 -21
View File
@@ -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) {