Refactor code structure for improved readability and maintainability
This commit is contained in:
@@ -298,7 +298,54 @@ Phases 1–3 complete & verified (see [requirements.md](requirements.md)). `pnpm
|
||||
DEPLOY: `docs/DEPLOY-MANUAL.md` — 3 no-Git paths: (A) Vercel CLI `vercel --prod` from repo root, Root Dir=apps/web (recommended; keeps SSR+nonce CSP+AI route); (B) self-host Node (`pnpm build` → `pnpm --filter web start` on Render/Railway/VPS); (C) static export to any host — requires `output:'export'` + DISABLING the nonce middleware/layout headers (security trade-off, no server AI route). SDK ships as source via transpilePackages so "whole project" = build the demo (no npm publish).
|
||||
AI: `apps/web/src/app/api/ai/edit/route.ts` REFACTORED to a pluggable backend via `AI_EDIT_PROVIDER` (auto|local|huggingface|gemini|none). auto = local→huggingface→gemini. `local` = own Stable Diffusion (AUTOMATIC1111/Forge/SD.Next `/sdapi/v1/img2img`, env LOCAL_SD_URL+denoise/steps/sampler; free+private). `huggingface` = free Inference API instruct-pix2pix (HF_API_TOKEN, HF_IMAGE_MODEL; 503 cold-start handled). `gemini` kept (needs BILLED key for image). MAX_BASE64 lowered to 4M (under Vercel 4.5MB). Server-side fetches → no CSP change. KEY POINT documented in `docs/AI-SETUP.md`: analysis (objects/faces/OCR/embeddings via TF.js/face-api/tesseract/transformers) + Remove Background (@imgly) already run FREE in-browser, no key — only generative pixel edits needed a backend. `.env.example` rewritten with Supabase + all AI options. typecheck + build pass. Did NOT git commit.
|
||||
|
||||
- S-runpod (2026-07-18): **RunPod endpoint integration (model-api-spec.docx) — image edits + detection, all server-proxied.**
|
||||
Planned via an `ultracode` Workflow (5 parallel readers → Opus blueprint). GOAL: make the app talk to the 15 RunPod
|
||||
endpoints (server still down; wired + typecheck/build-verified, not yet live-tested against real endpoints).
|
||||
NEW server-only client module `apps/web/src/lib/runpod/{types,base64,client,endpoints}.ts`: `runpodCall` transport
|
||||
(bearer `RUNPOD_API_KEY`, `{input}` envelope, `/runsync` inline + `/run`→`/status/{id}` poll under a 55s budget),
|
||||
`pickOutputImage` (handles image|image_png|images[0]|nested), and one typed `rp*` fn per endpoint. `RunpodError.status`
|
||||
surfaces as the editor's red error text.
|
||||
EDIT ROUTE `apps/web/src/app/api/ai/edit/route.ts`: added `'runpod'` provider (auto-preferred when RUNPOD_API_KEY +
|
||||
an SD URL set). `editRunPod` op→endpoint map: restore→#7 ESRGAN(+face), upscale→#7, colorize→#8 DDColor, prompt→#10
|
||||
SD img2img, replace-sky→#9 inpaint IF mask else #10 low-strength, magic-eraser/generative-fill→#9 masked inpaint
|
||||
(400 if no mask). Body now accepts `maskBase64` + `params` (negativePrompt/strength/steps/seed/guidanceScale, all
|
||||
clamped); image+mask share the ONE ~4MB body cap. RUNPOD_ONLY_OPS (upscale/magic-eraser/generative-fill) 400 on
|
||||
non-runpod backends.
|
||||
MASK FLOW (the SD 3.5 ask): client `createDemoAIProvider.generativeEdit` rasterizes `op.mask` (ImageData) → PNG
|
||||
matched to the downscaled image dims via new `lib/ai/imageEncode.ts maskToBase64`, strips the non-serializable mask
|
||||
from the wire op, sends `maskBase64`. Output interpreted unchanged (base64→Blob→aiResultUrl→Save), so no UI change
|
||||
needed to SHOW results. STILL MISSING: a brush-mask UI in PhotoEditor to CREATE masks (magic-eraser/generative-fill
|
||||
have no button yet + return the 400 until then); replace-sky degrades to img2img until a mask/sky-seg exists.
|
||||
DETECTION #1: new `apps/web/src/app/api/ai/classify/route.ts` + `lib/ai/runpodYoloProvider.ts` (POSTs image+dims,
|
||||
normalizes YOLO boxes→DetectedObject fractions 0..1 — `normalizeBox` assumes ultralytics xyxy pixels, VERIFY vs real
|
||||
endpoint). Gated by `NEXT_PUBLIC_APG_RUNPOD_DETECT=true`, else in-browser COCO-SSD (also the automatic fallback).
|
||||
ALREADY LOCAL (no RunPod): screenshot #4 / geo #15 (exifr+classify.ts), tag-rename #5 (state), bg-remove #6 (@imgly).
|
||||
SCAFFOLD ONLY (typed rp* + env, NO route/UI): tilt #2, voice #3, audio-denoise #12; DOC-ONLY: video #13/#14 (need
|
||||
out-of-band upload+queue, incompatible with Vercel 4.5MB/60s + in-browser bake). All flagged in docs/AI-SETUP.md.
|
||||
CSP UNCHANGED (server-side fetch bypasses browser CSP; middleware excludes /api/*). ENV: `.env.example` RunPod block;
|
||||
docs/AI-SETUP.md "Option 0 — RunPod" (op→endpoint table + mask/detection notes) + docs/DEPLOY.md server-env table
|
||||
(RUNPOD_* never NEXT_PUBLIC). VERIFIED: `pnpm -r typecheck` clean (both pkgs); `pnpm build` passes, both new routes
|
||||
compiled (/api/ai/edit, /api/ai/classify dynamic), /gallery 199KB, only the pre-existing benign face-api warning.
|
||||
Did NOT git commit. NEXT: brush-mask UI (unlocks true #9 for eraser/fill + real sky mask); outpaint #11 (client pads
|
||||
canvas+mask→generative-fill); #5 alias rename map; then Vercel deploy once endpoints are up. VERIFY normalizeBox +
|
||||
the `{input}`/`output` field names against the live endpoints before trusting detection/edit results.
|
||||
ADVERSARIAL REVIEW (ultracode Workflow, 3 dims find→verify, 20 agents; 12 CONFIRMED, 5 refuted) — ALL 12 FIXED &
|
||||
re-verified (typecheck + build pass, both routes compiled, /gallery 199KB): (1) normalizeBox rewritten to handle
|
||||
object-form xyxy `{x1,y1,x2,y2}`/`{left,top,right,bottom}` (ultralytics tojson shape — was dropped!) + array boxes
|
||||
keyed correctly (xyxy vs xywh/COCO bbox=[x,y,w,h] vs generic box=xyxy); (2) label priority fixed — string `name`
|
||||
before numeric `class` index (was returning "0" for ultralytics), empty-string label falls through to class_<id>;
|
||||
(3) envNum treats blank env ('' → 0) as unset; (4) rpUpscale no longer lets RUNPOD_UPSCALE_SCALE override the
|
||||
per-request factor (removed the env knob); (5) client envelope: a raw /runsync body with a `status` field but no
|
||||
`id` is now correctly treated as inline output (was misrouted to poll→throw); (6) poll loop clamps sleep+abort to
|
||||
the remaining budget so it can't overshoot Vercel 60s; (7) pickOutputImage validates strings under generic
|
||||
output/result/data keys (looksLikeImageData) so a status/id string isn't returned as an image; (8) auto provider
|
||||
detects runpod on ANY RUNPOD_*_URL (not just SD) so an upscale/colorize-only deploy works; (9) maskToBase64 sets
|
||||
imageSmoothingEnabled=false so binary masks stay crisp (no grey halo). Dropped RUNPOD_UPSCALE_SCALE from .env.example.
|
||||
REFUTED (correctly, no change): status-URL query-string, combined-body-cap rejecting masks, replace-sky client
|
||||
strength, 'mask' in op undefined-crash, mask-op has no UI button (all non-bugs / already-safe).
|
||||
|
||||
## Key files
|
||||
- RunPod integration (S-runpod): `apps/web/src/lib/runpod/*` (server) + `apps/web/src/app/api/ai/{edit,classify}/route.ts` + `apps/web/src/lib/ai/{runpodYoloProvider,imageEncode}.ts` (client). Env: `RUNPOD_API_KEY` + `RUNPOD_*_URL`.
|
||||
- Advanced video editor (S-advanced): `packages/photo-sdk/src/lib/videoBake.ts` + `lib/videoTimeline.ts` + `components/editor/VideoEditor.tsx`.
|
||||
- AI edit backend (S-deploy-ai): `apps/web/src/app/api/ai/edit/route.ts` (env AI_EDIT_PROVIDER).
|
||||
- SDK: `packages/photo-sdk/src/{components,store,adapters,ai,lib,icons,styles}`.
|
||||
|
||||
+17
-1
@@ -13,7 +13,23 @@
|
||||
# NOT required to run the app: object detection, faces, OCR, and Remove Background
|
||||
# all run FREE in-browser with no key. Full guide: docs/AI-SETUP.md
|
||||
# ---------------------------------------------------------------------------
|
||||
# AI_EDIT_PROVIDER=auto # auto | local | huggingface | gemini | none
|
||||
# AI_EDIT_PROVIDER=auto # auto | runpod | local | huggingface | gemini | none
|
||||
#
|
||||
# Option 0 — RunPod serverless GPU endpoints (one per model; key stays server-side).
|
||||
# `auto` picks this first when RUNPOD_API_KEY + an SD URL are set. Deploy endpoints so
|
||||
# their URLs end in /runsync. Full map (op → endpoint) in docs/AI-SETUP.md.
|
||||
# AI_EDIT_PROVIDER=runpod
|
||||
# RUNPOD_API_KEY=rpa_xxxxxxxx
|
||||
# RUNPOD_SD_IMG2IMG_URL=https://api.runpod.ai/v2/<id>/runsync # #10 prompt (img2img)
|
||||
# RUNPOD_SD_INPAINT_URL=https://api.runpod.ai/v2/<id>/runsync # #9 replace-sky / magic-eraser / generative-fill (masked)
|
||||
# RUNPOD_UPSCALE_URL=https://api.runpod.ai/v2/<id>/runsync # #7 restore / upscale (Real-ESRGAN)
|
||||
# RUNPOD_COLORIZE_URL=https://api.runpod.ai/v2/<id>/runsync # #8 colorize (DDColor)
|
||||
# RUNPOD_BG_REMOVE_URL=https://api.runpod.ai/v2/<id>/runsync # #6 (optional; default is in-browser @imgly)
|
||||
# RUNPOD_YOLO_URL=https://api.runpod.ai/v2/<id>/runsync # #1 detection (needs the flag below)
|
||||
# NEXT_PUBLIC_APG_RUNPOD_DETECT=false # true = use #1 YOLO server detection (else free in-browser COCO-SSD)
|
||||
# Optional SD tuning: RUNPOD_SD_STEPS RUNPOD_SD_STRENGTH RUNPOD_SD_GUIDANCE RUNPOD_SD_NEGATIVE_PROMPT
|
||||
# Scaffold-only (typed client, no route/UI yet): RUNPOD_TILT_URL (#2) RUNPOD_STT_URL (#3)
|
||||
# RUNPOD_AUDIO_DENOISE_URL (#12) RUNPOD_VIDEO_FPS_URL (#13) RUNPOD_VIDEO_UPSCALE_URL (#14)
|
||||
#
|
||||
# Option 1 — your own local Stable Diffusion (AUTOMATIC1111/Forge/SD.Next; free + private):
|
||||
# LOCAL_SD_URL=http://127.0.0.1:7860
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
import { type NextRequest, NextResponse } from 'next/server';
|
||||
|
||||
import { RunpodError } from '../../../../lib/runpod/client';
|
||||
import { rpDetect } from '../../../../lib/runpod/endpoints';
|
||||
|
||||
export const runtime = 'nodejs';
|
||||
export const maxDuration = 60;
|
||||
|
||||
const MAX_BASE64 = 4_000_000; // ~3 MB decoded — under serverless body limits
|
||||
|
||||
/**
|
||||
* Object-detection proxy for the RunPod YOLO construction-material classifier (#1).
|
||||
* The key + endpoint URL stay server-side. The client (runpodYoloProvider) calls
|
||||
* this only when NEXT_PUBLIC_APG_RUNPOD_DETECT is on; otherwise detection runs
|
||||
* fully in-browser (COCO-SSD) with no server round-trip. Returns the SDK's
|
||||
* DetectedObject[] shape (box as 0..1 fractions) so it drops straight into the
|
||||
* existing Objects browser / smart albums / search.
|
||||
*/
|
||||
export async function POST(req: NextRequest) {
|
||||
if (!process.env.RUNPOD_API_KEY || !process.env.RUNPOD_YOLO_URL) {
|
||||
return NextResponse.json(
|
||||
{ error: 'RunPod detection is not configured (set RUNPOD_API_KEY + RUNPOD_YOLO_URL).' },
|
||||
{ status: 503 },
|
||||
);
|
||||
}
|
||||
|
||||
let body: unknown;
|
||||
try {
|
||||
body = await req.json();
|
||||
} catch {
|
||||
return NextResponse.json({ error: 'Invalid request body.' }, { status: 400 });
|
||||
}
|
||||
|
||||
const { imageBase64, width, height } = (body ?? {}) as {
|
||||
imageBase64?: unknown;
|
||||
width?: unknown;
|
||||
height?: unknown;
|
||||
};
|
||||
if (typeof imageBase64 !== 'string' || imageBase64.length === 0 || imageBase64.length > MAX_BASE64) {
|
||||
return NextResponse.json({ error: 'Invalid or oversized image.' }, { status: 400 });
|
||||
}
|
||||
const w = Number(width);
|
||||
const h = Number(height);
|
||||
|
||||
try {
|
||||
const objects = await rpDetect(
|
||||
imageBase64,
|
||||
Number.isFinite(w) && w > 0 ? w : 1,
|
||||
Number.isFinite(h) && h > 0 ? h : 1,
|
||||
);
|
||||
return NextResponse.json({ objects });
|
||||
} catch (err) {
|
||||
const message = err instanceof Error ? err.message : 'Detection failed.';
|
||||
const status = err instanceof RunpodError ? err.status : 502;
|
||||
return NextResponse.json({ error: message }, { status });
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
import { type NextRequest, NextResponse } from 'next/server';
|
||||
|
||||
import { RunpodError } from '../../../../lib/runpod/client';
|
||||
import { rpDenoiseAudio } from '../../../../lib/runpod/endpoints';
|
||||
|
||||
export const runtime = 'nodejs';
|
||||
export const maxDuration = 60; // cold-start denoise worker can take a while
|
||||
|
||||
/**
|
||||
* Audio noise-removal proxy. Accepts base64 WAV (48 kHz mono PCM16, produced
|
||||
* in-browser by lib/audioCapture) and returns a cleaned base64 WAV. Calls the
|
||||
* RunPod audio-denoise endpoint (RUNPOD_AUDIO_DENOISE_URL) — key stays
|
||||
* server-side. Used before transcription on noisy sites. See docs/AI-SETUP.md.
|
||||
*/
|
||||
|
||||
const MAX_BASE64 = 12_000_000; // ~9 MB decoded WAV
|
||||
|
||||
export async function POST(req: NextRequest) {
|
||||
let body: { audio?: unknown };
|
||||
try {
|
||||
body = await req.json();
|
||||
} catch {
|
||||
return NextResponse.json({ error: 'Invalid JSON body.' }, { status: 400 });
|
||||
}
|
||||
|
||||
const audio = typeof body.audio === 'string' ? body.audio : '';
|
||||
if (!audio) return NextResponse.json({ error: 'Missing audio.' }, { status: 400 });
|
||||
if (audio.length > MAX_BASE64) {
|
||||
return NextResponse.json({ error: 'Audio too long — keep it under ~30s.' }, { status: 413 });
|
||||
}
|
||||
|
||||
try {
|
||||
const { audioB64 } = await rpDenoiseAudio(audio);
|
||||
return NextResponse.json({ audio: audioB64 });
|
||||
} catch (e) {
|
||||
const msg =
|
||||
e instanceof RunpodError ? e.message : e instanceof Error ? e.message : 'Denoise failed.';
|
||||
return NextResponse.json({ error: msg }, { status: 502 });
|
||||
}
|
||||
}
|
||||
@@ -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) {
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
import { type NextRequest, NextResponse } from 'next/server';
|
||||
|
||||
import { RunpodError } from '../../../../lib/runpod/client';
|
||||
import { rpTilt } from '../../../../lib/runpod/endpoints';
|
||||
|
||||
export const runtime = 'nodejs';
|
||||
export const maxDuration = 60; // cold-start tilt worker can take a while
|
||||
|
||||
/**
|
||||
* Camera-tilt estimation proxy. Accepts a base64 image and returns
|
||||
* {rollDegrees, pitchDegrees, fovDegrees} from the RunPod tilt endpoint
|
||||
* (RUNPOD_TILT_URL) so the editor can auto-straighten. Requires the tilt
|
||||
* endpoint to be deployed; errors clearly if RUNPOD_TILT_URL is unset.
|
||||
*/
|
||||
|
||||
const MAX_BASE64 = 4_000_000; // ~3 MB decoded
|
||||
|
||||
export async function POST(req: NextRequest) {
|
||||
let body: { image?: unknown };
|
||||
try {
|
||||
body = await req.json();
|
||||
} catch {
|
||||
return NextResponse.json({ error: 'Invalid JSON body.' }, { status: 400 });
|
||||
}
|
||||
|
||||
const image = typeof body.image === 'string' ? body.image : '';
|
||||
if (!image) return NextResponse.json({ error: 'Missing image.' }, { status: 400 });
|
||||
if (image.length > MAX_BASE64) {
|
||||
return NextResponse.json({ error: 'Image too large.' }, { status: 413 });
|
||||
}
|
||||
|
||||
try {
|
||||
const tilt = await rpTilt(image);
|
||||
return NextResponse.json(tilt);
|
||||
} catch (e) {
|
||||
const msg =
|
||||
e instanceof RunpodError
|
||||
? e.message
|
||||
: e instanceof Error
|
||||
? e.message
|
||||
: 'Tilt estimate failed.';
|
||||
return NextResponse.json({ error: msg }, { status: 502 });
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
import { type NextRequest, NextResponse } from 'next/server';
|
||||
|
||||
import { RunpodError } from '../../../../lib/runpod/client';
|
||||
import { rpTranscribe } from '../../../../lib/runpod/endpoints';
|
||||
|
||||
export const runtime = 'nodejs';
|
||||
export const maxDuration = 60; // cold-start STT worker can take a while
|
||||
|
||||
/**
|
||||
* Speech-to-text proxy for voice annotations. Accepts base64 WAV (16 kHz mono
|
||||
* PCM16, produced in-browser by lib/audioCapture) and returns the transcript.
|
||||
* Calls the RunPod voice-to-text endpoint (RUNPOD_STT_URL) — the key stays
|
||||
* server-side. See docs/AI-SETUP.md.
|
||||
*/
|
||||
|
||||
const MAX_BASE64 = 8_000_000; // ~6 MB decoded WAV — stays under serverless body limits
|
||||
|
||||
export async function POST(req: NextRequest) {
|
||||
let body: { audio?: unknown; language?: unknown };
|
||||
try {
|
||||
body = await req.json();
|
||||
} catch {
|
||||
return NextResponse.json({ error: 'Invalid JSON body.' }, { status: 400 });
|
||||
}
|
||||
|
||||
const audio = typeof body.audio === 'string' ? body.audio : '';
|
||||
const language = typeof body.language === 'string' ? body.language : undefined;
|
||||
if (!audio) return NextResponse.json({ error: 'Missing audio.' }, { status: 400 });
|
||||
if (audio.length > MAX_BASE64) {
|
||||
return NextResponse.json({ error: 'Audio too long — keep it under ~30s.' }, { status: 413 });
|
||||
}
|
||||
|
||||
try {
|
||||
const { transcript, segments } = await rpTranscribe(audio, { language, punctuation: true });
|
||||
return NextResponse.json({ transcript, segments });
|
||||
} catch (e) {
|
||||
const msg =
|
||||
e instanceof RunpodError
|
||||
? e.message
|
||||
: e instanceof Error
|
||||
? e.message
|
||||
: 'Transcription failed.';
|
||||
return NextResponse.json({ error: msg }, { status: 502 });
|
||||
}
|
||||
}
|
||||
@@ -2,17 +2,24 @@ import type { AIProvider } from '@photo-gallery/sdk';
|
||||
|
||||
import { createClipProvider } from './clipProvider';
|
||||
import { createFaceProvider } from './faceProvider';
|
||||
import { imageToBase64, maskToBase64 } from './imageEncode';
|
||||
import { createOCRProvider } from './ocrProvider';
|
||||
import { createRunpodYoloProvider } from './runpodYoloProvider';
|
||||
import { createTensorflowProvider } from './tensorflowProvider';
|
||||
|
||||
/**
|
||||
* The demo's AI provider, combining free in-browser capabilities + Gemini:
|
||||
* The demo's AI provider, combining free in-browser capabilities + a
|
||||
* server-proxied generative-edit backend (RunPod / local SD / Hugging Face / Gemini):
|
||||
* - object detection: TensorFlow.js COCO-SSD, fully in-browser (no key)
|
||||
* - face detection + recognition: face-api.js, in-browser → clustered into People
|
||||
* - OCR: tesseract.js, in-browser → searchable text + the Documents album (no key)
|
||||
* - background removal: @imgly, in-browser WASM (no key) — see generativeEdit
|
||||
* - other generative edits: Gemini, proxied through /api/ai/edit so the API key
|
||||
* stays server-side and is never exposed to the browser.
|
||||
* - other generative edits: proxied through /api/ai/edit (RunPod / local SD /
|
||||
* Hugging Face / Gemini, chosen by env) so the API key stays server-side.
|
||||
*
|
||||
* Object detection uses the in-browser COCO-SSD by default, or the RunPod YOLO
|
||||
* construction-material classifier (#1) when NEXT_PUBLIC_APG_RUNPOD_DETECT=true
|
||||
* (proxied through /api/ai/classify, with COCO-SSD as the automatic fallback).
|
||||
*/
|
||||
export function createDemoAIProvider(): AIProvider {
|
||||
const tf = createTensorflowProvider();
|
||||
@@ -20,9 +27,15 @@ export function createDemoAIProvider(): AIProvider {
|
||||
const ocr = createOCRProvider();
|
||||
const clip = createClipProvider();
|
||||
|
||||
// Opt-in server-side detection; defaults to free in-browser COCO-SSD.
|
||||
const detectObjects =
|
||||
process.env.NEXT_PUBLIC_APG_RUNPOD_DETECT === 'true'
|
||||
? createRunpodYoloProvider(tf).detectObjects
|
||||
: tf.detectObjects;
|
||||
|
||||
return {
|
||||
name: 'demo-ai (coco-ssd + face-api + tesseract + clip + gemini)',
|
||||
detectObjects: tf.detectObjects,
|
||||
name: 'demo-ai (coco-ssd/yolo + face-api + tesseract + clip + runpod-edit)',
|
||||
detectObjects,
|
||||
detectFaces: face.detectFaces,
|
||||
ocr: ocr.ocr,
|
||||
embedImage: clip.embedImage,
|
||||
@@ -30,18 +43,91 @@ export function createDemoAIProvider(): AIProvider {
|
||||
|
||||
async generativeEdit(item, image, op) {
|
||||
// Remove Background runs fully in-browser (free, no key) via @imgly — works
|
||||
// even when Gemini is unavailable. Other ops go through the Gemini route.
|
||||
// even with no backend. Other ops go through the /api/ai/edit route.
|
||||
if (op.type === 'remove-background') {
|
||||
// Prefer the RunPod U²-Net endpoint when enabled; fall back to in-browser
|
||||
// @imgly if it's off or the request fails, so this always produces a result.
|
||||
if (process.env.NEXT_PUBLIC_APG_RUNPOD_BG === 'true') {
|
||||
try {
|
||||
const { data, mimeType } = imageToBase64(image, 1600);
|
||||
const res = await fetch('/api/ai/edit', {
|
||||
method: 'POST',
|
||||
headers: { 'content-type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
imageBase64: data,
|
||||
mimeType,
|
||||
op: { type: 'remove-background' },
|
||||
params: {},
|
||||
}),
|
||||
});
|
||||
if (res.ok) {
|
||||
const { imageBase64: out, mimeType: outMime } = (await res.json()) as {
|
||||
imageBase64: string;
|
||||
mimeType?: string;
|
||||
};
|
||||
return base64ToBlob(out, outMime || 'image/png');
|
||||
}
|
||||
} catch {
|
||||
/* fall through to the in-browser remover */
|
||||
}
|
||||
}
|
||||
const inputBlob = await canvasBlob(image, 1600);
|
||||
const { removeBackground } = await import('@imgly/background-removal');
|
||||
return removeBackground(inputBlob, { output: { format: 'image/png' } });
|
||||
}
|
||||
|
||||
const { data, mimeType } = imageToBase64(image, 1280);
|
||||
// Outpaint / expand-canvas: pad the image with a neutral border, mark that
|
||||
// border WHITE in the mask, and run it through the same inpaint path as
|
||||
// generative-fill — no extra backend route needed.
|
||||
if (op.type === 'outpaint') {
|
||||
const { imageBase64: padded, maskBase64: border } = padForOutpaint(
|
||||
image,
|
||||
typeof op.factor === 'number' ? op.factor : 1.5,
|
||||
);
|
||||
const outParams: Record<string, unknown> = {
|
||||
strength: typeof op.strength === 'number' ? op.strength : 0.85,
|
||||
};
|
||||
const outRes = await fetch('/api/ai/edit', {
|
||||
method: 'POST',
|
||||
headers: { 'content-type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
imageBase64: padded,
|
||||
mimeType: 'image/png',
|
||||
op: {
|
||||
type: 'generative-fill',
|
||||
prompt:
|
||||
op.prompt ||
|
||||
'Extend and continue the scene naturally, matching lighting, colors and perspective.',
|
||||
},
|
||||
maskBase64: border,
|
||||
params: outParams,
|
||||
}),
|
||||
});
|
||||
if (!outRes.ok) {
|
||||
const err = (await outRes.json().catch(() => ({}))) as { error?: string };
|
||||
throw new Error(err.error || `AI request failed (${outRes.status}).`);
|
||||
}
|
||||
const outJson = (await outRes.json()) as { imageBase64: string; mimeType?: string };
|
||||
return base64ToBlob(outJson.imageBase64, outJson.mimeType || 'image/png');
|
||||
}
|
||||
|
||||
const { data, mimeType, width, height } = imageToBase64(image, 1280);
|
||||
// SD 3.5 masked ops (#9) carry an ImageData mask — rasterize it to a PNG
|
||||
// matched to the (downscaled) image dims, and strip it from the op since
|
||||
// ImageData is not JSON-serializable.
|
||||
const maskBase64 = 'mask' in op ? maskToBase64(op.mask, width, height) : undefined;
|
||||
const wireOp: Record<string, unknown> = { type: op.type };
|
||||
if ('prompt' in op && typeof op.prompt === 'string') wireOp.prompt = op.prompt;
|
||||
if ('factor' in op && typeof op.factor === 'number') wireOp.factor = op.factor;
|
||||
|
||||
// Forward the "edit strength" slider (0..1) so the backend can scale the edit.
|
||||
const params: Record<string, unknown> = {};
|
||||
if ('strength' in op && typeof op.strength === 'number') params.strength = op.strength;
|
||||
|
||||
const res = await fetch('/api/ai/edit', {
|
||||
method: 'POST',
|
||||
headers: { 'content-type': 'application/json' },
|
||||
body: JSON.stringify({ imageBase64: data, mimeType, op }),
|
||||
body: JSON.stringify({ imageBase64: data, mimeType, op: wireOp, maskBase64, params }),
|
||||
});
|
||||
if (!res.ok) {
|
||||
const err = (await res.json().catch(() => ({}))) as { error?: string };
|
||||
@@ -53,27 +139,60 @@ export function createDemoAIProvider(): AIProvider {
|
||||
};
|
||||
return base64ToBlob(imageBase64, outMime || 'image/png');
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/** Draw an image to a canvas (downscaled) and return base64 JPEG (no data: prefix). */
|
||||
function imageToBase64(
|
||||
image: ImageBitmap | HTMLImageElement,
|
||||
maxDim: number,
|
||||
): { data: string; mimeType: string } {
|
||||
const w = (image as HTMLImageElement).naturalWidth || (image as ImageBitmap).width;
|
||||
const h = (image as HTMLImageElement).naturalHeight || (image as ImageBitmap).height;
|
||||
const scale = Math.min(1, maxDim / Math.max(w, h));
|
||||
const cw = Math.max(1, Math.round(w * scale));
|
||||
const ch = Math.max(1, Math.round(h * scale));
|
||||
const canvas = document.createElement('canvas');
|
||||
canvas.width = cw;
|
||||
canvas.height = ch;
|
||||
const ctx = canvas.getContext('2d');
|
||||
if (!ctx) throw new Error('Canvas not supported.');
|
||||
ctx.drawImage(image as CanvasImageSource, 0, 0, cw, ch);
|
||||
const dataUrl = canvas.toDataURL('image/jpeg', 0.9);
|
||||
return { data: dataUrl.split(',')[1] ?? '', mimeType: 'image/jpeg' };
|
||||
// Voice annotation: record → (optional denoise) → transcribe. Both proxy
|
||||
// through server routes so the RunPod key stays server-side.
|
||||
async transcribeAudio(audioBase64) {
|
||||
const res = await fetch('/api/ai/transcribe', {
|
||||
method: 'POST',
|
||||
headers: { 'content-type': 'application/json' },
|
||||
body: JSON.stringify({ audio: audioBase64 }),
|
||||
});
|
||||
if (!res.ok) {
|
||||
const err = (await res.json().catch(() => ({}))) as { error?: string };
|
||||
throw new Error(err.error || `Transcription failed (${res.status}).`);
|
||||
}
|
||||
const { transcript } = (await res.json()) as { transcript?: string };
|
||||
return (transcript ?? '').trim();
|
||||
},
|
||||
|
||||
async denoiseAudio(audioBase64) {
|
||||
const res = await fetch('/api/ai/denoise', {
|
||||
method: 'POST',
|
||||
headers: { 'content-type': 'application/json' },
|
||||
body: JSON.stringify({ audio: audioBase64 }),
|
||||
});
|
||||
if (!res.ok) {
|
||||
const err = (await res.json().catch(() => ({}))) as { error?: string };
|
||||
throw new Error(err.error || `Denoise failed (${res.status}).`);
|
||||
}
|
||||
const { audio } = (await res.json()) as { audio?: string };
|
||||
return audio ?? audioBase64;
|
||||
},
|
||||
|
||||
// Camera-tilt estimation is opt-in (needs the RunPod tilt endpoint deployed);
|
||||
// gate it so the editor's Auto-straighten button only appears when configured.
|
||||
estimateTilt:
|
||||
process.env.NEXT_PUBLIC_APG_RUNPOD_TILT === 'true'
|
||||
? async (_item, image) => {
|
||||
const { data } = imageToBase64(image, 1024);
|
||||
const res = await fetch('/api/ai/tilt', {
|
||||
method: 'POST',
|
||||
headers: { 'content-type': 'application/json' },
|
||||
body: JSON.stringify({ image: data }),
|
||||
});
|
||||
if (!res.ok) {
|
||||
const err = (await res.json().catch(() => ({}))) as { error?: string };
|
||||
throw new Error(err.error || `Tilt estimate failed (${res.status}).`);
|
||||
}
|
||||
return (await res.json()) as {
|
||||
rollDegrees: number;
|
||||
pitchDegrees: number;
|
||||
fovDegrees: number;
|
||||
};
|
||||
}
|
||||
: undefined,
|
||||
};
|
||||
}
|
||||
|
||||
/** Draw an image to a canvas (downscaled) and return a JPEG Blob. */
|
||||
@@ -100,3 +219,55 @@ function base64ToBlob(base64: string, mime: string): Blob {
|
||||
for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i);
|
||||
return new Blob([bytes], { type: mime });
|
||||
}
|
||||
|
||||
/**
|
||||
* Pad an image with a neutral border for outpaint and return {paddedImage, mask}
|
||||
* as base64 PNG — the border is WHITE in the mask (regenerate), the original
|
||||
* image area BLACK (keep). Capped at 1280px on the long side.
|
||||
*/
|
||||
function padForOutpaint(
|
||||
image: ImageBitmap | HTMLImageElement,
|
||||
factor: number,
|
||||
): { imageBase64: string; maskBase64: string } {
|
||||
const w = (image as HTMLImageElement).naturalWidth || (image as ImageBitmap).width;
|
||||
const h = (image as HTMLImageElement).naturalHeight || (image as ImageBitmap).height;
|
||||
const f = Math.max(1.1, Math.min(2, factor));
|
||||
const maxDim = 1280;
|
||||
let pw = Math.round(w * f);
|
||||
let ph = Math.round(h * f);
|
||||
const scale = Math.min(1, maxDim / Math.max(pw, ph));
|
||||
pw = Math.max(16, Math.round(pw * scale));
|
||||
ph = Math.max(16, Math.round(ph * scale));
|
||||
const iw = Math.max(1, Math.round(w * scale));
|
||||
const ih = Math.max(1, Math.round(h * scale));
|
||||
const ox = Math.floor((pw - iw) / 2);
|
||||
const oy = Math.floor((ph - ih) / 2);
|
||||
|
||||
const imgCanvas = document.createElement('canvas');
|
||||
imgCanvas.width = pw;
|
||||
imgCanvas.height = ph;
|
||||
const ictx = imgCanvas.getContext('2d');
|
||||
if (!ictx) throw new Error('Canvas not supported.');
|
||||
// Fill the new border with a blurred, stretched copy of the photo so the model
|
||||
// has real color/context to continue from — flat gray gives it nothing and it
|
||||
// leaves the border gray. Then draw the original sharp on top, centered.
|
||||
ictx.filter = 'blur(28px)';
|
||||
ictx.drawImage(image as CanvasImageSource, 0, 0, pw, ph);
|
||||
ictx.filter = 'none';
|
||||
ictx.drawImage(image as CanvasImageSource, ox, oy, iw, ih);
|
||||
|
||||
const maskCanvas = document.createElement('canvas');
|
||||
maskCanvas.width = pw;
|
||||
maskCanvas.height = ph;
|
||||
const mctx = maskCanvas.getContext('2d');
|
||||
if (!mctx) throw new Error('Canvas not supported.');
|
||||
mctx.fillStyle = '#ffffff';
|
||||
mctx.fillRect(0, 0, pw, ph);
|
||||
mctx.fillStyle = '#000000';
|
||||
mctx.fillRect(ox, oy, iw, ih);
|
||||
|
||||
return {
|
||||
imageBase64: imgCanvas.toDataURL('image/png').split(',')[1] ?? '',
|
||||
maskBase64: maskCanvas.toDataURL('image/png').split(',')[1] ?? '',
|
||||
};
|
||||
}
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
/**
|
||||
* Shared client-side image encoding helpers for the AI providers. Downscale an
|
||||
* image to a JPEG base64 (keeping the exact output dims), and rasterize a mask
|
||||
* to a PNG base64 matched to those same dims — SD 3.5 (#9) requires the image
|
||||
* and mask to be identical pixel sizes.
|
||||
*/
|
||||
|
||||
export interface EncodedImage {
|
||||
/** base64 JPEG (no data: prefix). */
|
||||
data: string;
|
||||
mimeType: string;
|
||||
/** Actual pixel dims of the encoded image (after downscale). */
|
||||
width: number;
|
||||
height: number;
|
||||
}
|
||||
|
||||
/** Draw an image to a downscaled canvas and return base64 JPEG + its dims. */
|
||||
export function imageToBase64(image: ImageBitmap | HTMLImageElement, maxDim: number): EncodedImage {
|
||||
const w = (image as HTMLImageElement).naturalWidth || (image as ImageBitmap).width;
|
||||
const h = (image as HTMLImageElement).naturalHeight || (image as ImageBitmap).height;
|
||||
const scale = Math.min(1, maxDim / Math.max(w, h));
|
||||
const cw = Math.max(1, Math.round(w * scale));
|
||||
const ch = Math.max(1, Math.round(h * scale));
|
||||
const canvas = document.createElement('canvas');
|
||||
canvas.width = cw;
|
||||
canvas.height = ch;
|
||||
const ctx = canvas.getContext('2d');
|
||||
if (!ctx) throw new Error('Canvas not supported.');
|
||||
ctx.drawImage(image as CanvasImageSource, 0, 0, cw, ch);
|
||||
const dataUrl = canvas.toDataURL('image/jpeg', 0.9);
|
||||
return { data: dataUrl.split(',')[1] ?? '', mimeType: 'image/jpeg', width: cw, height: ch };
|
||||
}
|
||||
|
||||
/**
|
||||
* Rasterize a mask (ImageData, white = region to regenerate) to a PNG base64
|
||||
* scaled to targetW×targetH so it matches the encoded image exactly.
|
||||
*/
|
||||
export function maskToBase64(mask: ImageData, targetW: number, targetH: number): string {
|
||||
const tmp = document.createElement('canvas');
|
||||
tmp.width = mask.width;
|
||||
tmp.height = mask.height;
|
||||
const tctx = tmp.getContext('2d');
|
||||
if (!tctx) throw new Error('Canvas not supported.');
|
||||
tctx.putImageData(mask, 0, 0);
|
||||
|
||||
const out = document.createElement('canvas');
|
||||
out.width = targetW;
|
||||
out.height = targetH;
|
||||
const octx = out.getContext('2d');
|
||||
if (!octx) throw new Error('Canvas not supported.');
|
||||
// Nearest-neighbour, not bilinear — keep the mask strictly binary so SD 3.5 gets
|
||||
// crisp white(regenerate)/black(keep) edges instead of an anti-aliased grey halo.
|
||||
octx.imageSmoothingEnabled = false;
|
||||
octx.drawImage(tmp, 0, 0, targetW, targetH);
|
||||
return out.toDataURL('image/png').split(',')[1] ?? '';
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
import type { AIProvider, DetectedObject, MediaItem } from '@photo-gallery/sdk';
|
||||
|
||||
import { imageToBase64 } from './imageEncode';
|
||||
|
||||
/**
|
||||
* Object detection backed by the RunPod YOLO construction-material classifier
|
||||
* (#1), proxied through /api/ai/classify so the RunPod key stays server-side.
|
||||
*
|
||||
* Opt-in via `NEXT_PUBLIC_APG_RUNPOD_DETECT=true` (see createDemoAIProvider). If
|
||||
* the server call fails (endpoint down / not configured) it falls back to the
|
||||
* given in-browser provider (COCO-SSD), so detection never hard-fails.
|
||||
*/
|
||||
export function createRunpodYoloProvider(fallback: AIProvider): Pick<AIProvider, 'detectObjects'> {
|
||||
return {
|
||||
async detectObjects(item: MediaItem, image): Promise<DetectedObject[]> {
|
||||
try {
|
||||
const { data, mimeType, width, height } = imageToBase64(image, 1280);
|
||||
const res = await fetch('/api/ai/classify', {
|
||||
method: 'POST',
|
||||
headers: { 'content-type': 'application/json' },
|
||||
body: JSON.stringify({ imageBase64: data, mimeType, width, height }),
|
||||
});
|
||||
if (!res.ok) throw new Error(`classify failed (${res.status})`);
|
||||
const { objects } = (await res.json()) as { objects?: DetectedObject[] };
|
||||
if (Array.isArray(objects)) return objects;
|
||||
throw new Error('classify returned no objects');
|
||||
} catch {
|
||||
return fallback.detectObjects ? fallback.detectObjects(item, image) : [];
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
/**
|
||||
* base64 / data-URI helpers shared by the RunPod endpoint wrappers.
|
||||
* Different endpoints name their image field differently and some prefix a
|
||||
* `data:` URI — these helpers normalize both.
|
||||
*/
|
||||
|
||||
/** "data:image/png;base64,XXXX" -> "XXXX" (leaves a bare base64 string untouched). */
|
||||
export function stripDataUri(s: string): string {
|
||||
if (!s.startsWith('data:')) return s;
|
||||
const i = s.indexOf(',');
|
||||
return i === -1 ? s : s.slice(i + 1);
|
||||
}
|
||||
|
||||
/** Wrap a bare base64 string in a data: URI. */
|
||||
export function toDataUri(b64: string, mime: string): string {
|
||||
return `data:${mime};base64,${stripDataUri(b64)}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Pull the first image-like base64 string out of a RunPod endpoint's `output`,
|
||||
* regardless of which field name it used. Handles the common shapes:
|
||||
* "iVBOR..." (raw string)
|
||||
* { image: "..." } / { image_png } / { image_base64 }
|
||||
* { images: ["..."] } (array)
|
||||
* { output: { image: "..." } } (nested)
|
||||
* Mirrors the A1111 normalization already used for the local SD backend.
|
||||
*/
|
||||
export function pickOutputImage(output: unknown): string {
|
||||
const s = findImageString(output, false);
|
||||
if (!s) throw new Error('RunPod output did not contain an image.');
|
||||
return stripDataUri(s);
|
||||
}
|
||||
|
||||
/** Keys whose name implies the value IS the image — any non-empty string is accepted. */
|
||||
const IMAGE_KEYS = ['image_png', 'image', 'image_base64', 'images'] as const;
|
||||
/** Generic wrapper keys — a string here must actually look like image data. */
|
||||
const CONTAINER_KEYS = ['output', 'result', 'data'] as const;
|
||||
|
||||
/** A base64 image payload is long; a status/id string ("success", "job-abc") is short. */
|
||||
function looksLikeImageData(s: string): boolean {
|
||||
if (s.startsWith('data:image/')) return true;
|
||||
return s.length >= 256 && /^[A-Za-z0-9+/=\s]+$/.test(s.slice(0, 256));
|
||||
}
|
||||
|
||||
/**
|
||||
* `strict` is true when we descended through a generic wrapper key (output/result/
|
||||
* data), where a bare string could be a status/id rather than an image — so it must
|
||||
* pass `looksLikeImageData`. Under an explicit image key (or at top level) any
|
||||
* non-empty string is taken as the image.
|
||||
*/
|
||||
function findImageString(value: unknown, strict: boolean, depth = 0): string | undefined {
|
||||
if (depth > 5) return undefined;
|
||||
if (typeof value === 'string') {
|
||||
if (!value) return undefined;
|
||||
return !strict || looksLikeImageData(value) ? value : undefined;
|
||||
}
|
||||
if (Array.isArray(value)) {
|
||||
for (const el of value) {
|
||||
const s = findImageString(el, strict, depth + 1);
|
||||
if (s) return s;
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
if (value && typeof value === 'object') {
|
||||
const o = value as Record<string, unknown>;
|
||||
for (const key of IMAGE_KEYS) {
|
||||
if (key in o) {
|
||||
const s = findImageString(o[key], false, depth + 1);
|
||||
if (s) return s;
|
||||
}
|
||||
}
|
||||
for (const key of CONTAINER_KEYS) {
|
||||
if (key in o) {
|
||||
const s = findImageString(o[key], true, depth + 1);
|
||||
if (s) return s;
|
||||
}
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
/**
|
||||
* Low-level RunPod transport. Every model in the spec runs as its own RunPod
|
||||
* endpoint; this handles the common request envelope, bearer auth, and both
|
||||
* response modes:
|
||||
*
|
||||
* - `/runsync` (preferred): the job runs synchronously and the body already
|
||||
* contains `output` (or IS the output for a custom handler).
|
||||
* - `/run`: returns `{ id, status }`; we poll `/status/{id}` until the job
|
||||
* reaches a terminal state or the time budget (kept under Vercel's 60s
|
||||
* function cap) is exhausted.
|
||||
*
|
||||
* RUNPOD_API_KEY is read here so callers never handle the secret directly.
|
||||
*/
|
||||
|
||||
export class RunpodError extends Error {
|
||||
status: number;
|
||||
constructor(message: string, status = 502) {
|
||||
super(message);
|
||||
this.name = 'RunpodError';
|
||||
this.status = status;
|
||||
}
|
||||
}
|
||||
|
||||
export interface RunpodCallOpts {
|
||||
/** Human label used in error messages, e.g. "sd-inpaint". */
|
||||
name: string;
|
||||
/** Full endpoint URL from the per-model env var (…/runsync or …/run). */
|
||||
url: string;
|
||||
input: Record<string, unknown>;
|
||||
/** Total budget for the whole call incl. polling. Default 55s (< Vercel 60s). */
|
||||
timeoutMs?: number;
|
||||
pollIntervalMs?: number;
|
||||
}
|
||||
|
||||
interface RunpodEnvelope {
|
||||
id?: string;
|
||||
status?: string;
|
||||
output?: unknown;
|
||||
error?: unknown;
|
||||
}
|
||||
|
||||
export async function runpodCall<TOut = unknown>(opts: RunpodCallOpts): Promise<TOut> {
|
||||
const { name, url, input } = opts;
|
||||
const timeoutMs = opts.timeoutMs ?? 55_000;
|
||||
const pollIntervalMs = opts.pollIntervalMs ?? 1500;
|
||||
|
||||
const key = process.env.RUNPOD_API_KEY;
|
||||
if (!key) throw new RunpodError('RUNPOD_API_KEY is not set.', 500);
|
||||
|
||||
const deadline = Date.now() + timeoutMs;
|
||||
const authHeaders = { authorization: `Bearer ${key}` };
|
||||
|
||||
let res: Response;
|
||||
try {
|
||||
res = await fetch(url, {
|
||||
method: 'POST',
|
||||
headers: { 'content-type': 'application/json', ...authHeaders },
|
||||
body: JSON.stringify({ input }),
|
||||
signal: AbortSignal.timeout(timeoutMs),
|
||||
});
|
||||
} catch {
|
||||
throw new RunpodError(`Could not reach RunPod (${name}).`, 502);
|
||||
}
|
||||
if (!res.ok) {
|
||||
const detail = (await res.text().catch(() => '')).slice(0, 160);
|
||||
throw new RunpodError(`RunPod ${name} error (${res.status}). ${detail}`.trim(), 502);
|
||||
}
|
||||
|
||||
const data = (await res.json().catch(() => null)) as RunpodEnvelope | null;
|
||||
if (!data || typeof data !== 'object') {
|
||||
throw new RunpodError(`RunPod ${name} returned an invalid response.`, 502);
|
||||
}
|
||||
|
||||
// Terminal failure reported in a job envelope.
|
||||
if (data.status === 'FAILED' || data.status === 'CANCELLED') {
|
||||
throw new RunpodError(`RunPod ${name} job ${data.status.toLowerCase()}.`, 502);
|
||||
}
|
||||
// No job id → this is not an async envelope; the body itself is the output.
|
||||
// Covers custom /runsync handlers that return their result directly, even when
|
||||
// it carries a `status` field (e.g. "success").
|
||||
if (data.id === undefined) {
|
||||
return (data.output !== undefined ? data.output : data) as TOut;
|
||||
}
|
||||
// Job envelope that already carries a completed/inline output.
|
||||
if (data.output !== undefined && (data.status === undefined || data.status === 'COMPLETED')) {
|
||||
return data.output as TOut;
|
||||
}
|
||||
|
||||
// Async: poll /status/{id} until COMPLETED / FAILED / the time budget runs out.
|
||||
// Each wait + fetch is clamped to the remaining budget so the whole call stays
|
||||
// under `timeoutMs` (kept below Vercel's 60s function cap).
|
||||
const statusUrl = url.replace(/\/run(sync)?(\/?)$/, `/status/${data.id}`);
|
||||
for (;;) {
|
||||
const remaining = deadline - Date.now();
|
||||
if (remaining < 500) break; // not enough budget for another round
|
||||
await sleep(Math.min(pollIntervalMs, remaining));
|
||||
const left = deadline - Date.now();
|
||||
if (left <= 0) break;
|
||||
let sres: Response;
|
||||
try {
|
||||
sres = await fetch(statusUrl, {
|
||||
headers: authHeaders,
|
||||
signal: AbortSignal.timeout(Math.min(10_000, Math.max(1000, left))),
|
||||
});
|
||||
} catch {
|
||||
continue; // transient — keep polling until the deadline
|
||||
}
|
||||
if (!sres.ok) continue;
|
||||
const sdata = (await sres.json().catch(() => null)) as RunpodEnvelope | null;
|
||||
if (!sdata) continue;
|
||||
if (sdata.status === 'COMPLETED') return sdata.output as TOut;
|
||||
if (sdata.status === 'FAILED' || sdata.status === 'CANCELLED') {
|
||||
throw new RunpodError(`RunPod ${name} job ${sdata.status.toLowerCase()}.`, 502);
|
||||
}
|
||||
}
|
||||
throw new RunpodError(`RunPod ${name} timed out (raise the endpoint's speed or use /runsync).`, 504);
|
||||
}
|
||||
|
||||
function sleep(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
@@ -0,0 +1,283 @@
|
||||
/**
|
||||
* Typed wrappers for each RunPod endpoint in `model-api-spec.docx`.
|
||||
* Every function reads its own `RUNPOD_*_URL` env var, sends the exact `input`
|
||||
* contract the spec documents, and normalizes the response.
|
||||
*
|
||||
* Server-side only. Missing/invalid URLs throw a RunpodError(500) so an
|
||||
* unconfigured op surfaces as a clear message in the editor rather than a crash.
|
||||
*/
|
||||
|
||||
import { stripDataUri, pickOutputImage } from './base64';
|
||||
import { RunpodError, runpodCall } from './client';
|
||||
import type {
|
||||
Img2ImgReq,
|
||||
InpaintReq,
|
||||
RunpodDetection,
|
||||
RunpodImageResult,
|
||||
TiltResult,
|
||||
TranscriptResult,
|
||||
} from './types';
|
||||
|
||||
function envNum(v: string | undefined, fallback: number): number {
|
||||
// Treat a blank/whitespace env var as unset — Number('') is 0 (finite), which
|
||||
// would otherwise send e.g. strength:0 for `RUNPOD_SD_STRENGTH=`.
|
||||
if (v == null || v.trim() === '') return fallback;
|
||||
const n = Number(v);
|
||||
return Number.isFinite(n) ? n : fallback;
|
||||
}
|
||||
|
||||
function endpointUrl(envVar: string): string {
|
||||
const url = process.env[envVar];
|
||||
if (!url || !/^https?:\/\//i.test(url)) {
|
||||
throw new RunpodError(`${envVar} is not set (or is not an http(s) URL).`, 500);
|
||||
}
|
||||
return url;
|
||||
}
|
||||
|
||||
/** Run an image-in/image-out endpoint and normalize the result to base64 PNG. */
|
||||
async function imageOp(
|
||||
name: string,
|
||||
envVar: string,
|
||||
input: Record<string, unknown>,
|
||||
): Promise<RunpodImageResult> {
|
||||
const output = await runpodCall({ name, url: endpointUrl(envVar), input });
|
||||
return { imageBase64: pickOutputImage(output), mimeType: 'image/png' };
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Image endpoints
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/** #6 Background removal (U²-Net via rembg). model: u2net | u2netp | u2net_human_seg */
|
||||
export function rpRemoveBackground(imageB64: string, model?: string): Promise<RunpodImageResult> {
|
||||
return imageOp('background-removal', 'RUNPOD_BG_REMOVE_URL', {
|
||||
task: 'remove-bg',
|
||||
image: imageB64,
|
||||
...(model ? { model_name: model } : {}),
|
||||
});
|
||||
}
|
||||
|
||||
/** #7 Real-ESRGAN enhance/upscale (RealESRGAN_x4plus). The caller's scale (from the
|
||||
* op / restore pass) is authoritative — it is not overridden by any env default. */
|
||||
export function rpUpscale(imageB64: string, scale: 2 | 4, faceEnhance = false): Promise<RunpodImageResult> {
|
||||
return imageOp('upscale', 'RUNPOD_UPSCALE_URL', {
|
||||
task: 'upscale',
|
||||
image: imageB64,
|
||||
scale,
|
||||
face_enhance: faceEnhance,
|
||||
});
|
||||
}
|
||||
|
||||
/** #8 DDColor B&W → colorize. */
|
||||
export function rpColorize(imageB64: string, inputSize?: number): Promise<RunpodImageResult> {
|
||||
return imageOp('colorize', 'RUNPOD_COLORIZE_URL', {
|
||||
image: imageB64,
|
||||
...(inputSize ? { input_size: inputSize } : {}),
|
||||
});
|
||||
}
|
||||
|
||||
/** #9 SD 3.5 masked inpainting (sky fix / eraser / fill). Also #11 outpaint (pre-padded). */
|
||||
export function rpInpaint(p: InpaintReq): Promise<RunpodImageResult> {
|
||||
const input: Record<string, unknown> = {
|
||||
task: 'inpaint',
|
||||
image: p.imageB64,
|
||||
mask: p.maskB64,
|
||||
prompt: p.prompt,
|
||||
strength: p.strength ?? envNum(process.env.RUNPOD_SD_STRENGTH, 0.8),
|
||||
guidance_scale: p.guidanceScale ?? envNum(process.env.RUNPOD_SD_GUIDANCE, 7),
|
||||
num_inference_steps: p.steps ?? envNum(process.env.RUNPOD_SD_STEPS, 35),
|
||||
};
|
||||
const negative = p.negativePrompt ?? process.env.RUNPOD_SD_NEGATIVE_PROMPT;
|
||||
if (negative) input.negative_prompt = negative;
|
||||
if (p.seed != null) input.seed = p.seed;
|
||||
return imageOp('sd-inpaint', 'RUNPOD_SD_INPAINT_URL', input);
|
||||
}
|
||||
|
||||
/** #10 SD 3.5 general prompt edit (img2img, no mask). */
|
||||
export function rpImg2Img(p: Img2ImgReq): Promise<RunpodImageResult> {
|
||||
const input: Record<string, unknown> = {
|
||||
task: 'img2img',
|
||||
image: p.imageB64,
|
||||
prompt: p.prompt,
|
||||
strength: p.strength ?? envNum(process.env.RUNPOD_SD_STRENGTH, 0.6),
|
||||
guidance_scale: p.guidanceScale ?? envNum(process.env.RUNPOD_SD_GUIDANCE, 7),
|
||||
num_inference_steps: p.steps ?? envNum(process.env.RUNPOD_SD_STEPS, 35),
|
||||
};
|
||||
const negative = p.negativePrompt ?? process.env.RUNPOD_SD_NEGATIVE_PROMPT;
|
||||
if (negative) input.negative_prompt = negative;
|
||||
if (p.seed != null) input.seed = p.seed;
|
||||
return imageOp('sd-img2img', 'RUNPOD_SD_IMG2IMG_URL', input);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// #1 YOLO detection → SDK DetectedObject shape (box as fractions 0..1)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export async function rpDetect(
|
||||
imageB64: string,
|
||||
width: number,
|
||||
height: number,
|
||||
): Promise<RunpodDetection[]> {
|
||||
const output = await runpodCall<unknown>({
|
||||
name: 'yolo-detect',
|
||||
url: endpointUrl('RUNPOD_YOLO_URL'),
|
||||
input: { image: imageB64, task: 'detect' },
|
||||
});
|
||||
return normalizeDetections(output, width, height);
|
||||
}
|
||||
|
||||
function extractDetectionArray(output: unknown): unknown[] {
|
||||
if (Array.isArray(output)) return output;
|
||||
if (output && typeof output === 'object') {
|
||||
const o = output as Record<string, unknown>;
|
||||
for (const key of ['detections', 'predictions', 'objects', 'results', 'boxes']) {
|
||||
if (Array.isArray(o[key])) return o[key] as unknown[];
|
||||
}
|
||||
}
|
||||
return [];
|
||||
}
|
||||
|
||||
/** First non-empty STRING among the args (numbers ignored — a numeric `class` is an index, not a name). */
|
||||
function firstLabel(...vals: unknown[]): string | undefined {
|
||||
for (const v of vals) {
|
||||
if (typeof v === 'string' && v.trim()) return v.trim();
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function normalizeDetections(output: unknown, width: number, height: number): RunpodDetection[] {
|
||||
const out: RunpodDetection[] = [];
|
||||
for (const raw of extractDetectionArray(output)) {
|
||||
if (!raw || typeof raw !== 'object') continue;
|
||||
const o = raw as Record<string, unknown>;
|
||||
// Prefer a human-readable name (ultralytics tojson puts the string in `name`
|
||||
// and a numeric index in `class`); fall back to class_<id>. Using firstLabel
|
||||
// (not `??`) also means an explicit empty-string label doesn't get kept + dropped.
|
||||
const classId = o.class_id ?? (typeof o.class === 'number' ? o.class : undefined);
|
||||
const label = (
|
||||
firstLabel(o.label, o.name, o.class_name, typeof o.class === 'string' ? o.class : undefined) ??
|
||||
(classId != null ? `class_${classId}` : 'object')
|
||||
).toLowerCase();
|
||||
const confidence = Number(o.confidence ?? o.score ?? o.conf ?? 0) || 0;
|
||||
const box = normalizeBox(o, width, height);
|
||||
if (box && label) out.push({ label, confidence, box });
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
function num(v: unknown): number | null {
|
||||
const n = Number(v);
|
||||
return Number.isFinite(n) ? n : null;
|
||||
}
|
||||
|
||||
function asNum4(v: unknown): [number, number, number, number] | null {
|
||||
if (!Array.isArray(v) || v.length < 4) return null;
|
||||
const a = num(v[0]);
|
||||
const b = num(v[1]);
|
||||
const c = num(v[2]);
|
||||
const d = num(v[3]);
|
||||
return a === null || b === null || c === null || d === null ? null : [a, b, c, d];
|
||||
}
|
||||
|
||||
/**
|
||||
* Normalize a detection box to {x, y, width, height} as fractions 0..1 of the
|
||||
* image, from whatever shape the endpoint emits:
|
||||
* - `xyxy: [x1,y1,x2,y2]` (ultralytics) and generic `box: [...]` → corner form
|
||||
* - `xywh: [...]` and COCO `bbox: [x,y,w,h]` → x/y/width/height form
|
||||
* - object `{x1,y1,x2,y2}` / `{left,top,right,bottom}` (ultralytics tojson) → corners
|
||||
* - object `{x,y,width,height}` → x/y/width/height
|
||||
* Pixel values (any component > 1) are divided by the image dims; already-
|
||||
* normalized fractions pass through. VERIFY the shape against your live #1 endpoint.
|
||||
*/
|
||||
function normalizeBox(o: Record<string, unknown>, width: number, height: number): RunpodDetection['box'] | null {
|
||||
const W = width || 1;
|
||||
const H = height || 1;
|
||||
const clamp01 = (n: number) => Math.max(0, Math.min(1, n));
|
||||
const frac = (x: number, y: number, w: number, h: number): RunpodDetection['box'] => {
|
||||
if (Math.max(Math.abs(x), Math.abs(y), Math.abs(w), Math.abs(h)) > 1) {
|
||||
x /= W;
|
||||
y /= H;
|
||||
w /= W;
|
||||
h /= H;
|
||||
}
|
||||
return { x: clamp01(x), y: clamp01(y), width: clamp01(w), height: clamp01(h) };
|
||||
};
|
||||
const fromXyxy = (x1: number, y1: number, x2: number, y2: number) => frac(x1, y1, x2 - x1, y2 - y1);
|
||||
|
||||
// 1. Array boxes, interpreted by which key holds them.
|
||||
const xyxyArr = asNum4(o.xyxy);
|
||||
if (xyxyArr) return fromXyxy(xyxyArr[0], xyxyArr[1], xyxyArr[2], xyxyArr[3]);
|
||||
const xywhArr = asNum4(o.xywh) ?? asNum4(o.bbox); // COCO `bbox` is [x,y,w,h]
|
||||
if (xywhArr) return frac(xywhArr[0], xywhArr[1], xywhArr[2], xywhArr[3]);
|
||||
const boxArr = asNum4(o.box); // generic array box → assume corner form
|
||||
if (boxArr) return fromXyxy(boxArr[0], boxArr[1], boxArr[2], boxArr[3]);
|
||||
|
||||
// 2. Object boxes (either nested under `box` or directly on the detection).
|
||||
const src = o.box && typeof o.box === 'object' && !Array.isArray(o.box) ? (o.box as Record<string, unknown>) : o;
|
||||
const x1 = num(src.x1 ?? src.left);
|
||||
const y1 = num(src.y1 ?? src.top);
|
||||
const x2 = num(src.x2 ?? src.right);
|
||||
const y2 = num(src.y2 ?? src.bottom);
|
||||
if (x1 !== null && y1 !== null && x2 !== null && y2 !== null) return fromXyxy(x1, y1, x2, y2);
|
||||
|
||||
const x = num(src.x);
|
||||
const y = num(src.y);
|
||||
const w = num(src.width);
|
||||
const h = num(src.height);
|
||||
if (x !== null && y !== null && w !== null && h !== null) return frac(x, y, w, h);
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Scaffold-only endpoints (typed + wired to env, but no route/UI yet).
|
||||
// See docs/AI-SETUP.md → "Not yet wired". Kept so the client covers all 15
|
||||
// endpoints and a future feature can call them without new plumbing.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/** #2 Camera tilt (DeepSingleImageCalibration). */
|
||||
export async function rpTilt(imageB64: string): Promise<TiltResult> {
|
||||
const o = await runpodCall<Record<string, unknown>>({
|
||||
name: 'tilt',
|
||||
url: endpointUrl('RUNPOD_TILT_URL'),
|
||||
input: { image: imageB64 },
|
||||
});
|
||||
return {
|
||||
rollDegrees: Number(o.roll_degrees ?? 0) || 0,
|
||||
pitchDegrees: Number(o.pitch_degrees ?? 0) || 0,
|
||||
fovDegrees: Number(o.fov_degrees ?? 0) || 0,
|
||||
};
|
||||
}
|
||||
|
||||
/** #3 Voice-to-text (Parakeet). Audio must be WAV 16kHz mono PCM16. */
|
||||
export async function rpTranscribe(
|
||||
audioB64: string,
|
||||
opts?: { language?: string; timestamps?: boolean; punctuation?: boolean },
|
||||
): Promise<TranscriptResult> {
|
||||
const o = await runpodCall<Record<string, unknown>>({
|
||||
name: 'transcribe',
|
||||
url: endpointUrl('RUNPOD_STT_URL'),
|
||||
input: { audio: audioB64, task: 'transcribe', ...(opts ?? {}) },
|
||||
});
|
||||
const rawSegments = Array.isArray(o.segments) ? o.segments : [];
|
||||
const segments = rawSegments.map((s) => {
|
||||
const seg = (s ?? {}) as Record<string, unknown>;
|
||||
return {
|
||||
text: String(seg.text ?? ''),
|
||||
startSec: Number(seg.start_sec ?? 0) || 0,
|
||||
endSec: Number(seg.end_sec ?? 0) || 0,
|
||||
};
|
||||
});
|
||||
return { transcript: String(o.transcript ?? ''), segments: segments.length ? segments : undefined };
|
||||
}
|
||||
|
||||
/** #12 Audio noise removal (RNNoise). Audio must be WAV 48kHz mono 16-bit PCM. */
|
||||
export async function rpDenoiseAudio(audioB64: string): Promise<{ audioB64: string }> {
|
||||
const o = await runpodCall<Record<string, unknown>>({
|
||||
name: 'audio-denoise',
|
||||
url: endpointUrl('RUNPOD_AUDIO_DENOISE_URL'),
|
||||
input: { audio: audioB64, task: 'denoise' },
|
||||
});
|
||||
const a = o.audio ?? o.output ?? '';
|
||||
return { audioB64: stripDataUri(String(a)) };
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
/**
|
||||
* Request/response types for the RunPod endpoints described in
|
||||
* `model-api-spec.docx`. Server-side only — imported by the /api/ai/* routes,
|
||||
* never by client/SDK code (which must not see a RunPod URL or key).
|
||||
*/
|
||||
|
||||
/** Normalized image result returned by every image endpoint (raw base64, no data: prefix). */
|
||||
export interface RunpodImageResult {
|
||||
imageBase64: string;
|
||||
mimeType: string;
|
||||
}
|
||||
|
||||
/** #9 SD 3.5 masked inpainting (also #11 outpaint, pre-padded). white in mask = regenerate. */
|
||||
export interface InpaintReq {
|
||||
imageB64: string;
|
||||
maskB64: string;
|
||||
prompt: string;
|
||||
negativePrompt?: string;
|
||||
strength?: number;
|
||||
guidanceScale?: number;
|
||||
steps?: number;
|
||||
seed?: number;
|
||||
}
|
||||
|
||||
/** #10 SD 3.5 general prompt edit (img2img, no mask). */
|
||||
export interface Img2ImgReq {
|
||||
imageB64: string;
|
||||
prompt: string;
|
||||
negativePrompt?: string;
|
||||
strength?: number;
|
||||
guidanceScale?: number;
|
||||
steps?: number;
|
||||
seed?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* #1 YOLO construction-material classifier, normalized to the SDK's
|
||||
* `DetectedObject` shape (box as fractions 0..1 of the image).
|
||||
*/
|
||||
export interface RunpodDetection {
|
||||
label: string;
|
||||
confidence: number;
|
||||
box: { x: number; y: number; width: number; height: number };
|
||||
}
|
||||
|
||||
/** #2 DeepSingleImageCalibration — camera tilt (scaffold; no route/UI yet). */
|
||||
export interface TiltResult {
|
||||
rollDegrees: number;
|
||||
pitchDegrees: number;
|
||||
fovDegrees: number;
|
||||
}
|
||||
|
||||
/** #3 Parakeet voice-to-text (scaffold; no route/UI yet). */
|
||||
export interface TranscriptSegment {
|
||||
text: string;
|
||||
startSec: number;
|
||||
endSec: number;
|
||||
}
|
||||
export interface TranscriptResult {
|
||||
transcript: string;
|
||||
segments?: TranscriptSegment[];
|
||||
}
|
||||
@@ -137,5 +137,8 @@ export function buildSeedPhotos(now = Date.now()): MediaInput[] {
|
||||
});
|
||||
}
|
||||
|
||||
return items;
|
||||
// TESTING: seed only the demo video(s) — start with an EMPTY photo library so you
|
||||
// can upload 1–2 images and watch them get analyzed. To restore the full demo,
|
||||
// change this back to `return items;`.
|
||||
return items.filter((it) => it.kind === 'video');
|
||||
}
|
||||
|
||||
@@ -10,7 +10,6 @@ import { type NextRequest, NextResponse } from 'next/server';
|
||||
*/
|
||||
export function middleware(request: NextRequest) {
|
||||
const nonce = Buffer.from(crypto.randomUUID()).toString('base64');
|
||||
const isDev = process.env.NODE_ENV === 'development';
|
||||
|
||||
const csp = [
|
||||
"default-src 'self'",
|
||||
@@ -18,7 +17,10 @@ export function middleware(request: NextRequest) {
|
||||
// cdn.jsdelivr.net = defense for tesseract.js' worker bootstrap (the load-bearing
|
||||
// path is the blob: worker + connect-src fetch; under 'strict-dynamic' this host is
|
||||
// ignored, but it covers non-strict-dynamic fallback browsers).
|
||||
`script-src 'self' 'nonce-${nonce}' 'strict-dynamic' 'wasm-unsafe-eval' https://cdn.jsdelivr.net${isDev ? " 'unsafe-eval'" : ''}`,
|
||||
// 'unsafe-eval' is required by onnxruntime-web (Remove Background @imgly, CLIP
|
||||
// semantic search) which evals a JS glue string — 'wasm-unsafe-eval' alone
|
||||
// doesn't cover it. Trade-off: needed for the in-browser ML to run in production.
|
||||
`script-src 'self' 'nonce-${nonce}' 'strict-dynamic' 'wasm-unsafe-eval' 'unsafe-eval' https://cdn.jsdelivr.net`,
|
||||
// Leaflet & Framer Motion inject inline styles; keep style inline allowed.
|
||||
"style-src 'self' 'unsafe-inline'",
|
||||
// *.supabase.co serves uploaded media from Supabase Storage public URLs.
|
||||
@@ -30,7 +32,7 @@ export function middleware(request: NextRequest) {
|
||||
// cdn.jsdelivr.net = face-api models + tesseract.js worker/WASM-core/traineddata (OCR) + onnxruntime WASM;
|
||||
// huggingface.co + *.hf.co = transformers.js CLIP model config + weights, which
|
||||
// redirect to HF's Xet CDN (e.g. us.aws.cdn.hf.co) — semantic search.
|
||||
"connect-src 'self' blob: data: https://*.tile.openstreetmap.org https://server.arcgisonline.com https://picsum.photos https://fastly.picsum.photos https://storage.googleapis.com https://*.supabase.co https://staticimgly.com https://cdn.jsdelivr.net https://huggingface.co https://*.huggingface.co https://*.hf.co",
|
||||
"connect-src 'self' blob: data: https://*.tile.openstreetmap.org https://nominatim.openstreetmap.org https://server.arcgisonline.com https://picsum.photos https://fastly.picsum.photos https://storage.googleapis.com https://*.supabase.co https://staticimgly.com https://cdn.jsdelivr.net https://huggingface.co https://*.huggingface.co https://*.hf.co",
|
||||
"worker-src 'self' blob:",
|
||||
"frame-ancestors 'none'",
|
||||
"base-uri 'self'",
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
# CRM Chat + Translate Pod (self-bootstrapping)
|
||||
|
||||
One RunPod **Pod** running **Qwen2.5-14B** (chat/assistant, GPU) + **NLLB-200-1.3B**
|
||||
(translation, CPU). It **auto-starts on every boot/resume** — you only Stop/Resume;
|
||||
the CRM team's URLs never change and nothing needs editing.
|
||||
|
||||
## One-time: create the Pod
|
||||
RunPod → **Pods → Deploy**:
|
||||
|
||||
| Setting | Value |
|
||||
|---|---|
|
||||
| **GPU** | **RTX 3090 (24 GB)** — On-Demand / Secure Cloud (not Spot) |
|
||||
| **Template** | RunPod PyTorch (has python3 + CUDA) |
|
||||
| **Network Volume** | create one, **~50 GB**, mount at **`/workspace`** ← holds models + deps so Stop/Resume keeps them |
|
||||
| **Container Disk** | **40 GB** |
|
||||
| **Expose HTTP Ports** (Edit Template) | `8000,11434` |
|
||||
| **Docker / Start Command** | see below |
|
||||
|
||||
**Start Command** (paste exactly — this is the whole setup):
|
||||
```
|
||||
bash -c "curl -fsSL https://raw.githubusercontent.com/Sumit-Pluto/photo_gallery/main/chat-pod/start.sh | bash"
|
||||
```
|
||||
|
||||
First boot downloads Qwen (~16 GB) + converts NLLB → **~15–20 min**. Every Resume
|
||||
after that: models already on the volume → **ready in ~1–2 min**, same URLs.
|
||||
|
||||
## The two URLs to share with the CRM team
|
||||
After deploy, each exposed port has a stable proxy URL (stable as long as you
|
||||
**Stop/Resume**, never Terminate):
|
||||
```
|
||||
Chat (Qwen, OpenAI-compatible): https://<POD_ID>-11434.proxy.runpod.net/v1/chat/completions
|
||||
Translate (NLLB): https://<POD_ID>-8000.proxy.runpod.net/translate
|
||||
```
|
||||
|
||||
## API for the CRM
|
||||
**Chat / summarize** (OpenAI format, supports `stream:true`):
|
||||
```json
|
||||
POST /v1/chat/completions
|
||||
{ "model": "qwen2.5:14b-instruct-q8_0", "stream": true,
|
||||
"messages": [{ "role": "user", "content": "Summarize this thread in English: ..." }] }
|
||||
```
|
||||
**Translate before send / on receive:**
|
||||
```json
|
||||
POST /translate
|
||||
{ "text": "When can you deliver the cement?", "source": "eng_Latn", "target": "hin_Deva" }
|
||||
-> { "translation": "..." }
|
||||
```
|
||||
|
||||
## FLORES language codes (NLLB)
|
||||
`eng_Latn` English · `hin_Deva` Hindi · `spa_Latn` Spanish · `fra_Latn` French ·
|
||||
`deu_Latn` German · `arb_Arab` Arabic · `zho_Hans` Chinese · `rus_Cyrl` Russian ·
|
||||
`por_Latn` Portuguese · `ben_Beng` Bengali · `tam_Taml` Tamil · `tel_Telu` Telugu ·
|
||||
`mar_Deva` Marathi · `guj_Gujr` Gujarati · `pan_Guru` Punjabi · `jpn_Jpan` Japanese ·
|
||||
`kor_Hang` Korean · `vie_Latn` Vietnamese · `ind_Latn` Indonesian
|
||||
(full list: NLLB-200 FLORES-200 codes.) Your CRM maps each user's language → its code.
|
||||
|
||||
## Security (do this)
|
||||
- **Only call these URLs from your CRM backend**, never the browser (the proxy URLs are public).
|
||||
- Optional: set env **`CHAT_POD_KEY`** on the pod → the translate API then requires
|
||||
`Authorization: Bearer <key>`. (Ollama has no built-in auth — keep it backend-only or front it with a gateway.)
|
||||
|
||||
## Daily operation (what you actually do)
|
||||
- **Start working:** Pod → **Resume** → wait ~1–2 min → team uses the same URLs.
|
||||
- **Done:** Pod → **Stop** (stops GPU billing; keeps the volume + URL).
|
||||
- **Never Terminate** — that deletes the pod and changes the URL.
|
||||
|
||||
## Tuning
|
||||
- Lighter/faster model: set env `QWEN_MODEL=qwen2.5:14b` (Q4, ~9 GB) or `qwen2.5:7b`.
|
||||
- The script reads `QWEN_MODEL`, `NLLB_HF`, `CHAT_POD_KEY` from env if you want to override.
|
||||
@@ -0,0 +1,77 @@
|
||||
#!/usr/bin/env bash
|
||||
# =============================================================================
|
||||
# CRM chat + translate pod — self-bootstrapping start script.
|
||||
#
|
||||
# Set this as the Pod's Docker/Start Command (one line):
|
||||
# bash -c "curl -fsSL https://raw.githubusercontent.com/Sumit-Pluto/photo_gallery/main/chat-pod/start.sh | bash"
|
||||
#
|
||||
# It is IDEMPOTENT: safe to run on first boot AND on every Resume. Everything
|
||||
# heavy (models + python deps) lives on the /workspace VOLUME, so Stop/Resume
|
||||
# keeps it — the CRM team just needs the two proxy URLs, nothing to edit.
|
||||
#
|
||||
# Qwen2.5-14B (chat/assistant) -> GPU, Ollama, port 11434 (OpenAI-compatible)
|
||||
# NLLB-200-1.3B (translation) -> CPU, CTranslate2 + FastAPI, port 8000
|
||||
# =============================================================================
|
||||
set -uo pipefail
|
||||
|
||||
WORK=/workspace
|
||||
VENV="$WORK/venv"
|
||||
export OLLAMA_MODELS="$WORK/ollama"
|
||||
export OLLAMA_HOST=0.0.0.0
|
||||
export NLLB_HF="facebook/nllb-200-1.3B"
|
||||
export NLLB_CT2="$WORK/nllb-ct2"
|
||||
QWEN_MODEL="${QWEN_MODEL:-qwen2.5:14b-instruct-q8_0}" # 8-bit ~16GB; override via env
|
||||
|
||||
mkdir -p "$WORK" "$OLLAMA_MODELS"
|
||||
echo "[chat-pod] boot — models/deps on $WORK (persist across Stop/Resume)"
|
||||
|
||||
# --- 0. system deps: lspci lets Ollama auto-detect the GPU (else it may use CPU) ---
|
||||
if ! command -v lspci >/dev/null 2>&1; then
|
||||
echo "[chat-pod] installing pciutils (GPU detection)..."
|
||||
apt-get update -qq >/dev/null 2>&1 && apt-get install -y -qq pciutils >/dev/null 2>&1 || true
|
||||
fi
|
||||
|
||||
# --- 1. Ollama (installs to container disk; reinstalled cheaply if missing) ----
|
||||
if ! command -v ollama >/dev/null 2>&1; then
|
||||
echo "[chat-pod] installing Ollama..."
|
||||
curl -fsSL https://ollama.com/install.sh | sh
|
||||
fi
|
||||
|
||||
# --- 2. Python deps on the VOLUME (persist) -----------------------------------
|
||||
if [ ! -x "$VENV/bin/uvicorn" ]; then
|
||||
echo "[chat-pod] creating venv + installing NLLB deps (one time)..."
|
||||
python3 -m venv "$VENV"
|
||||
"$VENV/bin/pip" install -q --upgrade pip
|
||||
"$VENV/bin/pip" install -q ctranslate2 transformers sentencepiece fastapi uvicorn
|
||||
fi
|
||||
# torch (CPU) is required to CONVERT the NLLB model; ensure it exists on any venv.
|
||||
"$VENV/bin/python" -c "import torch" >/dev/null 2>&1 || {
|
||||
echo "[chat-pod] installing CPU torch (needed to convert NLLB)..."
|
||||
"$VENV/bin/pip" install -q torch --index-url https://download.pytorch.org/whl/cpu
|
||||
}
|
||||
|
||||
# --- 3. Fetch the translate server (always latest from repo) -------------------
|
||||
curl -fsSL https://raw.githubusercontent.com/Sumit-Pluto/photo_gallery/main/chat-pod/translate.py -o "$WORK/translate.py" \
|
||||
|| echo "[chat-pod] WARN: could not refresh translate.py (using existing)"
|
||||
|
||||
# --- 4. Start Ollama, wait for it, ensure Qwen is pulled (idempotent) ---------
|
||||
echo "[chat-pod] starting Ollama (GPU)..."
|
||||
ollama serve &
|
||||
for i in $(seq 1 30); do curl -s http://localhost:11434/api/tags >/dev/null 2>&1 && break; sleep 2; done
|
||||
echo "[chat-pod] ensuring model $QWEN_MODEL (first time downloads ~16GB)..."
|
||||
ollama pull "$QWEN_MODEL" || echo "[chat-pod] WARN: pull failed (will use cached if present)"
|
||||
|
||||
# --- 5. Convert NLLB once (cached on the volume after) ------------------------
|
||||
if [ ! -f "$NLLB_CT2/model.bin" ]; then
|
||||
echo "[chat-pod] converting NLLB -> CTranslate2 int8 (one time)..."
|
||||
rm -rf "$NLLB_CT2" # clear any partial/failed conversion
|
||||
"$VENV/bin/ct2-transformers-converter" --model "$NLLB_HF" --output_dir "$NLLB_CT2" --quantization int8
|
||||
fi
|
||||
|
||||
# --- 6. Start the NLLB translate API (CPU) -----------------------------------
|
||||
echo "[chat-pod] starting NLLB translate API (CPU) on :8000..."
|
||||
cd "$WORK"
|
||||
"$VENV/bin/uvicorn" translate:app --host 0.0.0.0 --port 8000 &
|
||||
|
||||
echo "[chat-pod] READY -> Qwen :11434 (chat) | NLLB :8000 (translate)"
|
||||
wait
|
||||
@@ -0,0 +1,56 @@
|
||||
"""
|
||||
NLLB-200 translation API (CPU, CTranslate2 int8). Part of the CRM chat pod.
|
||||
Runs on port 8000 alongside Ollama/Qwen (port 11434) on the same pod.
|
||||
|
||||
POST /translate {"text": "...", "source": "eng_Latn", "target": "hin_Deva"}
|
||||
-> {"translation": "..."}
|
||||
GET /health -> {"ok": true}
|
||||
|
||||
Optional auth: set env CHAT_POD_KEY, then callers must send
|
||||
`Authorization: Bearer <key>`.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import ctranslate2
|
||||
import transformers
|
||||
from fastapi import FastAPI, Header, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
CT2 = os.environ.get("NLLB_CT2", "/workspace/nllb-ct2")
|
||||
HF = os.environ.get("NLLB_HF", "facebook/nllb-200-1.3B")
|
||||
API_KEY = os.environ.get("CHAT_POD_KEY", "")
|
||||
|
||||
_translator = ctranslate2.Translator(CT2, device="cpu", compute_type="int8")
|
||||
_tok = transformers.AutoTokenizer.from_pretrained(HF)
|
||||
|
||||
app = FastAPI(title="CRM Translate (NLLB-200)")
|
||||
|
||||
|
||||
class Req(BaseModel):
|
||||
text: str
|
||||
source: str # FLORES code, e.g. "eng_Latn"
|
||||
target: str # FLORES code, e.g. "hin_Deva"
|
||||
|
||||
|
||||
def _check(authorization: str) -> None:
|
||||
if API_KEY and authorization != f"Bearer {API_KEY}":
|
||||
raise HTTPException(status_code=401, detail="unauthorized")
|
||||
|
||||
|
||||
@app.post("/translate")
|
||||
def translate(r: Req, authorization: str = Header(default="")):
|
||||
_check(authorization)
|
||||
text = (r.text or "").strip()
|
||||
if not text:
|
||||
return {"translation": ""}
|
||||
_tok.src_lang = r.source
|
||||
source = _tok.convert_ids_to_tokens(_tok.encode(text))
|
||||
results = _translator.translate_batch([source], target_prefix=[[r.target]])
|
||||
hyp = results[0].hypotheses[0][1:] # drop the leading target-language token
|
||||
return {"translation": _tok.decode(_tok.convert_tokens_to_ids(hyp))}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {"ok": True}
|
||||
+59
-1
@@ -30,7 +30,65 @@ These edits *generate new pixels* from a prompt: **Restore, Colorize, Replace Sk
|
||||
Prompt**. Pick a backend with the `AI_EDIT_PROVIDER` env var (default `auto`). The server route
|
||||
(`/api/ai/edit`) proxies to it so no key ever reaches the browser.
|
||||
|
||||
`auto` chooses the first configured of: **local → huggingface → gemini**.
|
||||
`auto` chooses the first configured of: **runpod → local → huggingface → gemini**.
|
||||
|
||||
### Option 0 — RunPod serverless GPU endpoints (each model is its own endpoint) 🚀
|
||||
|
||||
This is the production backend for the models in `model-api-spec.docx`. Each model runs as a
|
||||
separate RunPod endpoint; the app proxies to them **server-side** (`/api/ai/edit`,
|
||||
`/api/ai/classify`) so `RUNPOD_API_KEY` and the URLs are **never** exposed to the browser.
|
||||
|
||||
```env
|
||||
AI_EDIT_PROVIDER=runpod
|
||||
RUNPOD_API_KEY=rpa_xxxxxxxx
|
||||
RUNPOD_SD_IMG2IMG_URL=https://api.runpod.ai/v2/<id>/runsync # #10
|
||||
RUNPOD_SD_INPAINT_URL=https://api.runpod.ai/v2/<id>/runsync # #9
|
||||
RUNPOD_UPSCALE_URL=https://api.runpod.ai/v2/<id>/runsync # #7
|
||||
RUNPOD_COLORIZE_URL=https://api.runpod.ai/v2/<id>/runsync # #8
|
||||
RUNPOD_BG_REMOVE_URL=https://api.runpod.ai/v2/<id>/runsync # #6 (optional; default is in-browser @imgly)
|
||||
RUNPOD_YOLO_URL=https://api.runpod.ai/v2/<id>/runsync # #1 (also set NEXT_PUBLIC_APG_RUNPOD_DETECT=true)
|
||||
```
|
||||
|
||||
**Which editor action hits which endpoint** (the app maps them for you):
|
||||
|
||||
| Editor action | `op.type` | RunPod endpoint | Notes |
|
||||
|---|---|---|---|
|
||||
| Restore & Enhance | `restore` | #7 Real-ESRGAN (`RUNPOD_UPSCALE_URL`) | ×4 + face-enhance |
|
||||
| Upscale | `upscale` | #7 Real-ESRGAN | ×2 / ×4 (`op.factor`) |
|
||||
| Colorize | `colorize` | #8 DDColor (`RUNPOD_COLORIZE_URL`) | |
|
||||
| Apply Prompt | `prompt` | #10 SD 3.5 img2img (`RUNPOD_SD_IMG2IMG_URL`) | whole-image, no mask |
|
||||
| Replace Sky | `replace-sky` | #9 SD 3.5 inpaint **if a mask is sent**, else #10 img2img (low strength) | see masking below |
|
||||
| Magic Eraser | `magic-eraser` | #9 SD 3.5 inpaint (`RUNPOD_SD_INPAINT_URL`) | **mask required** |
|
||||
| Generative Fill | `generative-fill` | #9 SD 3.5 inpaint | **mask + prompt required** |
|
||||
|
||||
**Endpoint I/O contract:** the app POSTs `{ input: { image, mask?, prompt?, strength?, guidance_scale?,
|
||||
num_inference_steps?, seed?, … } }` and reads the image back from `output` (any of `image`,
|
||||
`image_png`, `images[0]`). Deploy your handlers to accept that `input` shape (or adjust
|
||||
`apps/web/src/lib/runpod/endpoints.ts`). `/runsync` is preferred; `/run` + `/status/{id}` polling is
|
||||
supported as a fallback but must finish inside Vercel's 60 s function limit.
|
||||
|
||||
**Masking (SD 3.5 #9):** masked ops send a `maskBase64` PNG **the same pixel size as the image**
|
||||
(white = regenerate, black = keep). The client rasterizes the editor's `ImageData` mask to match the
|
||||
downscaled image automatically (`apps/web/src/lib/ai/imageEncode.ts` → `maskToBase64`). **Magic
|
||||
Eraser / Generative Fill therefore need a selection** — until a brush-mask UI is added to the editor
|
||||
they return a clear "needs a mask/selection" error; **Replace Sky** works today by degrading to a
|
||||
low-strength img2img (#10) when no mask is present (add sky-segmentation or a drawn mask to get true
|
||||
masked #9).
|
||||
|
||||
**Detection (#1):** set `NEXT_PUBLIC_APG_RUNPOD_DETECT=true` to route object detection to the YOLO
|
||||
construction-material classifier via `/api/ai/classify`; it falls back to in-browser COCO-SSD if the
|
||||
endpoint is unreachable. Verify the box format in `endpoints.ts` (`normalizeBox`) against your model —
|
||||
it assumes ultralytics `xyxy` pixel coordinates.
|
||||
|
||||
**Already local (no RunPod needed):** screenshot detection (#4), EXIF geolocation (#15), and the tag
|
||||
rename lookup (#5) run in the app/browser — see `lib/classify.ts`, `lib/media.ts`. Background removal
|
||||
(#6) also runs free in-browser by default.
|
||||
|
||||
**Not yet wired (typed client + env only):** camera tilt (#2), voice-to-text (#3), audio denoise
|
||||
(#12), and video frame-rate/resolution (#13/#14) have `rp*` functions in `endpoints.ts` and env vars,
|
||||
but **no route or UI** — they have no consuming feature yet. #13/#14 in particular need an out-of-band
|
||||
pipeline (presigned upload → async job → webhook), because per-frame video can't fit Vercel's 4.5 MB
|
||||
body / 60 s limits or the in-browser video export path. Add these when the corresponding feature is scheduled.
|
||||
|
||||
### Option 1 — Your own local Stable Diffusion (free + private, recommended) 🏆
|
||||
|
||||
|
||||
+19
-7
@@ -60,12 +60,21 @@ install the whole workspace). Do not commit `.env.local` (it's gitignored).
|
||||
**Server-only (optional — do NOT prefix with `NEXT_PUBLIC`):**
|
||||
| Name | Value |
|
||||
|---|---|
|
||||
| `GEMINI_API_KEY` | a Gemini key *(only for the AI generative-edit button)* |
|
||||
| `AI_EDIT_PROVIDER` | `runpod` *(or `local` / `huggingface` / `gemini` / `auto`)* |
|
||||
| `RUNPOD_API_KEY` | your RunPod API key *(shared by every RunPod endpoint)* |
|
||||
| `RUNPOD_SD_IMG2IMG_URL` | #10 prompt endpoint URL *(…/runsync)* |
|
||||
| `RUNPOD_SD_INPAINT_URL` | #9 masked sky/eraser/fill endpoint URL |
|
||||
| `RUNPOD_UPSCALE_URL` | #7 restore/upscale endpoint URL |
|
||||
| `RUNPOD_COLORIZE_URL` | #8 colorize endpoint URL |
|
||||
| `RUNPOD_YOLO_URL` | #1 detection endpoint URL *(+ `NEXT_PUBLIC_APG_RUNPOD_DETECT=true`)* |
|
||||
| `GEMINI_API_KEY` | a Gemini key *(alternative generative-edit backend)* |
|
||||
| `GEMINI_IMAGE_MODEL` | `gemini-2.5-flash-image` *(optional)* |
|
||||
|
||||
Without the Supabase vars the app still runs, but on browser-local storage only (data won't sync
|
||||
across devices). Without `GEMINI_API_KEY` the generative-edit route returns 503 while all the free
|
||||
in-browser AI keeps working.
|
||||
The RunPod URLs and `RUNPOD_API_KEY` are read **only** in the `/api/ai/*` route handlers (server
|
||||
side) — never send them to the browser (no `NEXT_PUBLIC_` prefix). Full endpoint map + the op →
|
||||
endpoint table is in [`docs/AI-SETUP.md`](./AI-SETUP.md). Without the Supabase vars the app still
|
||||
runs on browser-local storage; without any AI-edit provider the generative-edit route returns 503
|
||||
while all the free in-browser AI (detection, faces, OCR, background removal) keeps working.
|
||||
|
||||
Optional theming/feature flags (`NEXT_PUBLIC_APG_*`) are documented in [`docs/ENV.md`](./ENV.md).
|
||||
|
||||
@@ -90,9 +99,12 @@ they're intentionally not stored.
|
||||
|
||||
## Gotchas / limits (free tier)
|
||||
|
||||
- **Gemini AI edit + Vercel body limit:** Vercel Hobby caps request bodies at ~4.5 MB. The AI-edit
|
||||
route is capped accordingly and large images are downscaled client-side first; very large images
|
||||
may still be rejected. The free in-browser AI is unaffected.
|
||||
- **AI edit + Vercel body limit:** Vercel Hobby caps request bodies at ~4.5 MB. The AI-edit route is
|
||||
capped accordingly and large images are downscaled client-side first (≤1280 px). For masked SD 3.5
|
||||
ops the **image + mask share one body**, so they're budgeted together; a mask PNG is mostly flat
|
||||
black/white and compresses small, so this holds. Very large images may still be rejected. Prefer
|
||||
RunPod `/runsync` endpoints so a job returns within the 60 s function limit. The free in-browser AI
|
||||
is unaffected.
|
||||
- **Single-user demo:** the Supabase adapter does a full-state sync (no auth). For multi-user, add
|
||||
Supabase Auth and scope rows per user before going to production (the RLS policies are open for the
|
||||
demo — tighten them with `auth.uid()`).
|
||||
|
||||
@@ -43,10 +43,24 @@ export function createLocalStorageAdapter(key: string = STORAGE_KEY): StorageAda
|
||||
: [],
|
||||
}))
|
||||
: [];
|
||||
// Keep only string→string entries — never trust persisted data blindly.
|
||||
const labelAliases: Record<string, string> =
|
||||
parsed.labelAliases && typeof parsed.labelAliases === 'object'
|
||||
? Object.fromEntries(
|
||||
Object.entries(parsed.labelAliases).filter(
|
||||
([k, v]) => typeof k === 'string' && typeof v === 'string',
|
||||
),
|
||||
)
|
||||
: {};
|
||||
const deletedLabels: string[] = Array.isArray(parsed.deletedLabels)
|
||||
? (parsed.deletedLabels as unknown[]).filter((l): l is string => typeof l === 'string')
|
||||
: [];
|
||||
return {
|
||||
media,
|
||||
albums,
|
||||
people: Array.isArray(parsed.people) ? parsed.people : [],
|
||||
labelAliases,
|
||||
deletedLabels,
|
||||
version: typeof parsed.version === 'number' ? parsed.version : CURRENT_VERSION,
|
||||
};
|
||||
} catch {
|
||||
|
||||
@@ -4,6 +4,12 @@ export interface PersistedState {
|
||||
media: MediaItem[];
|
||||
albums: Album[];
|
||||
people: Person[];
|
||||
/**
|
||||
* User's permanent object-tag renames (canonical lowercased detector label →
|
||||
* chosen label). Optional for backward compatibility with pre-rename data.
|
||||
*/
|
||||
labelAliases?: Record<string, string>;
|
||||
deletedLabels?: string[];
|
||||
version: number;
|
||||
}
|
||||
|
||||
|
||||
@@ -33,18 +33,46 @@ export interface AIProvider {
|
||||
image: ImageBitmap | HTMLImageElement,
|
||||
op: GenerativeEditOp,
|
||||
): Promise<Blob>;
|
||||
|
||||
/**
|
||||
* Transcribe recorded speech (base64 WAV, 16 kHz mono PCM16) to text.
|
||||
* Powers voice input for photo annotations / comments.
|
||||
*/
|
||||
transcribeAudio?(audioBase64: string): Promise<string>;
|
||||
|
||||
/**
|
||||
* Denoise recorded audio (base64 WAV) → cleaned base64 WAV. Useful before
|
||||
* transcription when the recording was made on a noisy site.
|
||||
*/
|
||||
denoiseAudio?(audioBase64: string): Promise<string>;
|
||||
|
||||
/** Estimate camera tilt (roll/pitch/fov, degrees) for auto-straightening. */
|
||||
estimateTilt?(item: MediaItem, image: ImageBitmap | HTMLImageElement): Promise<TiltEstimate>;
|
||||
}
|
||||
|
||||
/** Camera-tilt estimate (degrees). `rollDegrees` = in-plane rotation to correct. */
|
||||
export interface TiltEstimate {
|
||||
rollDegrees: number;
|
||||
pitchDegrees: number;
|
||||
fovDegrees: number;
|
||||
}
|
||||
|
||||
export type GenerativeEditOp =
|
||||
| { type: 'remove-background' }
|
||||
| { type: 'replace-sky'; prompt?: string }
|
||||
| { type: 'magic-eraser'; mask: ImageData }
|
||||
| { type: 'generative-fill'; prompt: string; mask: ImageData }
|
||||
| { type: 'replace-sky'; prompt?: string; strength?: number }
|
||||
| { type: 'magic-eraser'; mask: ImageData; strength?: number }
|
||||
| { type: 'generative-fill'; prompt: string; mask: ImageData; strength?: number }
|
||||
| { type: 'restore' }
|
||||
| { type: 'upscale'; factor: 2 | 4 }
|
||||
| { type: 'colorize' }
|
||||
/** Free-form natural-language edit instruction. */
|
||||
| { type: 'prompt'; prompt: string };
|
||||
/** Outpaint / expand-canvas: pad the image and generatively fill the new border. */
|
||||
| { type: 'outpaint'; prompt?: string; factor?: number; strength?: number }
|
||||
/**
|
||||
* Free-form natural-language edit instruction. `strength` (0..1) controls how
|
||||
* strongly the edit is applied (subtle → strong); the backend maps it to the
|
||||
* appropriate model knob.
|
||||
*/
|
||||
| { type: 'prompt'; prompt: string; strength?: number };
|
||||
|
||||
/** Cosine similarity between two equal-length vectors. */
|
||||
export function cosineSimilarity(a: number[], b: number[]): number {
|
||||
|
||||
@@ -4,12 +4,13 @@ import { useEffect, useRef } from 'react';
|
||||
|
||||
import type { AIProvider } from '../ai/types';
|
||||
import { PET_LABELS } from '../constants';
|
||||
import { resolveLabel } from '../lib/smartAlbums';
|
||||
import { sanitizeOcrText } from '../lib/text';
|
||||
import { useGallery, useGalleryStoreApi } from '../store/context';
|
||||
import type { MediaItem } from '../types';
|
||||
|
||||
const CONCURRENCY = 2;
|
||||
const CONFIDENCE = 0.5;
|
||||
const CONFIDENCE = 0.15;
|
||||
const START_DELAY_MS = 1200;
|
||||
|
||||
/**
|
||||
@@ -33,14 +34,16 @@ export function AIAnalyzer({ provider }: { provider: AIProvider }) {
|
||||
provider.detectObjects || provider.detectFaces || provider.ocr || provider.embedImage;
|
||||
if (!ready || !canDetect || runningRef.current) return;
|
||||
|
||||
// Each capability has its OWN backfill gate (not the shared analyzedAt), so a
|
||||
// capability added in a later session re-processes items analyzed before it
|
||||
// existed: objects key off analyzedAt; faces off `faces === undefined`; OCR
|
||||
// off `ocrText === undefined`.
|
||||
// Analyze each image EXACTLY ONCE: every capability is gated on `!analyzedAt`,
|
||||
// so once an item has been analyzed (and analyzedAt persisted) it is never
|
||||
// re-processed on later reloads — even if an individual result field didn't
|
||||
// round-trip through storage. Fresh uploads (no analyzedAt) run everything.
|
||||
const needsObjects = (m: MediaItem) => !!provider.detectObjects && !m.analyzedAt;
|
||||
const needsFaces = (m: MediaItem) => !!provider.detectFaces && m.faces === undefined;
|
||||
const needsOcr = (m: MediaItem) => !!provider.ocr && m.ocrText === undefined;
|
||||
const needsEmbedding = (m: MediaItem) => !!provider.embedImage && m.embedding === undefined;
|
||||
const needsFaces = (m: MediaItem) =>
|
||||
!!provider.detectFaces && !m.analyzedAt && m.faces === undefined;
|
||||
const needsOcr = (m: MediaItem) => !!provider.ocr && !m.analyzedAt && m.ocrText === undefined;
|
||||
const needsEmbedding = (m: MediaItem) =>
|
||||
!!provider.embedImage && !m.analyzedAt && m.embedding === undefined;
|
||||
// Collect items still needing analysis, NEWEST FIRST — so a freshly uploaded or
|
||||
// captured photo is tagged immediately instead of waiting behind the whole library.
|
||||
const collectPending = () =>
|
||||
@@ -99,8 +102,17 @@ export function AIAnalyzer({ provider }: { provider: AIProvider }) {
|
||||
const patch: Partial<MediaItem> = { analyzedAt: Date.now() };
|
||||
if (objects) {
|
||||
patch.objects = objects;
|
||||
// Map each detected label through the user's rename map so future uploads
|
||||
// are stored/grouped under the renamed label instead of a fresh class.
|
||||
const aliases = api.getState().labelAliases;
|
||||
const deleted = api.getState().deletedLabels;
|
||||
patch.objectLabels = [
|
||||
...new Set(objects.filter((o) => o.confidence >= CONFIDENCE).map((o) => o.label)),
|
||||
...new Set(
|
||||
objects
|
||||
.filter((o) => o.confidence >= CONFIDENCE)
|
||||
.map((o) => resolveLabel(o.label, aliases))
|
||||
.filter((l) => !deleted.includes(l)),
|
||||
),
|
||||
];
|
||||
}
|
||||
if (faces) patch.faces = faces;
|
||||
|
||||
@@ -5,7 +5,9 @@ import { useEffect, useRef, useState } from 'react';
|
||||
import { sourceLabel } from '../lib/classify';
|
||||
import { editFilterCss, editTransformCss } from '../lib/edits';
|
||||
import { formatBytes, formatDate, formatDuration, formatTime } from '../lib/format';
|
||||
import { blobToWavBase64, startRecording, wavBase64ToBlob, type Recorder } from '../lib/audioCapture';
|
||||
import { Icon } from '../icons';
|
||||
import { useAIProvider } from './aiContext';
|
||||
import { useGallery, useGalleryStoreApi } from '../store/context';
|
||||
import type { MediaItem } from '../types';
|
||||
|
||||
@@ -49,6 +51,13 @@ export function InfoPanel() {
|
||||
{formatDate(item.takenAt)} · {formatTime(item.takenAt)}
|
||||
</div>
|
||||
|
||||
{item.kind === 'image' && !item.analyzedAt ? (
|
||||
<div className="apg-info__analyzing">
|
||||
<span className="apg-info__spinner" aria-hidden />
|
||||
Analyzing image…
|
||||
</div>
|
||||
) : null}
|
||||
|
||||
<Row label="Kind" value={item.kind === 'video' ? 'Video' : 'Photo'} />
|
||||
<Row label="Source" value={sourceLabel(item.source)} />
|
||||
<Row label="Format" value={`${item.mime} (${ext})`} />
|
||||
@@ -70,6 +79,10 @@ export function InfoPanel() {
|
||||
.join(' · ')}
|
||||
/>
|
||||
) : null}
|
||||
{item.exif?.LensModel ? <Row label="Lens" value={String(item.exif.LensModel)} /> : null}
|
||||
{item.exif?.Orientation ? (
|
||||
<Row label="Orientation" value={orientationLabel(item.exif.Orientation)} />
|
||||
) : null}
|
||||
|
||||
{item.objectLabels.length ? (
|
||||
<Chips
|
||||
@@ -113,8 +126,13 @@ export function InfoPanel() {
|
||||
lng={item.location.lng}
|
||||
onOpen={() => api.getState().focusMap({ lat: item.location!.lat, lng: item.location!.lng })}
|
||||
/>
|
||||
<AddressLookup
|
||||
lat={item.location.lat}
|
||||
lng={item.location.lng}
|
||||
onOpenMap={() => api.getState().focusMap({ lat: item.location!.lat, lng: item.location!.lng })}
|
||||
/>
|
||||
<div style={{ fontSize: 11, color: 'var(--apg-text-tertiary)', marginTop: 4 }}>
|
||||
Click the map to open it in full.
|
||||
Click the map or address to open it in full.
|
||||
</div>
|
||||
</>
|
||||
) : null}
|
||||
@@ -221,6 +239,49 @@ function Comments({ item }: { item: MediaItem }) {
|
||||
return window.localStorage.getItem('apg:comment-author') || 'You';
|
||||
});
|
||||
|
||||
const provider = useAIProvider();
|
||||
const canVoice = Boolean(provider?.transcribeAudio);
|
||||
const canDenoise = Boolean(provider?.denoiseAudio);
|
||||
const [recording, setRecording] = useState(false);
|
||||
const [denoise, setDenoise] = useState(false);
|
||||
const [voiceStatus, setVoiceStatus] = useState<string | null>(null);
|
||||
const recorderRef = useRef<Recorder | null>(null);
|
||||
|
||||
const startVoice = async () => {
|
||||
setVoiceStatus(null);
|
||||
try {
|
||||
recorderRef.current = await startRecording();
|
||||
setRecording(true);
|
||||
} catch (e) {
|
||||
setVoiceStatus(e instanceof Error ? e.message : 'Microphone unavailable.');
|
||||
}
|
||||
};
|
||||
|
||||
const stopVoice = async () => {
|
||||
const rec = recorderRef.current;
|
||||
recorderRef.current = null;
|
||||
setRecording(false);
|
||||
if (!rec || !provider?.transcribeAudio) return;
|
||||
try {
|
||||
const blob = await rec.stop();
|
||||
let wav16: string;
|
||||
if (denoise && provider.denoiseAudio) {
|
||||
setVoiceStatus('Reducing noise…');
|
||||
const wav48 = await blobToWavBase64(blob, 48000);
|
||||
const cleaned = await provider.denoiseAudio(wav48);
|
||||
wav16 = await blobToWavBase64(wavBase64ToBlob(cleaned), 16000);
|
||||
} else {
|
||||
wav16 = await blobToWavBase64(blob, 16000);
|
||||
}
|
||||
setVoiceStatus('Transcribing…');
|
||||
const spoken = (await provider.transcribeAudio(wav16)).trim();
|
||||
if (spoken) setText((prev) => (prev ? `${prev} ${spoken}` : spoken));
|
||||
setVoiceStatus(null);
|
||||
} catch (e) {
|
||||
setVoiceStatus(e instanceof Error ? e.message : 'Could not transcribe audio.');
|
||||
}
|
||||
};
|
||||
|
||||
const post = () => {
|
||||
const t = text.trim();
|
||||
if (!t) return;
|
||||
@@ -294,6 +355,36 @@ function Comments({ item }: { item: MediaItem }) {
|
||||
rows={2}
|
||||
maxLength={2000}
|
||||
/>
|
||||
{canVoice ? (
|
||||
<div className="apg-voice">
|
||||
<button
|
||||
type="button"
|
||||
className={`apg-btn apg-btn--small apg-voice__mic${recording ? ' apg-voice__mic--rec' : ''}`}
|
||||
onClick={recording ? stopVoice : startVoice}
|
||||
aria-label={recording ? 'Stop recording' : 'Record a voice comment'}
|
||||
title={recording ? 'Stop & transcribe' : 'Speak your comment'}
|
||||
>
|
||||
<Icon name={recording ? 'check' : 'mic'} size={14} />
|
||||
{recording ? 'Stop' : 'Speak'}
|
||||
</button>
|
||||
{canDenoise ? (
|
||||
<label
|
||||
className="apg-voice__denoise"
|
||||
title="Clean up background noise before transcribing"
|
||||
>
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={denoise}
|
||||
onChange={(e) => setDenoise(e.target.checked)}
|
||||
/>
|
||||
Reduce noise
|
||||
</label>
|
||||
) : null}
|
||||
<span className="apg-voice__status" aria-live="polite">
|
||||
{recording ? '● Listening…' : (voiceStatus ?? '')}
|
||||
</span>
|
||||
</div>
|
||||
) : null}
|
||||
<button
|
||||
type="button"
|
||||
className="apg-btn apg-btn--small apg-comment-form__post"
|
||||
@@ -316,6 +407,22 @@ function Row({ label, value }: { label: string; value: string }) {
|
||||
);
|
||||
}
|
||||
|
||||
// EXIF Orientation codes 1–8 → human text (construction shots are often sideways).
|
||||
const ORIENTATION_LABELS: Record<number, string> = {
|
||||
1: 'Normal',
|
||||
2: 'Mirrored horizontal',
|
||||
3: 'Rotated 180°',
|
||||
4: 'Mirrored vertical',
|
||||
5: 'Mirrored + 90° CCW',
|
||||
6: 'Rotated 90° CW',
|
||||
7: 'Mirrored + 90° CW',
|
||||
8: 'Rotated 90° CCW',
|
||||
};
|
||||
function orientationLabel(v: string | number): string {
|
||||
const n = typeof v === 'number' ? v : parseInt(String(v), 10);
|
||||
return ORIENTATION_LABELS[n] ?? String(v);
|
||||
}
|
||||
|
||||
function Chips({
|
||||
label,
|
||||
items,
|
||||
@@ -345,6 +452,72 @@ function Chips({
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* "Show address" button → reverse-geocodes the GPS coords to a human-readable
|
||||
* address (free OpenStreetMap Nominatim). Clicking the resolved address opens the
|
||||
* full Map at that location.
|
||||
*/
|
||||
function AddressLookup({ lat, lng, onOpenMap }: { lat: number; lng: number; onOpenMap: () => void }) {
|
||||
const [address, setAddress] = useState<string | null>(null);
|
||||
const [state, setState] = useState<'idle' | 'loading' | 'error'>('idle');
|
||||
|
||||
const lookup = async () => {
|
||||
setState('loading');
|
||||
try {
|
||||
const res = await fetch(
|
||||
`https://nominatim.openstreetmap.org/reverse?lat=${lat}&lon=${lng}&format=json&zoom=18&addressdetails=1`,
|
||||
{ headers: { Accept: 'application/json' } },
|
||||
);
|
||||
const data = (await res.json()) as { display_name?: string };
|
||||
if (data?.display_name) {
|
||||
setAddress(data.display_name);
|
||||
setState('idle');
|
||||
} else {
|
||||
setState('error');
|
||||
}
|
||||
} catch {
|
||||
setState('error');
|
||||
}
|
||||
};
|
||||
|
||||
if (address) {
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={onOpenMap}
|
||||
title="Open in Map"
|
||||
style={{
|
||||
display: 'block',
|
||||
textAlign: 'left',
|
||||
width: '100%',
|
||||
marginTop: 6,
|
||||
padding: '6px 8px',
|
||||
background: 'var(--apg-bg-elevated)',
|
||||
border: '1px solid var(--apg-glass-border, rgba(255,255,255,0.1))',
|
||||
borderRadius: 8,
|
||||
color: 'var(--apg-text)',
|
||||
fontSize: 12,
|
||||
cursor: 'pointer',
|
||||
}}
|
||||
>
|
||||
📍 {address}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
className="apg-btn apg-btn--small"
|
||||
onClick={() => void lookup()}
|
||||
disabled={state === 'loading'}
|
||||
style={{ marginTop: 6 }}
|
||||
>
|
||||
{state === 'loading' ? 'Looking up…' : state === 'error' ? 'Retry address' : 'Show address'}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
/** Small non-interactive Leaflet map for the location preview (click to open full Map). */
|
||||
function MiniMap({ lat, lng, onOpen }: { lat: number; lng: number; onOpen?: () => void }) {
|
||||
const ref = useRef<HTMLDivElement>(null);
|
||||
|
||||
@@ -7,7 +7,7 @@ import { Icon, type IconName } from '../icons';
|
||||
import { useGallery, useGalleryStoreApi } from '../store/context';
|
||||
import type { Album, ViewId } from '../types';
|
||||
import { closeContextMenu, openContextMenu } from './ContextMenu';
|
||||
import { openSecuritySettings, promptAlbumName } from './modals';
|
||||
import { confirmAction, openSecuritySettings, promptAlbumName } from './modals';
|
||||
|
||||
interface RowProps {
|
||||
icon: IconName;
|
||||
@@ -161,6 +161,38 @@ export function Sidebar() {
|
||||
]);
|
||||
};
|
||||
|
||||
// Right-click an auto object album → permanently rename its tag (e.g. car → excavator).
|
||||
const objectMenu = (album: Album) => (e: React.MouseEvent) => {
|
||||
e.preventDefault();
|
||||
const label = album.id.slice('sys:obj:'.length);
|
||||
openContextMenu(e.clientX, e.clientY, [
|
||||
{
|
||||
label: 'Rename Tag',
|
||||
icon: 'tag',
|
||||
onClick: () =>
|
||||
promptAlbumName(
|
||||
'Rename Tag',
|
||||
album.name,
|
||||
(name) => api.getState().renameLabel(label, name),
|
||||
{ placeholder: 'Tag name' },
|
||||
),
|
||||
},
|
||||
{
|
||||
label: 'Delete Tag',
|
||||
icon: 'trash',
|
||||
danger: true,
|
||||
onClick: () =>
|
||||
confirmAction({
|
||||
title: 'Delete Tag',
|
||||
message: `Remove the "${album.name}" tag? It's deleted from all photos and won't be created again.`,
|
||||
confirmLabel: 'Delete',
|
||||
danger: true,
|
||||
onConfirm: () => api.getState().deleteLabel(label),
|
||||
}),
|
||||
},
|
||||
]);
|
||||
};
|
||||
|
||||
const newAlbum = () =>
|
||||
promptAlbumName('New Album', '', (name) => {
|
||||
const id = api.getState().createAlbum(name);
|
||||
@@ -182,7 +214,7 @@ export function Sidebar() {
|
||||
<Row icon="video" label="Videos" view="videos" />
|
||||
<Row icon="screenshot" label="Screenshots" view="screenshots" />
|
||||
<Row icon="document" label="Documents" view="sys:documents" />
|
||||
<Row icon="person-circle" label="People & Pets" view="people" />
|
||||
<Row icon="person-circle" label="People" view="people" />
|
||||
<Row
|
||||
icon="trash"
|
||||
label="Recently Deleted"
|
||||
@@ -251,6 +283,7 @@ export function Sidebar() {
|
||||
label={a.name}
|
||||
view={a.id as ViewId}
|
||||
indent
|
||||
onContextMenu={objectMenu(a)}
|
||||
/>
|
||||
))
|
||||
: null}
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
'use client';
|
||||
|
||||
import { type PointerEvent as ReactPointerEvent, useEffect, useRef, useState } from 'react';
|
||||
|
||||
import { Icon } from '../../icons';
|
||||
|
||||
/**
|
||||
* Full-screen brush overlay for masked AI edits (Magic Eraser / Generative Fill).
|
||||
* The user paints over a region; on Apply we emit a binary ImageData mask where
|
||||
* painted pixels are WHITE (regenerate) and everything else is BLACK (keep) —
|
||||
* exactly what the inpaint backend expects (see maskToBase64 / rpInpaint).
|
||||
*
|
||||
* The paint canvas is sized to the image's aspect ratio (long side = BASE), then
|
||||
* scaled with CSS to fit the viewport, so the returned mask lines up with the
|
||||
* photo regardless of screen size.
|
||||
*/
|
||||
const BASE = 640;
|
||||
|
||||
export function MaskBrush({
|
||||
src,
|
||||
aspect,
|
||||
title,
|
||||
onCancel,
|
||||
onApply,
|
||||
}: {
|
||||
src: string;
|
||||
aspect: number; // width / height
|
||||
title: string;
|
||||
onCancel: () => void;
|
||||
onApply: (mask: ImageData) => void;
|
||||
}) {
|
||||
const canvasRef = useRef<HTMLCanvasElement>(null);
|
||||
const paintingRef = useRef(false);
|
||||
const lastRef = useRef<{ x: number; y: number } | null>(null);
|
||||
const [brush, setBrush] = useState(48);
|
||||
const [dirty, setDirty] = useState(false);
|
||||
|
||||
const cw = aspect >= 1 ? BASE : Math.round(BASE * aspect);
|
||||
const ch = aspect >= 1 ? Math.round(BASE / aspect) : BASE;
|
||||
|
||||
// Escape cancels.
|
||||
useEffect(() => {
|
||||
const onKey = (e: KeyboardEvent) => {
|
||||
if (e.key === 'Escape') onCancel();
|
||||
};
|
||||
window.addEventListener('keydown', onKey);
|
||||
return () => window.removeEventListener('keydown', onKey);
|
||||
}, [onCancel]);
|
||||
|
||||
const ctx = () => canvasRef.current?.getContext('2d') ?? null;
|
||||
|
||||
const toCanvas = (e: ReactPointerEvent) => {
|
||||
const c = canvasRef.current;
|
||||
if (!c) return { x: 0, y: 0 };
|
||||
const r = c.getBoundingClientRect();
|
||||
return {
|
||||
x: ((e.clientX - r.left) / r.width) * c.width,
|
||||
y: ((e.clientY - r.top) / r.height) * c.height,
|
||||
};
|
||||
};
|
||||
|
||||
const paintTo = (x: number, y: number) => {
|
||||
const c = ctx();
|
||||
if (!c) return;
|
||||
c.fillStyle = 'rgba(255,60,60,0.55)';
|
||||
c.strokeStyle = 'rgba(255,60,60,0.55)';
|
||||
c.lineWidth = brush;
|
||||
c.lineCap = 'round';
|
||||
const last = lastRef.current;
|
||||
if (last) {
|
||||
c.beginPath();
|
||||
c.moveTo(last.x, last.y);
|
||||
c.lineTo(x, y);
|
||||
c.stroke();
|
||||
}
|
||||
c.beginPath();
|
||||
c.arc(x, y, brush / 2, 0, Math.PI * 2);
|
||||
c.fill();
|
||||
lastRef.current = { x, y };
|
||||
};
|
||||
|
||||
const onDown = (e: ReactPointerEvent) => {
|
||||
e.preventDefault();
|
||||
(e.target as HTMLElement).setPointerCapture?.(e.pointerId);
|
||||
paintingRef.current = true;
|
||||
lastRef.current = null;
|
||||
const p = toCanvas(e);
|
||||
paintTo(p.x, p.y);
|
||||
setDirty(true);
|
||||
};
|
||||
const onMove = (e: ReactPointerEvent) => {
|
||||
if (!paintingRef.current) return;
|
||||
const p = toCanvas(e);
|
||||
paintTo(p.x, p.y);
|
||||
};
|
||||
const onUp = () => {
|
||||
paintingRef.current = false;
|
||||
lastRef.current = null;
|
||||
};
|
||||
|
||||
const clear = () => {
|
||||
const c = ctx();
|
||||
if (c && canvasRef.current) c.clearRect(0, 0, canvasRef.current.width, canvasRef.current.height);
|
||||
setDirty(false);
|
||||
};
|
||||
|
||||
const apply = () => {
|
||||
const c = ctx();
|
||||
const canvas = canvasRef.current;
|
||||
if (!c || !canvas) return;
|
||||
const painted = c.getImageData(0, 0, canvas.width, canvas.height);
|
||||
// Binarize: any painted (alpha) pixel → opaque white, else opaque black.
|
||||
const out = new ImageData(canvas.width, canvas.height);
|
||||
for (let i = 0; i < painted.data.length; i += 4) {
|
||||
const on = painted.data[i + 3]! > 10;
|
||||
const v = on ? 255 : 0;
|
||||
out.data[i] = v;
|
||||
out.data[i + 1] = v;
|
||||
out.data[i + 2] = v;
|
||||
out.data[i + 3] = 255;
|
||||
}
|
||||
onApply(out);
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="apg-maskbrush" role="dialog" aria-label={title}>
|
||||
<div className="apg-maskbrush__title">{title}</div>
|
||||
<div className="apg-maskbrush__stage" style={{ aspectRatio: `${cw} / ${ch}` }}>
|
||||
<img className="apg-maskbrush__img" src={src} alt="" draggable={false} />
|
||||
<canvas
|
||||
ref={canvasRef}
|
||||
width={cw}
|
||||
height={ch}
|
||||
className="apg-maskbrush__canvas"
|
||||
onPointerDown={onDown}
|
||||
onPointerMove={onMove}
|
||||
onPointerUp={onUp}
|
||||
onPointerLeave={onUp}
|
||||
/>
|
||||
</div>
|
||||
<div className="apg-maskbrush__bar">
|
||||
<label className="apg-maskbrush__brush">
|
||||
Brush
|
||||
<input
|
||||
type="range"
|
||||
min={12}
|
||||
max={120}
|
||||
step={2}
|
||||
value={brush}
|
||||
onChange={(e) => setBrush(Number(e.target.value))}
|
||||
/>
|
||||
</label>
|
||||
<button type="button" className="apg-btn apg-btn--small" onClick={clear} disabled={!dirty}>
|
||||
<Icon name="trash" size={14} /> Clear
|
||||
</button>
|
||||
<div style={{ flex: 1 }} />
|
||||
<button type="button" className="apg-btn apg-btn--small" onClick={onCancel}>
|
||||
Cancel
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
className="apg-btn apg-btn--primary apg-btn--small"
|
||||
onClick={apply}
|
||||
disabled={!dirty}
|
||||
>
|
||||
Apply
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -17,6 +17,8 @@ import { addToAlbumPicker, confirmAction } from '../modals';
|
||||
import { useAIProvider } from '../aiContext';
|
||||
import { Annotations, type AnnotationTool } from './Annotations';
|
||||
import { CropBox, type CropRect } from './CropBox';
|
||||
import { MaskBrush } from './MaskBrush';
|
||||
import { VoiceButton } from './VoiceButton';
|
||||
|
||||
const RATIOS: Array<{ label: string; value: number | null }> = [
|
||||
{ label: 'Free', value: null },
|
||||
@@ -77,6 +79,14 @@ const AI_OPS: Array<{ label: string; op: GenerativeEditOp; icon: 'wand' | 'image
|
||||
{ label: 'Replace Sky', op: { type: 'replace-sky' }, icon: 'image' },
|
||||
];
|
||||
|
||||
/** Ops whose result varies with the "edit strength" slider (the backend maps it per model). */
|
||||
const STRENGTH_OPS = new Set<GenerativeEditOp['type']>([
|
||||
'prompt',
|
||||
'replace-sky',
|
||||
'magic-eraser',
|
||||
'generative-fill',
|
||||
]);
|
||||
|
||||
export function PhotoEditor() {
|
||||
const api = useGalleryStoreApi();
|
||||
const editorId = useGallery((s) => s.editorId);
|
||||
@@ -91,6 +101,10 @@ export function PhotoEditor() {
|
||||
const [aiError, setAiError] = useState<string | null>(null);
|
||||
const [aiResultUrl, setAiResultUrl] = useState<string | null>(null);
|
||||
const [aiPrompt, setAiPrompt] = useState('');
|
||||
const [aiStrength, setAiStrength] = useState(0.5);
|
||||
const [maskMode, setMaskMode] = useState<'magic-eraser' | 'generative-fill' | null>(null);
|
||||
const [tiltBusy, setTiltBusy] = useState(false);
|
||||
const [tiltError, setTiltError] = useState<string | null>(null);
|
||||
const [annTool, setAnnTool] = useState<AnnotationTool>('rect');
|
||||
const [annColor, setAnnColor] = useState<string>('#ff3b30');
|
||||
const [cropRatio, setCropRatio] = useState<number | null>(null);
|
||||
@@ -146,7 +160,11 @@ export function PhotoEditor() {
|
||||
setAiError(null);
|
||||
try {
|
||||
const img = await loadCrossOriginImage(item.src);
|
||||
const blob = await provider.generativeEdit(item, img, op);
|
||||
// Attach the current "edit strength" to ops that support it.
|
||||
const opToRun = STRENGTH_OPS.has(op.type)
|
||||
? ({ ...op, strength: aiStrength } as GenerativeEditOp)
|
||||
: op;
|
||||
const blob = await provider.generativeEdit(item, img, opToRun);
|
||||
aiBlobRef.current = blob;
|
||||
setAiResultUrl((prev) => {
|
||||
if (prev) URL.revokeObjectURL(prev);
|
||||
@@ -159,6 +177,22 @@ export function PhotoEditor() {
|
||||
}
|
||||
};
|
||||
|
||||
const autoStraighten = async () => {
|
||||
if (!provider?.estimateTilt) return;
|
||||
setTiltBusy(true);
|
||||
setTiltError(null);
|
||||
try {
|
||||
const img = await loadCrossOriginImage(item.src);
|
||||
const t = await provider.estimateTilt(item, img);
|
||||
const roll = Math.max(-45, Math.min(45, Math.round(t.rollDegrees)));
|
||||
setEdits((e) => ({ ...e, straighten: roll }));
|
||||
} catch (err) {
|
||||
setTiltError(err instanceof Error ? err.message : 'Could not estimate tilt.');
|
||||
} finally {
|
||||
setTiltBusy(false);
|
||||
}
|
||||
};
|
||||
|
||||
const clearAI = () => {
|
||||
aiBlobRef.current = null;
|
||||
setAiError(null);
|
||||
@@ -379,7 +413,7 @@ export function PhotoEditor() {
|
||||
}}
|
||||
>
|
||||
<span className="apg-ai-spinner" style={{ width: 26, height: 26 }} />
|
||||
<span style={{ fontSize: 13 }}>Generating with Gemini…</span>
|
||||
<span style={{ fontSize: 13 }}>Generating…</span>
|
||||
</div>
|
||||
) : null}
|
||||
</div>
|
||||
@@ -526,6 +560,20 @@ export function PhotoEditor() {
|
||||
Reset Straighten
|
||||
</button>
|
||||
) : null}
|
||||
{provider?.estimateTilt ? (
|
||||
<button
|
||||
type="button"
|
||||
className="apg-editor__tab"
|
||||
disabled={tiltBusy}
|
||||
onClick={() => void autoStraighten()}
|
||||
>
|
||||
<Icon name="wand" size={15} />{' '}
|
||||
{tiltBusy ? 'Analyzing tilt…' : 'Auto-straighten (fix camera tilt)'}
|
||||
</button>
|
||||
) : null}
|
||||
{tiltError ? (
|
||||
<p style={{ color: '#ff6b6b', fontSize: 12, margin: 0 }}>{tiltError}</p>
|
||||
) : null}
|
||||
</div>
|
||||
) : null}
|
||||
|
||||
@@ -547,6 +595,28 @@ export function PhotoEditor() {
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
{provider?.transcribeAudio ? (
|
||||
<VoiceButton
|
||||
label="Speak → add text"
|
||||
onText={(t) => {
|
||||
setAnnotations([
|
||||
...annotations,
|
||||
{
|
||||
id: nanoid(8),
|
||||
shape: 'text',
|
||||
color: annColor,
|
||||
strokeWidth: 2,
|
||||
x1: 0.08,
|
||||
y1: 0.08,
|
||||
x2: 0.55,
|
||||
y2: 0.17,
|
||||
text: t,
|
||||
},
|
||||
]);
|
||||
setAnnTool('select');
|
||||
}}
|
||||
/>
|
||||
) : null}
|
||||
<div>
|
||||
<div style={{ fontSize: 12, color: '#9b9ba1', marginBottom: 6 }}>Color</div>
|
||||
<div style={{ display: 'flex', gap: 8, flexWrap: 'wrap' }}>
|
||||
@@ -596,7 +666,7 @@ export function PhotoEditor() {
|
||||
{tab === 'ai' ? (
|
||||
<div style={{ display: 'flex', flexDirection: 'column', gap: 10 }}>
|
||||
<p style={{ color: '#9b9ba1', fontSize: 12, margin: '0 0 4px' }}>
|
||||
Generative edits powered by Gemini. Results replace the photo when you press Save.
|
||||
Generative edits run through your configured AI backend. Results replace the photo when you press Save.
|
||||
</p>
|
||||
{AI_OPS.map((a) => (
|
||||
<button
|
||||
@@ -610,6 +680,35 @@ export function PhotoEditor() {
|
||||
</button>
|
||||
))}
|
||||
|
||||
<div className="apg-slider-row" style={{ marginTop: 6 }}>
|
||||
<div className="apg-slider-row__head">
|
||||
<span>Edit strength</span>
|
||||
<span>{Math.round(aiStrength * 100)}%</span>
|
||||
</div>
|
||||
<input
|
||||
className="apg-slider"
|
||||
type="range"
|
||||
min={0}
|
||||
max={1}
|
||||
step={0.05}
|
||||
value={aiStrength}
|
||||
disabled={aiBusy}
|
||||
onChange={(e) => setAiStrength(Number(e.target.value))}
|
||||
/>
|
||||
<div
|
||||
style={{
|
||||
display: 'flex',
|
||||
justifyContent: 'space-between',
|
||||
fontSize: 11,
|
||||
color: '#9b9ba1',
|
||||
marginTop: 2,
|
||||
}}
|
||||
>
|
||||
<span>Subtle</span>
|
||||
<span>Strong</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div style={{ marginTop: 6 }}>
|
||||
<input
|
||||
className="apg-modal__input"
|
||||
@@ -633,6 +732,36 @@ export function PhotoEditor() {
|
||||
</button>
|
||||
</div>
|
||||
|
||||
<div style={{ height: 1, background: 'rgba(255,255,255,0.1)', margin: '4px 0' }} />
|
||||
<p style={{ color: '#9b9ba1', fontSize: 12, margin: 0 }}>
|
||||
Brush & expand tools — inpaint / outpaint.
|
||||
</p>
|
||||
<button
|
||||
type="button"
|
||||
className="apg-editor__tab"
|
||||
disabled={aiBusy}
|
||||
onClick={() => setMaskMode('magic-eraser')}
|
||||
>
|
||||
<Icon name="wand" size={16} /> Magic Eraser (remove an object)
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
className="apg-editor__tab"
|
||||
disabled={aiBusy}
|
||||
title="Paint an area, then it fills it — type a prompt above to control what appears (optional)"
|
||||
onClick={() => setMaskMode('generative-fill')}
|
||||
>
|
||||
<Icon name="image" size={16} /> Generative Fill (paint + prompt)
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
className="apg-editor__tab"
|
||||
disabled={aiBusy}
|
||||
onClick={() => void runAI({ type: 'outpaint', prompt: aiPrompt.trim() || undefined })}
|
||||
>
|
||||
<Icon name="crop" size={16} /> Expand Image (Outpaint)
|
||||
</button>
|
||||
|
||||
{aiResultUrl ? (
|
||||
<button type="button" className="apg-editor__tab" onClick={clearAI} disabled={aiBusy}>
|
||||
Discard AI result
|
||||
@@ -645,6 +774,28 @@ export function PhotoEditor() {
|
||||
) : null}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{maskMode ? (
|
||||
<MaskBrush
|
||||
src={item.src}
|
||||
aspect={imgAspect}
|
||||
title={
|
||||
maskMode === 'magic-eraser'
|
||||
? 'Paint over what to remove'
|
||||
: 'Paint the area to replace (uses your prompt)'
|
||||
}
|
||||
onCancel={() => setMaskMode(null)}
|
||||
onApply={(mask) => {
|
||||
const m = maskMode;
|
||||
setMaskMode(null);
|
||||
if (m === 'generative-fill') {
|
||||
void runAI({ type: 'generative-fill', prompt: aiPrompt.trim() || 'fill naturally', mask });
|
||||
} else {
|
||||
void runAI({ type: 'magic-eraser', mask });
|
||||
}
|
||||
}}
|
||||
/>
|
||||
) : null}
|
||||
</motion.div>
|
||||
</AnimatePresence>
|
||||
);
|
||||
|
||||
@@ -9,6 +9,9 @@ import { useFocusTrap } from '../../hooks/useFocusTrap';
|
||||
import { editFilterCss } from '../../lib/edits';
|
||||
import { summarizeEdits } from '../../lib/versions';
|
||||
import { bakeVideo } from '../../lib/videoBake';
|
||||
import { blobToWavBase64, wavBase64ToBlob } from '../../lib/audioCapture';
|
||||
import { useAIProvider } from '../aiContext';
|
||||
import { VoiceButton } from './VoiceButton';
|
||||
import {
|
||||
normalizeSegments,
|
||||
outputDuration,
|
||||
@@ -91,6 +94,9 @@ export function VideoEditor() {
|
||||
const [duration, setDuration] = useState(0);
|
||||
const [playhead, setPlayhead] = useState(0);
|
||||
const [selOverlay, setSelOverlay] = useState<string | null>(null);
|
||||
const provider = useAIProvider();
|
||||
const [denoiseBusy, setDenoiseBusy] = useState(false);
|
||||
const [denoiseErr, setDenoiseErr] = useState<string | null>(null);
|
||||
const [previewH, setPreviewH] = useState(360);
|
||||
const [baking, setBaking] = useState(false);
|
||||
const [progress, setProgress] = useState(0);
|
||||
@@ -209,6 +215,26 @@ export function VideoEditor() {
|
||||
opacity: 1,
|
||||
rotation: 0,
|
||||
});
|
||||
const runVideoDenoise = async () => {
|
||||
if (!provider?.denoiseAudio) return;
|
||||
setDenoiseBusy(true);
|
||||
setDenoiseErr(null);
|
||||
try {
|
||||
const resp = await fetch(item.src);
|
||||
const blob = await resp.blob();
|
||||
// Decode the video's audio track → 48 kHz mono WAV → RunPod denoise → clean WAV.
|
||||
const wav48 = await blobToWavBase64(blob, 48000);
|
||||
const cleaned = await provider.denoiseAudio(wav48);
|
||||
const url = URL.createObjectURL(wavBase64ToBlob(cleaned));
|
||||
update({ audio: { ...edits.audio, denoisedSrc: url } });
|
||||
} catch (e) {
|
||||
setDenoiseErr(
|
||||
e instanceof Error ? e.message : 'Could not clean the audio (keep clips under ~30s).',
|
||||
);
|
||||
} finally {
|
||||
setDenoiseBusy(false);
|
||||
}
|
||||
};
|
||||
const addKeyframe = (id: string) => {
|
||||
const o = overlays.find((x) => x.id === id);
|
||||
if (!o) return;
|
||||
@@ -314,6 +340,21 @@ export function VideoEditor() {
|
||||
};
|
||||
|
||||
const filterCss = editFilterCss(edits);
|
||||
// Live preview for Crop & Rotate. videoBake applies these on export; without
|
||||
// mirroring them on the preview the Rotate/Flip/Crop buttons look like they do
|
||||
// nothing. Quarter-turn rotations are scaled to fit the frame.
|
||||
const pvRot = (((edits.rotation ?? 0) % 360) + 360) % 360;
|
||||
const pvQuarter = pvRot === 90 || pvRot === 270;
|
||||
const pvAspect = (item.width || 16) / (item.height || 9);
|
||||
const pvFit = pvQuarter ? Math.min(pvAspect, 1 / pvAspect) : 1;
|
||||
const previewTransform =
|
||||
pvRot || edits.flipH || edits.flipV
|
||||
? `rotate(${pvRot}deg) scale(${(edits.flipH ? -1 : 1) * pvFit}, ${(edits.flipV ? -1 : 1) * pvFit})`
|
||||
: undefined;
|
||||
const pvCrop = edits.crop;
|
||||
const previewClip = pvCrop
|
||||
? `inset(${(pvCrop.y * 100).toFixed(3)}% ${((1 - pvCrop.x - pvCrop.width) * 100).toFixed(3)}% ${((1 - pvCrop.y - pvCrop.height) * 100).toFixed(3)}% ${(pvCrop.x * 100).toFixed(3)}%)`
|
||||
: undefined;
|
||||
const sel = overlays.find((o) => o.id === selOverlay) ?? null;
|
||||
|
||||
return (
|
||||
@@ -360,7 +401,15 @@ export function VideoEditor() {
|
||||
controls
|
||||
playsInline
|
||||
crossOrigin="anonymous"
|
||||
style={{ maxWidth: '100%', maxHeight: '70vh', display: 'block', filter: filterCss || undefined }}
|
||||
style={{
|
||||
maxWidth: '100%',
|
||||
maxHeight: '70vh',
|
||||
display: 'block',
|
||||
filter: filterCss || undefined,
|
||||
transform: previewTransform,
|
||||
clipPath: previewClip,
|
||||
transition: 'transform 0.15s ease',
|
||||
}}
|
||||
onLoadedMetadata={(e) => {
|
||||
const d = e.currentTarget.duration || 0;
|
||||
setDuration(d);
|
||||
@@ -651,6 +700,13 @@ export function VideoEditor() {
|
||||
onChange={(e) => patchOverlay(sel.id, { text: e.target.value })}
|
||||
/>
|
||||
</label>
|
||||
<VoiceButton
|
||||
label="Speak → text"
|
||||
onText={(t) => {
|
||||
const cur = sel.text && sel.text !== 'Your text' ? sel.text : '';
|
||||
patchOverlay(sel.id, { text: cur ? `${cur} ${t}` : t });
|
||||
}}
|
||||
/>
|
||||
<div style={{ display: 'flex', gap: 6, margin: '4px 0' }}>
|
||||
{COLORS.map((c) => (
|
||||
<button
|
||||
@@ -794,6 +850,33 @@ export function VideoEditor() {
|
||||
onChange={(e) => update({ audio: { ...edits.audio, muted: e.target.checked } })} />
|
||||
Mute original audio
|
||||
</label>
|
||||
{provider?.denoiseAudio ? (
|
||||
edits.audio?.denoisedSrc ? (
|
||||
<div style={{ display: 'flex', alignItems: 'center', gap: 8, fontSize: 13 }}>
|
||||
<span style={{ color: '#34c759' }}>✓ Background noise reduced (AI)</span>
|
||||
<button
|
||||
type="button"
|
||||
className="apg-btn apg-btn--small"
|
||||
onClick={() => update({ audio: { ...edits.audio, denoisedSrc: undefined } })}
|
||||
>
|
||||
Undo
|
||||
</button>
|
||||
</div>
|
||||
) : (
|
||||
<button
|
||||
type="button"
|
||||
className="apg-btn apg-btn--primary"
|
||||
disabled={denoiseBusy}
|
||||
onClick={() => void runVideoDenoise()}
|
||||
>
|
||||
<Icon name="wand" size={14} />{' '}
|
||||
{denoiseBusy ? 'Cleaning audio…' : 'Reduce background noise (AI)'}
|
||||
</button>
|
||||
)
|
||||
) : null}
|
||||
{denoiseErr ? (
|
||||
<p style={{ color: '#ff6b6b', fontSize: 12, margin: 0 }}>{denoiseErr}</p>
|
||||
) : null}
|
||||
{!edits.audio?.muted ? (
|
||||
<label className="apg-vedit__row">
|
||||
<span>Original volume {Math.round((edits.audio?.originalVolume ?? 1) * 100)}%</span>
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
'use client';
|
||||
|
||||
import { useRef, useState } from 'react';
|
||||
|
||||
import { blobToWavBase64, startRecording, wavBase64ToBlob, type Recorder } from '../../lib/audioCapture';
|
||||
import { Icon } from '../../icons';
|
||||
import { useAIProvider } from '../aiContext';
|
||||
|
||||
/**
|
||||
* Reusable dictation button: record → (optional AI denoise) → transcribe → onText.
|
||||
* Renders nothing if the AI provider can't transcribe. Used by the image-markup
|
||||
* Text tool, the video Text overlay, and the Info comment box.
|
||||
*/
|
||||
export function VoiceButton({
|
||||
onText,
|
||||
denoise = false,
|
||||
label = 'Speak',
|
||||
size = 14,
|
||||
}: {
|
||||
onText: (text: string) => void;
|
||||
denoise?: boolean;
|
||||
label?: string;
|
||||
size?: number;
|
||||
}) {
|
||||
const provider = useAIProvider();
|
||||
const [recording, setRecording] = useState(false);
|
||||
const [status, setStatus] = useState<string | null>(null);
|
||||
const recRef = useRef<Recorder | null>(null);
|
||||
|
||||
if (!provider?.transcribeAudio) return null;
|
||||
|
||||
const start = async () => {
|
||||
setStatus(null);
|
||||
try {
|
||||
recRef.current = await startRecording();
|
||||
setRecording(true);
|
||||
} catch (e) {
|
||||
setStatus(e instanceof Error ? e.message : 'Microphone unavailable.');
|
||||
}
|
||||
};
|
||||
|
||||
const stop = async () => {
|
||||
const rec = recRef.current;
|
||||
recRef.current = null;
|
||||
setRecording(false);
|
||||
if (!rec || !provider.transcribeAudio) return;
|
||||
try {
|
||||
const blob = await rec.stop();
|
||||
let wav: string;
|
||||
if (denoise && provider.denoiseAudio) {
|
||||
setStatus('Reducing noise…');
|
||||
const w48 = await blobToWavBase64(blob, 48000);
|
||||
const cleaned = await provider.denoiseAudio(w48);
|
||||
wav = await blobToWavBase64(wavBase64ToBlob(cleaned), 16000);
|
||||
} else {
|
||||
wav = await blobToWavBase64(blob, 16000);
|
||||
}
|
||||
setStatus('Transcribing…');
|
||||
const t = (await provider.transcribeAudio(wav)).trim();
|
||||
if (t) onText(t);
|
||||
setStatus(null);
|
||||
} catch (e) {
|
||||
setStatus(e instanceof Error ? e.message : 'Could not transcribe.');
|
||||
}
|
||||
};
|
||||
|
||||
return (
|
||||
<span className="apg-voice">
|
||||
<button
|
||||
type="button"
|
||||
className={`apg-btn apg-btn--small apg-voice__mic${recording ? ' apg-voice__mic--rec' : ''}`}
|
||||
onClick={recording ? stop : start}
|
||||
aria-label={recording ? 'Stop recording' : 'Dictate text'}
|
||||
title={recording ? 'Stop & transcribe' : 'Speak to type'}
|
||||
>
|
||||
<Icon name={recording ? 'check' : 'mic'} size={size} />
|
||||
{recording ? 'Stop' : label}
|
||||
</button>
|
||||
{recording || status ? (
|
||||
<span className="apg-voice__status" aria-live="polite">
|
||||
{recording ? '● Listening…' : status}
|
||||
</span>
|
||||
) : null}
|
||||
</span>
|
||||
);
|
||||
}
|
||||
@@ -4,11 +4,13 @@ import { useState } from 'react';
|
||||
|
||||
import { formatDay } from '../../lib/format';
|
||||
import { groupByTime } from '../../lib/grouping';
|
||||
import { resolveLabel } from '../../lib/smartAlbums';
|
||||
import { Icon, type IconName } from '../../icons';
|
||||
import { useGallery, useGalleryStoreApi } from '../../store/context';
|
||||
import { albumMedia, liveMedia, objectLabelCounts } from '../../store/selectors';
|
||||
import type { MediaItem, ViewId } from '../../types';
|
||||
import { promptAlbumName } from '../modals';
|
||||
import { openContextMenu } from '../ContextMenu';
|
||||
import { confirmAction, promptAlbumName } from '../modals';
|
||||
|
||||
function SectionHeader({
|
||||
title,
|
||||
@@ -97,6 +99,7 @@ export function CollectionsView() {
|
||||
const api = useGalleryStoreApi();
|
||||
const media = useGallery((s) => s.media);
|
||||
const albums = useGallery((s) => s.albums);
|
||||
const labelAliases = useGallery((s) => s.labelAliases);
|
||||
const live = liveMedia(media);
|
||||
|
||||
const first = (pred: (m: MediaItem) => boolean) => live.find(pred);
|
||||
@@ -105,10 +108,38 @@ export function CollectionsView() {
|
||||
|
||||
const recentDays = groupByTime(live, 'day').slice(0, 8);
|
||||
const featured = [...live].sort((a, b) => b.takenAt - a.takenAt).slice(0, 12);
|
||||
const objectEntries = [...objectLabelCounts(media).entries()]
|
||||
const objectEntries = [...objectLabelCounts(media, labelAliases).entries()]
|
||||
.sort((a, b) => b[1] - a[1])
|
||||
.slice(0, 14);
|
||||
|
||||
// Right-click an object card → permanently rename its tag (car → excavator).
|
||||
const renameTag = (label: string) => (e: React.MouseEvent) => {
|
||||
e.preventDefault();
|
||||
openContextMenu(e.clientX, e.clientY, [
|
||||
{
|
||||
label: 'Rename Tag',
|
||||
icon: 'tag',
|
||||
onClick: () =>
|
||||
promptAlbumName('Rename Tag', label, (name) => api.getState().renameLabel(label, name), {
|
||||
placeholder: 'Tag name',
|
||||
}),
|
||||
},
|
||||
{
|
||||
label: 'Delete Tag',
|
||||
icon: 'trash',
|
||||
danger: true,
|
||||
onClick: () =>
|
||||
confirmAction({
|
||||
title: 'Delete Tag',
|
||||
message: `Remove the "${label}" tag? It's deleted from all photos and won't be created again.`,
|
||||
confirmLabel: 'Delete',
|
||||
danger: true,
|
||||
onConfirm: () => api.getState().deleteLabel(label),
|
||||
}),
|
||||
},
|
||||
]);
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="apg-scroll">
|
||||
<div className="apg-collections">
|
||||
@@ -184,10 +215,15 @@ export function CollectionsView() {
|
||||
type="button"
|
||||
className="apg-pinned-card"
|
||||
onClick={() => api.getState().setView(`sys:obj:${label}` as ViewId)}
|
||||
onContextMenu={renameTag(label)}
|
||||
aria-label={`${count} photos containing ${label}`}
|
||||
>
|
||||
{(() => {
|
||||
const cover = live.find((m) => m.objectLabels.includes(label));
|
||||
// Match on the resolved label so the cover works for items still
|
||||
// stored under the original detector label.
|
||||
const cover = live.find((m) =>
|
||||
m.objectLabels.some((l) => resolveLabel(l, labelAliases) === label),
|
||||
);
|
||||
return cover ? <img src={cover.thumbnail ?? cover.src} alt="" draggable={false} /> : null;
|
||||
})()}
|
||||
<span className="apg-pinned-card__label" style={{ textTransform: 'capitalize' }}>
|
||||
|
||||
@@ -55,7 +55,8 @@ export type IconName =
|
||||
| 'pip'
|
||||
| 'tag'
|
||||
| 'document'
|
||||
| 'pin';
|
||||
| 'pin'
|
||||
| 'mic';
|
||||
|
||||
export interface IconProps extends SVGProps<SVGSVGElement> {
|
||||
name: IconName;
|
||||
@@ -83,6 +84,13 @@ export function Icon({ name, size = 20, ...rest }: IconProps) {
|
||||
}
|
||||
|
||||
const paths: Record<IconName, JSX.Element> = {
|
||||
mic: (
|
||||
<>
|
||||
<rect x="9" y="3" width="6" height="11" rx="3" />
|
||||
<path d="M6 11a6 6 0 0 0 12 0" />
|
||||
<path d="M12 17v3M9 20.5h6" />
|
||||
</>
|
||||
),
|
||||
library: (
|
||||
<>
|
||||
<rect x="3" y="6" width="18" height="13" rx="2.5" />
|
||||
|
||||
@@ -0,0 +1,138 @@
|
||||
/**
|
||||
* Microphone capture + WAV (PCM16) encoding — browser-only, no external deps.
|
||||
*
|
||||
* Powers the voice-annotation UI: record speech, optionally run it through the
|
||||
* audio-denoise model, then transcribe. Records via MediaRecorder, then decodes
|
||||
* and resamples with the Web Audio API to the mono sample rate each model wants
|
||||
* (16 kHz for speech-to-text, 48 kHz for denoise) and encodes 16-bit PCM WAV.
|
||||
*
|
||||
* All browser globals are touched at call time (never module top-level), so
|
||||
* importing this on the server is safe.
|
||||
*/
|
||||
|
||||
export interface Recorder {
|
||||
/** Stop recording, release the mic, and resolve with the recorded audio blob. */
|
||||
stop(): Promise<Blob>;
|
||||
/** Abort without producing a blob (still releases the mic). */
|
||||
cancel(): void;
|
||||
}
|
||||
|
||||
/** Begin recording from the default microphone. Rejects if mic access is denied. */
|
||||
export async function startRecording(): Promise<Recorder> {
|
||||
if (typeof navigator === 'undefined' || !navigator.mediaDevices?.getUserMedia) {
|
||||
throw new Error('Microphone is not available in this browser.');
|
||||
}
|
||||
const stream = await navigator.mediaDevices.getUserMedia({ audio: true });
|
||||
const recorder = new MediaRecorder(stream);
|
||||
const chunks: BlobPart[] = [];
|
||||
recorder.ondataavailable = (e) => {
|
||||
if (e.data && e.data.size > 0) chunks.push(e.data);
|
||||
};
|
||||
const release = () => stream.getTracks().forEach((t) => t.stop());
|
||||
recorder.start();
|
||||
|
||||
const collect = () => new Blob(chunks, { type: recorder.mimeType || 'audio/webm' });
|
||||
|
||||
return {
|
||||
stop() {
|
||||
return new Promise<Blob>((resolve) => {
|
||||
recorder.onstop = () => {
|
||||
release();
|
||||
resolve(collect());
|
||||
};
|
||||
try {
|
||||
recorder.stop();
|
||||
} catch {
|
||||
release();
|
||||
resolve(collect());
|
||||
}
|
||||
});
|
||||
},
|
||||
cancel() {
|
||||
try {
|
||||
recorder.stop();
|
||||
} catch {
|
||||
/* ignore */
|
||||
}
|
||||
release();
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Decode any recorded audio blob, downmix to mono, resample to `sampleRate`, and
|
||||
* return base64 (no `data:` prefix) of a 16-bit PCM WAV.
|
||||
*/
|
||||
export async function blobToWavBase64(blob: Blob, sampleRate: number): Promise<string> {
|
||||
const arrayBuf = await blob.arrayBuffer();
|
||||
const AudioCtx =
|
||||
typeof window !== 'undefined'
|
||||
? window.AudioContext ||
|
||||
(window as unknown as { webkitAudioContext?: typeof AudioContext }).webkitAudioContext
|
||||
: undefined;
|
||||
if (!AudioCtx) throw new Error('Web Audio is not supported in this browser.');
|
||||
|
||||
const decodeCtx = new AudioCtx();
|
||||
let decoded: AudioBuffer;
|
||||
try {
|
||||
// slice(0) — decodeAudioData detaches the buffer; keep a copy.
|
||||
decoded = await decodeCtx.decodeAudioData(arrayBuf.slice(0));
|
||||
} finally {
|
||||
void decodeCtx.close();
|
||||
}
|
||||
|
||||
// Downmix (multi-channel → 1) + resample via an OfflineAudioContext at the target rate.
|
||||
const frames = Math.max(1, Math.round(decoded.duration * sampleRate));
|
||||
const offline = new OfflineAudioContext(1, frames, sampleRate);
|
||||
const source = offline.createBufferSource();
|
||||
source.buffer = decoded;
|
||||
source.connect(offline.destination);
|
||||
source.start();
|
||||
const rendered = await offline.startRendering();
|
||||
return base64FromBytes(new Uint8Array(encodeWavPcm16(rendered.getChannelData(0), sampleRate)));
|
||||
}
|
||||
|
||||
/** Turn base64 WAV (as returned by the denoise model) back into a Blob. */
|
||||
export function wavBase64ToBlob(base64: string): Blob {
|
||||
const clean = base64.includes(',') ? base64.split(',', 2)[1]! : base64;
|
||||
const bin = atob(clean);
|
||||
const bytes = new Uint8Array(bin.length);
|
||||
for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i);
|
||||
return new Blob([bytes], { type: 'audio/wav' });
|
||||
}
|
||||
|
||||
function encodeWavPcm16(samples: Float32Array, sampleRate: number): ArrayBuffer {
|
||||
const buffer = new ArrayBuffer(44 + samples.length * 2);
|
||||
const view = new DataView(buffer);
|
||||
const writeStr = (offset: number, s: string) => {
|
||||
for (let i = 0; i < s.length; i++) view.setUint8(offset + i, s.charCodeAt(i));
|
||||
};
|
||||
writeStr(0, 'RIFF');
|
||||
view.setUint32(4, 36 + samples.length * 2, true);
|
||||
writeStr(8, 'WAVE');
|
||||
writeStr(12, 'fmt ');
|
||||
view.setUint32(16, 16, true); // fmt chunk size
|
||||
view.setUint16(20, 1, true); // audio format = PCM
|
||||
view.setUint16(22, 1, true); // channels = mono
|
||||
view.setUint32(24, sampleRate, true);
|
||||
view.setUint32(28, sampleRate * 2, true); // byte rate = sampleRate * blockAlign
|
||||
view.setUint16(32, 2, true); // block align = channels * bytesPerSample
|
||||
view.setUint16(34, 16, true); // bits per sample
|
||||
writeStr(36, 'data');
|
||||
view.setUint32(40, samples.length * 2, true);
|
||||
let offset = 44;
|
||||
for (let i = 0; i < samples.length; i++, offset += 2) {
|
||||
const s = Math.max(-1, Math.min(1, samples[i]!));
|
||||
view.setInt16(offset, s < 0 ? s * 0x8000 : s * 0x7fff, true);
|
||||
}
|
||||
return buffer;
|
||||
}
|
||||
|
||||
function base64FromBytes(bytes: Uint8Array): string {
|
||||
let bin = '';
|
||||
const chunk = 0x8000;
|
||||
for (let i = 0; i < bytes.length; i += chunk) {
|
||||
bin += String.fromCharCode(...bytes.subarray(i, i + chunk));
|
||||
}
|
||||
return btoa(bin);
|
||||
}
|
||||
@@ -52,6 +52,9 @@ export function createMediaItem(input: MediaInput): MediaItem {
|
||||
caption: input.caption,
|
||||
objects: input.objects,
|
||||
faces: input.faces,
|
||||
// Preserve OCR text through the normalization path — without this it's dropped
|
||||
// on every reload, so the analyzer re-runs OCR on the whole library each time.
|
||||
ocrText: input.ocrText,
|
||||
analyzedAt: input.analyzedAt,
|
||||
colorPalette: input.colorPalette,
|
||||
blurScore: input.blurScore,
|
||||
@@ -210,7 +213,15 @@ async function readExif(file: File): Promise<ExifMeta> {
|
||||
out.location = { lat: data.latitude, lng: data.longitude };
|
||||
}
|
||||
const exif: Record<string, string | number> = {};
|
||||
for (const k of ['Make', 'Model', 'LensModel', 'ISO', 'FNumber', 'FocalLength'] as const) {
|
||||
for (const k of [
|
||||
'Make',
|
||||
'Model',
|
||||
'LensModel',
|
||||
'ISO',
|
||||
'FNumber',
|
||||
'FocalLength',
|
||||
'Orientation',
|
||||
] as const) {
|
||||
const v = data[k];
|
||||
if (typeof v === 'string' || typeof v === 'number') exif[k] = v;
|
||||
}
|
||||
|
||||
@@ -1,7 +1,19 @@
|
||||
import type { Album, MediaItem, SmartRule, SmartRuleSet } from '../types';
|
||||
|
||||
/** The user's permanent tag-rename map: canonical lowercased detector label → chosen label. */
|
||||
export type LabelAliases = Record<string, string>;
|
||||
|
||||
/**
|
||||
* Resolve a detected object label through the user's rename alias map. Keys are
|
||||
* canonical lowercased detector labels ("car"); the value is the label the user
|
||||
* renamed it to ("excavator"). Returns the original label when no alias applies.
|
||||
*/
|
||||
export function resolveLabel(label: string, aliases: LabelAliases = {}): string {
|
||||
return aliases[label.trim().toLowerCase()] ?? label;
|
||||
}
|
||||
|
||||
/** Evaluate a single smart rule against a media item. */
|
||||
function evalRule(item: MediaItem, rule: SmartRule): boolean {
|
||||
function evalRule(item: MediaItem, rule: SmartRule, aliases: LabelAliases = {}): boolean {
|
||||
const { field, op, value } = rule;
|
||||
|
||||
const get = (): unknown => {
|
||||
@@ -18,7 +30,9 @@ function evalRule(item: MediaItem, rule: SmartRule): boolean {
|
||||
case 'isLivePhoto': return Boolean(item.isLivePhoto);
|
||||
case 'isPanorama': return Boolean(item.isPanorama);
|
||||
case 'tag': return item.tags;
|
||||
case 'object': return item.objectLabels;
|
||||
// Resolve object labels through the rename map so a photo detected as the
|
||||
// original label still matches the renamed album's rule.
|
||||
case 'object': return item.objectLabels.map((l) => resolveLabel(l, aliases));
|
||||
case 'person': return item.personIds;
|
||||
default: return undefined;
|
||||
}
|
||||
@@ -48,18 +62,26 @@ function evalRule(item: MediaItem, rule: SmartRule): boolean {
|
||||
}
|
||||
|
||||
/** Does an item satisfy a rule set (AND/OR over its rules)? */
|
||||
export function matchesRuleSet(item: MediaItem, ruleSet: SmartRuleSet): boolean {
|
||||
export function matchesRuleSet(
|
||||
item: MediaItem,
|
||||
ruleSet: SmartRuleSet,
|
||||
aliases: LabelAliases = {},
|
||||
): boolean {
|
||||
if (item.deletedAt) return false; // trashed items never appear in smart albums
|
||||
if (ruleSet.rules.length === 0) return true;
|
||||
return ruleSet.match === 'all'
|
||||
? ruleSet.rules.every((r) => evalRule(item, r))
|
||||
: ruleSet.rules.some((r) => evalRule(item, r));
|
||||
? ruleSet.rules.every((r) => evalRule(item, r, aliases))
|
||||
: ruleSet.rules.some((r) => evalRule(item, r, aliases));
|
||||
}
|
||||
|
||||
/** Resolve the live member ids of a smart album from the full library. */
|
||||
export function resolveSmartAlbum(album: Album, items: MediaItem[]): string[] {
|
||||
export function resolveSmartAlbum(
|
||||
album: Album,
|
||||
items: MediaItem[],
|
||||
aliases: LabelAliases = {},
|
||||
): string[] {
|
||||
if (!album.ruleSet) return [];
|
||||
return items.filter((i) => matchesRuleSet(i, album.ruleSet!)).map((i) => i.id);
|
||||
return items.filter((i) => matchesRuleSet(i, album.ruleSet!, aliases)).map((i) => i.id);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -141,12 +163,18 @@ function titleCase(s: string): string {
|
||||
*
|
||||
* @param minCount only surface a label once it appears on at least this many photos.
|
||||
*/
|
||||
export function objectSmartAlbums(media: MediaItem[], now: number, minCount = 1): Album[] {
|
||||
export function objectSmartAlbums(
|
||||
media: MediaItem[],
|
||||
now: number,
|
||||
aliases: LabelAliases = {},
|
||||
minCount = 1,
|
||||
): Album[] {
|
||||
const counts = new Map<string, number>();
|
||||
for (const m of media) {
|
||||
if (m.deletedAt || m.hidden) continue;
|
||||
// De-dupe labels within one item so a single photo counts once per label.
|
||||
for (const label of new Set(m.objectLabels)) {
|
||||
// Resolve labels through the rename map, then de-dupe within one item so a
|
||||
// single photo counts once per (renamed) label and renamed labels regroup.
|
||||
for (const label of new Set(m.objectLabels.map((l) => resolveLabel(l, aliases)))) {
|
||||
const key = label.trim().toLowerCase();
|
||||
if (!key) continue;
|
||||
counts.set(key, (counts.get(key) ?? 0) + 1);
|
||||
|
||||
@@ -31,6 +31,7 @@ export function summarizeEdits(edits?: EditState): string[] {
|
||||
if (edits.overlays?.some((o) => o.kind === 'text')) c.push('Text');
|
||||
}
|
||||
if (edits.audio?.muted) c.push('Muted original audio');
|
||||
if (edits.audio?.denoisedSrc) c.push('Reduced audio noise (AI)');
|
||||
if (edits.audio?.musicSrc) c.push('Added music');
|
||||
if (edits.audio?.fadeIn || edits.audio?.fadeOut) c.push('Audio fade');
|
||||
return c.length ? c : ['Edited'];
|
||||
|
||||
@@ -173,7 +173,10 @@ export async function bakeVideo(
|
||||
try {
|
||||
const vNode = ac.createMediaElementSource(video);
|
||||
const vGain = ac.createGain();
|
||||
vGain.gain.value = edits.audio?.muted ? 0 : (edits.audio?.originalVolume ?? 1);
|
||||
// When an AI-denoised track was produced, mute the original and play the clean
|
||||
// one instead (added below); otherwise use the original at its set volume.
|
||||
const hasClean = Boolean(edits.audio?.denoisedSrc);
|
||||
vGain.gain.value = edits.audio?.muted || hasClean ? 0 : (edits.audio?.originalVolume ?? 1);
|
||||
vNode.connect(vGain).connect(master);
|
||||
} catch {
|
||||
/* element may have no audio track */
|
||||
@@ -192,6 +195,21 @@ export async function bakeVideo(
|
||||
}
|
||||
}
|
||||
|
||||
// AI-denoised original audio (RunPod): play it in place of the muted original.
|
||||
let denoised: HTMLAudioElement | null = null;
|
||||
if (edits.audio?.denoisedSrc && !edits.audio?.muted) {
|
||||
denoised = new Audio(edits.audio.denoisedSrc);
|
||||
denoised.crossOrigin = 'anonymous';
|
||||
try {
|
||||
const dNode = ac.createMediaElementSource(denoised);
|
||||
const dGain = ac.createGain();
|
||||
dGain.gain.value = edits.audio.originalVolume ?? 1;
|
||||
dNode.connect(dGain).connect(master);
|
||||
} catch {
|
||||
/* ignore */
|
||||
}
|
||||
}
|
||||
|
||||
// 6) Recorder over canvas video + mixed audio.
|
||||
const stream = canvas.captureStream(fps);
|
||||
const audioTrack = dest.stream.getAudioTracks()[0];
|
||||
@@ -283,6 +301,7 @@ export async function bakeVideo(
|
||||
await seekTo(video, seg.start);
|
||||
video.playbackRate = seg.speed || 1;
|
||||
if (music && seg === segs[0]) await music.play().catch(() => {});
|
||||
if (denoised && seg === segs[0]) await denoised.play().catch(() => {});
|
||||
await video.play().catch(() => {});
|
||||
const segOut = (seg.end - seg.start) / (seg.speed || 1);
|
||||
// Guard against a stalled decoder (autoplay blocked, boundary glitch).
|
||||
@@ -310,6 +329,7 @@ export async function bakeVideo(
|
||||
}
|
||||
|
||||
music?.pause();
|
||||
denoised?.pause();
|
||||
if (recorder.state !== 'inactive') recorder.stop();
|
||||
await stopped;
|
||||
opts.signal?.removeEventListener('abort', onAbort);
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { matchesRuleSet, resolveSmartAlbum } from '../lib/smartAlbums';
|
||||
import { matchesRuleSet, resolveLabel, resolveSmartAlbum } from '../lib/smartAlbums';
|
||||
import type { LabelAliases } from '../lib/smartAlbums';
|
||||
import type { Album, MediaItem } from '../types';
|
||||
import type { GalleryState, GridFilter } from './store';
|
||||
|
||||
@@ -18,9 +19,9 @@ export function albumById(albums: Album[], id: string): Album | undefined {
|
||||
}
|
||||
|
||||
/** Resolve the members of any album (smart albums are computed live). */
|
||||
export function albumMedia(album: Album, media: MediaItem[]): MediaItem[] {
|
||||
export function albumMedia(album: Album, media: MediaItem[], aliases: LabelAliases = {}): MediaItem[] {
|
||||
if (album.kind === 'smart' && album.ruleSet) {
|
||||
const ids = new Set(resolveSmartAlbum(album, media));
|
||||
const ids = new Set(resolveSmartAlbum(album, media, aliases));
|
||||
return liveMedia(media)
|
||||
.filter((m) => ids.has(m.id))
|
||||
.sort((a, b) => b.takenAt - a.takenAt);
|
||||
@@ -73,7 +74,7 @@ export function searchMedia(items: MediaItem[], query: string): MediaItem[] {
|
||||
* grid filter, search and object-focus. Returns reverse-chronological order.
|
||||
*/
|
||||
export function mediaForView(state: GalleryState): MediaItem[] {
|
||||
const { media, albums, view, gridFilter, searchQuery, objectFocus, tagFocus, personFocus, searchPreset, semanticResults, recentlyViewed } =
|
||||
const { media, albums, view, gridFilter, searchQuery, objectFocus, tagFocus, personFocus, searchPreset, semanticResults, recentlyViewed, labelAliases } =
|
||||
state;
|
||||
let items: MediaItem[];
|
||||
// When a search ranks results (keyword + semantic), preserve that order
|
||||
@@ -103,7 +104,7 @@ export function mediaForView(state: GalleryState): MediaItem[] {
|
||||
items = liveMedia(media);
|
||||
} else if (view.startsWith('album:') || view.startsWith('sys:')) {
|
||||
const album = albumById(albums, view);
|
||||
items = album ? albumMedia(album, media) : [];
|
||||
items = album ? albumMedia(album, media, labelAliases) : [];
|
||||
} else if (view === 'favourites') {
|
||||
items = liveMedia(media).filter((m) => m.favorite);
|
||||
} else if (view === 'recently-saved') {
|
||||
@@ -120,7 +121,10 @@ export function mediaForView(state: GalleryState): MediaItem[] {
|
||||
items = applyGridFilter(items, gridFilter);
|
||||
}
|
||||
if (objectFocus) {
|
||||
items = items.filter((m) => m.objectLabels.includes(objectFocus));
|
||||
// Match on the resolved (renamed) label so focusing a renamed object still
|
||||
// surfaces photos whose stored label is the original detector label.
|
||||
const focus = resolveLabel(objectFocus, labelAliases);
|
||||
items = items.filter((m) => m.objectLabels.some((l) => resolveLabel(l, labelAliases) === focus));
|
||||
}
|
||||
if (tagFocus) {
|
||||
items = items.filter((m) => m.tags.includes(tagFocus));
|
||||
@@ -153,11 +157,14 @@ export function mediaForView(state: GalleryState): MediaItem[] {
|
||||
return items;
|
||||
}
|
||||
|
||||
/** Distinct object labels across the live library (for AI auto-albums). */
|
||||
export function objectLabelCounts(media: MediaItem[]): Map<string, number> {
|
||||
/** Distinct object labels across the live library (for AI auto-albums), resolved
|
||||
* through the user's rename map so renamed tags collapse into one bucket. */
|
||||
export function objectLabelCounts(media: MediaItem[], aliases: LabelAliases = {}): Map<string, number> {
|
||||
const counts = new Map<string, number>();
|
||||
for (const m of liveMedia(media)) {
|
||||
for (const label of m.objectLabels) {
|
||||
// De-dupe per item after resolving so two labels renamed to the same value
|
||||
// only count that photo once.
|
||||
for (const label of new Set(m.objectLabels.map((l) => resolveLabel(l, aliases)))) {
|
||||
counts.set(label, (counts.get(label) ?? 0) + 1);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -97,6 +97,14 @@ export interface GalleryState {
|
||||
media: MediaItem[];
|
||||
albums: Album[];
|
||||
people: Person[];
|
||||
/**
|
||||
* User's permanent object-tag renames. Key = canonical lowercased detector
|
||||
* label (e.g. "car"), value = the label the user renamed it to ("excavator").
|
||||
* Persisted, applied to future uploads, and resolved wherever labels are grouped.
|
||||
*/
|
||||
labelAliases: Record<string, string>;
|
||||
/** Object/material tags the user deleted — stripped from photos + hidden from future detection. */
|
||||
deletedLabels: string[];
|
||||
|
||||
// Navigation / view
|
||||
view: ViewId;
|
||||
@@ -200,6 +208,14 @@ export interface GalleryState {
|
||||
syncObjectAlbums: () => void;
|
||||
/** Rename a person/pet group (empty name clears it back to unnamed). */
|
||||
renamePerson: (id: PersonId, name: string) => void;
|
||||
/**
|
||||
* Permanently rename a detected object tag (e.g. "car" → "excavator"). Records a
|
||||
* persisted alias, backfills existing items' objectLabels, and rebuilds object
|
||||
* albums so the rename applies to past photos and every future upload.
|
||||
*/
|
||||
renameLabel: (from: string, to: string) => void;
|
||||
/** Permanently delete an object/material tag (removes it from photos + future detections). */
|
||||
deleteLabel: (label: string) => void;
|
||||
|
||||
// Recently Deleted lock
|
||||
setLockPassword: (password: string) => Promise<void>;
|
||||
@@ -287,10 +303,10 @@ export function createGalleryStore(options: CreateStoreOptions) {
|
||||
if (!adapter) return;
|
||||
if (persistTimer) clearTimeout(persistTimer);
|
||||
persistTimer = setTimeout(() => {
|
||||
const { media, albums, people } = get();
|
||||
const { media, albums, people, labelAliases, deletedLabels } = get();
|
||||
// Only persist user-owned albums; smart/system albums are regenerated.
|
||||
const userAlbums = albums.filter((a) => !a.system);
|
||||
void adapter!.save({ media, albums: userAlbums, people, version: 1 });
|
||||
void adapter!.save({ media, albums: userAlbums, people, labelAliases, deletedLabels, version: 1 });
|
||||
}, 400);
|
||||
};
|
||||
|
||||
@@ -302,6 +318,8 @@ export function createGalleryStore(options: CreateStoreOptions) {
|
||||
media: options.initialMedia ?? [],
|
||||
albums: [...defaultSystemAlbums(now), ...(options.initialAlbums ?? [])],
|
||||
people: options.initialPeople ?? [],
|
||||
labelAliases: {},
|
||||
deletedLabels: [],
|
||||
|
||||
view: 'library',
|
||||
libraryScale: 'all',
|
||||
@@ -349,6 +367,8 @@ export function createGalleryStore(options: CreateStoreOptions) {
|
||||
media: seedEmptyBackend ? seed : loaded.media,
|
||||
albums: [...defaultSystemAlbums(ts), ...loaded.albums.filter((al) => !al.system)],
|
||||
people: loaded.people ?? [],
|
||||
labelAliases: loaded.labelAliases ?? {},
|
||||
deletedLabels: loaded.deletedLabels ?? [],
|
||||
ready: true,
|
||||
aiAvailable: Boolean(ai),
|
||||
});
|
||||
@@ -427,6 +447,9 @@ export function createGalleryStore(options: CreateStoreOptions) {
|
||||
// tags objects/faces/OCR/embeddings in the background.
|
||||
get().addMedia(items);
|
||||
if (albumId) get().addToAlbum(albumId, items.map((i) => i.id));
|
||||
// Surface the just-uploaded photo with its Info panel open, so the user
|
||||
// watches the analysis run + sees its results without hunting for it.
|
||||
set({ infoOpen: true, selection: new Set([items[0]!.id]) });
|
||||
}
|
||||
return items.map((i) => i.id);
|
||||
},
|
||||
@@ -880,7 +903,7 @@ export function createGalleryStore(options: CreateStoreOptions) {
|
||||
// Replace ONLY the sys:obj:* set; keep default system albums + user albums.
|
||||
// Membership stays live via resolveSmartAlbum, so no per-item bookkeeping and
|
||||
// no persist() needed (system albums are stripped on save + regenerated on load).
|
||||
const generated = objectSmartAlbums(get().media, Date.now());
|
||||
const generated = objectSmartAlbums(get().media, Date.now(), get().labelAliases);
|
||||
set((s) => {
|
||||
const others = s.albums.filter((a) => !a.id.startsWith('sys:obj:'));
|
||||
// Skip the state write if the label SET is unchanged (avoids render churn).
|
||||
@@ -899,6 +922,66 @@ export function createGalleryStore(options: CreateStoreOptions) {
|
||||
}));
|
||||
persist();
|
||||
},
|
||||
renameLabel(from, to) {
|
||||
const key = from.trim().toLowerCase();
|
||||
const value = to.trim().toLowerCase();
|
||||
if (!key) return;
|
||||
set((s) => {
|
||||
const labelAliases: Record<string, string> = { ...s.labelAliases };
|
||||
// Empty / unchanged name clears the alias (renames the tag back to itself).
|
||||
const clearing = !value || value === key;
|
||||
if (clearing) {
|
||||
delete labelAliases[key];
|
||||
} else {
|
||||
labelAliases[key] = value;
|
||||
// Re-point any earlier aliases that resolved to `key` so renaming an
|
||||
// already-renamed tag keeps mapping the original detector label(s).
|
||||
for (const k of Object.keys(labelAliases)) {
|
||||
if (k !== key && labelAliases[k] === key) labelAliases[k] = value;
|
||||
}
|
||||
}
|
||||
// Backfill existing items so past photos regroup under the new label and
|
||||
// stay searchable by it (search reads objectLabels directly).
|
||||
const target = clearing ? key : value;
|
||||
const media = s.media.map((m) => {
|
||||
if (!m.objectLabels.some((l) => l.trim().toLowerCase() === key)) return m;
|
||||
const mapped = [
|
||||
...new Set(
|
||||
m.objectLabels.map((l) => (l.trim().toLowerCase() === key ? target : l)),
|
||||
),
|
||||
];
|
||||
return { ...m, objectLabels: mapped };
|
||||
});
|
||||
return { labelAliases, media };
|
||||
});
|
||||
// Rebuild the sys:obj:* albums so the sidebar/Objects re-title live.
|
||||
get().syncObjectAlbums();
|
||||
persist();
|
||||
},
|
||||
|
||||
deleteLabel(label) {
|
||||
const key = label.trim().toLowerCase();
|
||||
if (!key) return;
|
||||
set((s) => {
|
||||
const deletedLabels = s.deletedLabels.includes(key)
|
||||
? s.deletedLabels
|
||||
: [...s.deletedLabels, key];
|
||||
// Drop any aliases pointing at this label so it can't reappear via a rename.
|
||||
const labelAliases: Record<string, string> = { ...s.labelAliases };
|
||||
for (const k of Object.keys(labelAliases)) {
|
||||
if (k === key || labelAliases[k] === key) delete labelAliases[k];
|
||||
}
|
||||
// Strip the label from every photo so its album empties and it stops matching.
|
||||
const media = s.media.map((m) =>
|
||||
m.objectLabels.some((l) => l.trim().toLowerCase() === key)
|
||||
? { ...m, objectLabels: m.objectLabels.filter((l) => l.trim().toLowerCase() !== key) }
|
||||
: m,
|
||||
);
|
||||
return { deletedLabels, labelAliases, media };
|
||||
});
|
||||
get().syncObjectAlbums();
|
||||
persist();
|
||||
},
|
||||
|
||||
async setLockPassword(password) {
|
||||
const hash = await hashPassword(password);
|
||||
|
||||
@@ -1659,6 +1659,55 @@
|
||||
.apg-comment-form__author:focus,
|
||||
.apg-comment-form__input:focus { outline: none; border-color: var(--apg-accent); }
|
||||
.apg-comment-form__post { align-self: flex-end; }
|
||||
.apg-voice { display: flex; align-items: center; flex-wrap: wrap; gap: 8px; }
|
||||
.apg-voice__mic { display: inline-flex; align-items: center; gap: 5px; }
|
||||
.apg-voice__mic--rec {
|
||||
background: var(--apg-danger, #e5484d);
|
||||
border-color: var(--apg-danger, #e5484d);
|
||||
color: #fff;
|
||||
}
|
||||
.apg-voice__denoise {
|
||||
display: inline-flex; align-items: center; gap: 5px;
|
||||
font-size: 12px; opacity: 0.85; cursor: pointer; user-select: none;
|
||||
}
|
||||
.apg-voice__denoise input { accent-color: var(--apg-accent); }
|
||||
.apg-voice__status { font-size: 12px; opacity: 0.7; min-height: 1em; }
|
||||
.apg-info__analyzing {
|
||||
display: flex; align-items: center; gap: 8px;
|
||||
font-size: 13px; font-weight: 500; color: var(--apg-accent);
|
||||
padding: 8px 10px; margin: 2px 0 4px;
|
||||
background: color-mix(in srgb, var(--apg-accent) 10%, transparent);
|
||||
border-radius: 8px;
|
||||
}
|
||||
.apg-info__spinner {
|
||||
width: 14px; height: 14px; flex: none;
|
||||
border: 2px solid currentColor; border-top-color: transparent;
|
||||
border-radius: 50%; animation: apg-spin 0.8s linear infinite;
|
||||
}
|
||||
@keyframes apg-spin { to { transform: rotate(360deg); } }
|
||||
@media (prefers-reduced-motion: reduce) { .apg-info__spinner { animation: none; } }
|
||||
.apg-maskbrush {
|
||||
position: fixed; inset: 0; z-index: 200;
|
||||
background: rgba(0, 0, 0, 0.85);
|
||||
display: flex; flex-direction: column; align-items: center; justify-content: center;
|
||||
gap: 12px; padding: 20px;
|
||||
}
|
||||
.apg-maskbrush__title { color: #fff; font-size: 14px; font-weight: 600; }
|
||||
.apg-maskbrush__stage {
|
||||
position: relative; width: 100%;
|
||||
max-width: min(90vw, 900px); max-height: 68vh;
|
||||
}
|
||||
.apg-maskbrush__img,
|
||||
.apg-maskbrush__canvas {
|
||||
position: absolute; inset: 0; width: 100%; height: 100%; border-radius: 8px;
|
||||
}
|
||||
.apg-maskbrush__img { object-fit: contain; user-select: none; -webkit-user-drag: none; }
|
||||
.apg-maskbrush__canvas { cursor: crosshair; touch-action: none; }
|
||||
.apg-maskbrush__bar {
|
||||
display: flex; align-items: center; gap: 10px; flex-wrap: wrap;
|
||||
width: 100%; max-width: min(90vw, 900px);
|
||||
}
|
||||
.apg-maskbrush__brush { display: inline-flex; align-items: center; gap: 6px; color: #fff; font-size: 12px; }
|
||||
|
||||
/* ---- Versions & Audit browser view ---- */
|
||||
.apg-audit-list { display: flex; flex-direction: column; gap: 8px; }
|
||||
|
||||
@@ -91,6 +91,8 @@ export interface EditState {
|
||||
/** Master fade in / out over the whole exported clip, in seconds. */
|
||||
fadeIn?: number;
|
||||
fadeOut?: number;
|
||||
/** AI-denoised audio (RunPod); replaces the original track on export when set. */
|
||||
denoisedSrc?: string;
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
# RunPod Serverless worker: FLUX.2 [klein] image editing.
|
||||
# CUDA 12.4 base + torch cu124: RunPod's RTX 3090/4090 hosts run driver 550
|
||||
# (CUDA 12.4). A 12.8 base demands driver >= 12.8 and the container REFUSES to
|
||||
# start on a 550 host ("nvidia-container-cli: unsatisfied condition: cuda>=12.8").
|
||||
# 12.4 runs on driver 550+ and carries sm_86 (3090) + sm_89 (4090) kernels.
|
||||
# NOTE: this does NOT carry Blackwell (sm_120) kernels — keep this endpoint's GPU
|
||||
# pool to 3090/4090 (24 GB Ada/Ampere), not the PRO 6000 Blackwell cards.
|
||||
FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
python3 python3-pip git && \
|
||||
ln -sf /usr/bin/python3 /usr/bin/python && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# torch (CUDA 12.4 build) FIRST — matches the 12.4 base + driver 550 hosts.
|
||||
RUN pip install --upgrade pip && \
|
||||
pip install torch --index-url https://download.pytorch.org/whl/cu124
|
||||
COPY runpod-worker/requirements.txt .
|
||||
RUN pip install -r requirements.txt
|
||||
|
||||
COPY runpod-worker/handler.py .
|
||||
|
||||
CMD ["python3", "-u", "handler.py"]
|
||||
@@ -0,0 +1,49 @@
|
||||
# Instruct-Pix2Pix image-editing worker (RunPod Serverless)
|
||||
|
||||
Instruction-based image editing that powers the gallery editor's **Apply Prompt**
|
||||
button. Instruct-Pix2Pix is *trained to follow edit instructions* ("make it
|
||||
golden hour", "add snow", "turn the sky orange") and applies them while keeping
|
||||
the rest of the photo — unlike plain SD img2img.
|
||||
|
||||
Small (~2 GB), runs on your existing 24 GB endpoint, **no gated license, no HF
|
||||
token, no volume resize**. Input/output matches the app, so no app-code change.
|
||||
|
||||
## Contract
|
||||
|
||||
Request → `POST /runsync`:
|
||||
```json
|
||||
{ "input": { "image": "<base64>", "prompt": "<edit instruction>",
|
||||
"guidance_scale": 7.5, "num_inference_steps": 25, "seed": 123 } }
|
||||
```
|
||||
Response → `{ "output": { "image": "<base64 PNG>" } }` (or `{ "output": { "error": "..." } }`).
|
||||
|
||||
`strength` / `mask` from the app are ignored (this model conditions on the image
|
||||
via `image_guidance_scale`, not img2img noise or masks).
|
||||
|
||||
## Deploy (reuse your existing endpoint)
|
||||
|
||||
No RunPod settings to change — just push, and RunPod auto-rebuilds:
|
||||
```bash
|
||||
cd "advance-photo-gallery-web-sdk - Copy"
|
||||
git add runpod-worker
|
||||
git commit -m "Switch worker to Instruct-Pix2Pix"
|
||||
git push
|
||||
```
|
||||
Then **warm up once** (downloads ~2 GB — quick):
|
||||
```bash
|
||||
curl -X POST https://api.runpod.ai/v2/<ENDPOINT_ID>/runsync \
|
||||
-H "Content-Type: application/json" -H "Authorization: Bearer <API_KEY>" \
|
||||
-d @runpod-worker/warmup.json
|
||||
```
|
||||
Test **Apply Prompt** in the app — no Vercel change (same URL).
|
||||
|
||||
## Tuning (optional endpoint env vars)
|
||||
|
||||
| Var | Default | Purpose |
|
||||
|---|---|---|
|
||||
| `IP2P_MODEL` | `timbrooks/instruct-pix2pix` | the model |
|
||||
| `IP2P_IMAGE_GUIDANCE` | `1.5` | higher keeps more of the original; **lower toward ~1.2 if edits feel too weak** |
|
||||
| `IP2P_MAX_SIZE` | `768` | max long-edge px (SD1.5 works best ≤ 768) |
|
||||
|
||||
**Prompts work best as imperatives:** "make it a golden-hour sunset", "add snow",
|
||||
"turn it into winter" — not descriptions.
|
||||
@@ -0,0 +1,113 @@
|
||||
"""
|
||||
RunPod Serverless worker: FLUX.2 [klein] 4B instruction image editing (img2img).
|
||||
|
||||
Replaces the old Instruct-Pix2Pix worker. FLUX.2-klein is Apache-2.0 (no HF token),
|
||||
runs on a 24 GB card (RTX 3090), and is distilled → ~6 steps, fast. It needs a
|
||||
recent torch, which also carries kernels for the newest (Blackwell) GPUs — so it
|
||||
runs on whatever card RunPod assigns.
|
||||
|
||||
Matches the app's img2img contract (no app-code change needed):
|
||||
{"input": {
|
||||
"image": "<base64>", # required (raw base64 or data: URI)
|
||||
"prompt": "<edit instruction>", # required
|
||||
"strength": 0.6, # optional — the editor's "Edit strength" slider (mapped to guidance)
|
||||
"seed": 123 # optional
|
||||
}}
|
||||
Returns: {"image": "<base64 PNG>"} (or {"error": "..."}).
|
||||
|
||||
Tuning (env): FLUX2_MODEL (default black-forest-labs/FLUX.2-klein-4B),
|
||||
FLUX2_STEPS (default 6), FLUX2_MAX_SIZE (default 1024).
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
|
||||
_V = "/runpod-volume"
|
||||
if os.path.isdir(_V):
|
||||
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
|
||||
|
||||
import torch # noqa: E402
|
||||
import runpod # noqa: E402
|
||||
from PIL import Image # noqa: E402
|
||||
from diffusers import Flux2KleinPipeline # noqa: E402
|
||||
|
||||
MODEL = os.environ.get("FLUX2_MODEL", "black-forest-labs/FLUX.2-klein-4B")
|
||||
STEPS = int(os.environ.get("FLUX2_STEPS", "6"))
|
||||
MAX_SIZE = int(os.environ.get("FLUX2_MAX_SIZE", "1024"))
|
||||
|
||||
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
_pipe = None
|
||||
|
||||
|
||||
def _get_pipe():
|
||||
global _pipe
|
||||
if _pipe is None:
|
||||
pipe = Flux2KleinPipeline.from_pretrained(MODEL, torch_dtype=torch.bfloat16)
|
||||
total_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 if DEVICE == "cuda" else 0
|
||||
if total_gb >= 20:
|
||||
pipe = pipe.to("cuda")
|
||||
else:
|
||||
pipe.enable_model_cpu_offload()
|
||||
pipe.set_progress_bar_config(disable=True)
|
||||
_pipe = pipe
|
||||
return _pipe
|
||||
|
||||
|
||||
def _b64_to_image(data):
|
||||
if "," in data:
|
||||
data = data.split(",", 1)[1]
|
||||
return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGB")
|
||||
|
||||
|
||||
def _image_to_b64(img):
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
def _fit(img):
|
||||
w, h = img.size
|
||||
scale = min(1.0, MAX_SIZE / max(w, h))
|
||||
nw = max(16, (int(w * scale) // 16) * 16)
|
||||
nh = max(16, (int(h * scale) // 16) * 16)
|
||||
return img.resize((nw, nh), Image.LANCZOS)
|
||||
|
||||
|
||||
def handler(event):
|
||||
inp = (event or {}).get("input") or {}
|
||||
image_b64 = inp.get("image")
|
||||
prompt = (inp.get("prompt") or "").strip()
|
||||
if not image_b64:
|
||||
return {"error": "Missing 'image' (base64)."}
|
||||
if not prompt:
|
||||
return {"error": "Missing 'prompt' (edit instruction)."}
|
||||
|
||||
try:
|
||||
image = _fit(_b64_to_image(image_b64))
|
||||
# The editor's "Edit strength" slider (0..1) maps to guidance — FLUX.2 klein
|
||||
# likes low guidance (~1); stronger = higher.
|
||||
guidance = 1.0
|
||||
strength = inp.get("strength")
|
||||
if strength is not None:
|
||||
s = max(0.0, min(1.0, float(strength)))
|
||||
guidance = round(1.0 + s * 3.0, 2)
|
||||
seed = inp.get("seed")
|
||||
generator = torch.Generator(device="cpu").manual_seed(int(seed)) if seed is not None else None
|
||||
|
||||
result = _get_pipe()(
|
||||
prompt=prompt,
|
||||
image=image,
|
||||
height=image.height,
|
||||
width=image.width,
|
||||
guidance_scale=guidance,
|
||||
num_inference_steps=STEPS,
|
||||
generator=generator,
|
||||
).images[0]
|
||||
|
||||
return {"image": _image_to_b64(result)}
|
||||
except Exception as exc:
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,12 @@
|
||||
# FLUX.2 [klein] worker. torch is installed (CUDA 12.8 build) FIRST in the
|
||||
# Dockerfile so it carries Blackwell kernels; diffusers is from git for the
|
||||
# Flux2KleinPipeline. FLUX.2 uses a Mistral-based text encoder → transformers +
|
||||
# sentencepiece/protobuf.
|
||||
diffusers @ git+https://github.com/huggingface/diffusers.git
|
||||
transformers>=4.49.0
|
||||
accelerate>=1.2.0
|
||||
safetensors>=0.4.5
|
||||
sentencepiece>=0.2.0
|
||||
protobuf>=4.25.0
|
||||
Pillow==10.4.0
|
||||
runpod==1.7.9
|
||||
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"input": {
|
||||
"image": "PASTE_BASE64_IMAGE_HERE",
|
||||
"prompt": "a vibrant golden-hour sky, dramatic clouds, photorealistic",
|
||||
"strength": 0.6,
|
||||
"num_inference_steps": 25
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg==", "prompt": "turn the sky into a dramatic golden-hour sunset with warm orange clouds"}}
|
||||
@@ -0,0 +1,23 @@
|
||||
FROM python:3.11-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# ffmpeg (with arnndn/afftdn filters) + curl for model download
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends ffmpeg ca-certificates curl && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Download a standard RNNoise model for ffmpeg's arnndn filter.
|
||||
# Source: GregorR/rnnoise-models (beguiling-drafts, a widely used general model).
|
||||
# Non-fatal: if the download fails, handler.py falls back to afftdn.
|
||||
RUN mkdir -p /app/models && \
|
||||
( curl -fsSL -o /app/models/rnnoise.rnnn \
|
||||
https://raw.githubusercontent.com/GregorR/rnnoise-models/master/beguiling-drafts-2018-08-30/bd.rnnn \
|
||||
|| echo "WARNING: RNNoise model download failed; will fall back to afftdn" )
|
||||
|
||||
COPY workers/audio-denoise/requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY workers/audio-denoise/handler.py .
|
||||
|
||||
CMD ["python3", "-u", "handler.py"]
|
||||
@@ -0,0 +1,26 @@
|
||||
# Audio Denoise Worker (RNNoise / arnndn)
|
||||
|
||||
RunPod Serverless CPU worker that removes background noise from audio using
|
||||
ffmpeg's RNNoise filter (`arnndn`), with a built-in FFT denoiser (`afftdn`)
|
||||
as an automatic fallback.
|
||||
|
||||
## Contract
|
||||
- Input: `{"audio": "<base64>"}` — any format/sample rate. A `data:...;base64,`
|
||||
prefix is stripped automatically.
|
||||
- Output: `{"audio": "<base64 WAV>"}` — denoised, 48 kHz mono WAV.
|
||||
- Error: `{"error": "<message>"}`.
|
||||
|
||||
## Pipeline
|
||||
decode base64 -> temp input file -> `ffmpeg -i in -af arnndn=m=/app/models/rnnoise.rnnn -ar 48000 -ac 1 out.wav` -> read -> base64.
|
||||
|
||||
## Deploy
|
||||
- Dockerfile path: `/workers/audio-denoise/Dockerfile`
|
||||
- Build context: repo root (`docker build -f workers/audio-denoise/Dockerfile .`)
|
||||
- GPU: none (CPU-only worker)
|
||||
- App env var to set: `RUNPOD_AUDIO_DENOISE_URL` (your deployed endpoint URL)
|
||||
|
||||
## Model
|
||||
The Dockerfile downloads the `beguiling-drafts` RNNoise model from
|
||||
`GregorR/rnnoise-models` to `/app/models/rnnoise.rnnn`. If that download
|
||||
fails at build time, the handler transparently falls back to ffmpeg's
|
||||
built-in `afftdn` FFT denoiser (no model file needed).
|
||||
@@ -0,0 +1,70 @@
|
||||
_V = "/runpod-volume"
|
||||
import os
|
||||
if os.path.isdir(_V):
|
||||
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
|
||||
|
||||
import base64
|
||||
import subprocess
|
||||
import tempfile
|
||||
|
||||
import runpod
|
||||
|
||||
# Path where the Dockerfile places the RNNoise model file.
|
||||
RNNN_MODEL = "/app/models/rnnoise.rnnn"
|
||||
|
||||
|
||||
def _decode_b64(s):
|
||||
if isinstance(s, str) and "," in s and s.strip().lower().startswith("data:"):
|
||||
s = s.split(",", 1)[1]
|
||||
return base64.b64decode(s)
|
||||
|
||||
|
||||
def _denoise(in_path, out_path):
|
||||
"""Run ffmpeg. Prefer RNNoise (arnndn) when a model file is present,
|
||||
otherwise fall back to the built-in FFT denoiser (afftdn)."""
|
||||
if os.path.isfile(RNNN_MODEL):
|
||||
af = "arnndn=m=" + RNNN_MODEL
|
||||
else:
|
||||
af = "afftdn=nf=-25"
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y", "-hide_banner", "-nostdin",
|
||||
"-i", in_path,
|
||||
"-af", af,
|
||||
"-ar", "48000", "-ac", "1",
|
||||
"-f", "wav",
|
||||
out_path,
|
||||
]
|
||||
proc = subprocess.run(cmd, capture_output=True)
|
||||
if proc.returncode != 0:
|
||||
raise RuntimeError(
|
||||
"ffmpeg failed: " + proc.stderr.decode("utf-8", "replace")[-800:]
|
||||
)
|
||||
|
||||
|
||||
def handler(event):
|
||||
try:
|
||||
inp = (event or {}).get("input") or {}
|
||||
audio_b64 = inp.get("audio")
|
||||
if not audio_b64:
|
||||
return {"error": "Missing 'audio' in input"}
|
||||
|
||||
raw = _decode_b64(audio_b64)
|
||||
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
in_path = os.path.join(td, "in")
|
||||
out_path = os.path.join(td, "out.wav")
|
||||
with open(in_path, "wb") as f:
|
||||
f.write(raw)
|
||||
|
||||
_denoise(in_path, out_path)
|
||||
|
||||
with open(out_path, "rb") as f:
|
||||
out_bytes = f.read()
|
||||
|
||||
return {"audio": base64.b64encode(out_bytes).decode("ascii")}
|
||||
except Exception as exc:
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1 @@
|
||||
runpod==1.7.9
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,21 @@
|
||||
# RunPod Serverless worker: background removal (rembg / U²-Net), CPU-only.
|
||||
# Slim Python base — NO CUDA — so it never hits the driver/kernel issues the GPU
|
||||
# models did. rembg pulls opencv (needs libgl1/libglib). Build context = repo root.
|
||||
FROM python:3.11-slim
|
||||
|
||||
ENV PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
libgl1 libglib2.0-0 && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
RUN pip install --upgrade pip
|
||||
COPY workers/bg-remove/requirements.txt .
|
||||
RUN pip install -r requirements.txt
|
||||
|
||||
COPY workers/bg-remove/handler.py .
|
||||
|
||||
CMD ["python3", "-u", "handler.py"]
|
||||
@@ -0,0 +1,17 @@
|
||||
# Background-removal worker — rembg / U²-Net (CPU)
|
||||
|
||||
Replaces the unreliable in-browser @imgly remover with a RunPod endpoint. Set its
|
||||
`/runsync` URL as **`RUNPOD_BG_REMOVE_URL`**, and turn it on with
|
||||
**`NEXT_PUBLIC_APG_RUNPOD_BG=true`** (without the flag the app keeps using the
|
||||
in-browser fallback).
|
||||
|
||||
## Contract
|
||||
`{"input": {"image": "<base64>", "model_name": "u2net"?}}` → `{"image": "<base64 PNG, transparent background>"}`.
|
||||
|
||||
Models: `u2net` (default, general), `u2netp` (lighter), `u2net_human_seg` (people), `isnet-general-use` (sharper edges).
|
||||
|
||||
## RunPod setup
|
||||
1. New Serverless endpoint → connect this repo → **Dockerfile path** `workers/bg-remove/Dockerfile`.
|
||||
2. Any GPU is fine (it runs on CPU) — pick the cheapest. Max workers ≥ 1.
|
||||
3. Optional volume to cache the ~170 MB model (`U2NET_HOME` → `/runpod-volume/u2net`).
|
||||
4. Copy the `…/runsync` URL → Vercel `RUNPOD_BG_REMOVE_URL`, add `NEXT_PUBLIC_APG_RUNPOD_BG=true`, redeploy.
|
||||
@@ -0,0 +1,66 @@
|
||||
"""
|
||||
RunPod Serverless worker: background removal (U²-Net via rembg).
|
||||
|
||||
Replaces the flaky in-browser @imgly remover. Runs rembg on CPU (U²-Net is tiny,
|
||||
~1-3s) so it needs NO CUDA — it can't hit the driver/kernel problems the GPU
|
||||
models had, and it's very reliable.
|
||||
|
||||
Contract (matches the app's rpRemoveBackground — no app change needed):
|
||||
{"input": {
|
||||
"image": "<base64>", # required (raw base64 or data: URI)
|
||||
"model_name": "u2net" # optional: u2net | u2netp | u2net_human_seg | isnet-general-use
|
||||
}}
|
||||
Returns: {"image": "<base64 PNG, RGBA with the background transparent>"}.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
|
||||
_V = "/runpod-volume"
|
||||
if os.path.isdir(_V):
|
||||
os.environ.setdefault("U2NET_HOME", os.path.join(_V, "u2net"))
|
||||
|
||||
import runpod # noqa: E402
|
||||
from PIL import Image # noqa: E402
|
||||
from rembg import new_session, remove # noqa: E402
|
||||
|
||||
DEFAULT_MODEL = os.environ.get("BG_REMOVE_MODEL", "u2net")
|
||||
_sessions = {}
|
||||
|
||||
|
||||
def _get_session(name):
|
||||
if name not in _sessions:
|
||||
_sessions[name] = new_session(name)
|
||||
return _sessions[name]
|
||||
|
||||
|
||||
def _b64_to_image(data):
|
||||
if "," in data:
|
||||
data = data.split(",", 1)[1]
|
||||
return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGBA")
|
||||
|
||||
|
||||
def _image_to_b64(img):
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
def handler(event):
|
||||
inp = (event or {}).get("input") or {}
|
||||
image_b64 = inp.get("image")
|
||||
if not image_b64:
|
||||
return {"error": "Missing 'image' (base64)."}
|
||||
name = (inp.get("model_name") or DEFAULT_MODEL).strip() or DEFAULT_MODEL
|
||||
try:
|
||||
img = _b64_to_image(image_b64)
|
||||
out = remove(img, session=_get_session(name)) # RGBA, background alpha=0
|
||||
if out.mode != "RGBA":
|
||||
out = out.convert("RGBA")
|
||||
return {"image": _image_to_b64(out)}
|
||||
except Exception as exc:
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,6 @@
|
||||
# Background-removal worker (CPU). rembg pulls onnxruntime (CPU) + opencv itself.
|
||||
rembg==2.0.59
|
||||
onnxruntime==1.17.3
|
||||
numpy<2
|
||||
Pillow==10.4.0
|
||||
runpod==1.7.9
|
||||
@@ -0,0 +1 @@
|
||||
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAQAAAAEACAIAAADTED8xAAACAElEQVR42u3TQQ0AAAjEsJONCEQglTcaaFIFS5bqgbciAQYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABMIAKGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAbAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAGUAEDgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwAFwLd10YnS3DbrYAAAAASUVORK5CYII=", "model_name": "u2net"}}
|
||||
@@ -0,0 +1,24 @@
|
||||
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
PIP_NO_CACHE_DIR=1 \
|
||||
PYTHONUNBUFFERED=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
python3 python3-pip python3-dev \
|
||||
libgl1 libglib2.0-0 \
|
||||
git ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& ln -sf /usr/bin/python3 /usr/bin/python
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
COPY workers/colorize/requirements.txt .
|
||||
RUN pip3 install --upgrade pip && pip3 install -r requirements.txt
|
||||
# Install runpod LAST so modelscope's dependency pins can't break `import runpod`
|
||||
# at worker startup (the cause of "worker never becomes ready, no logs").
|
||||
RUN pip3 install runpod==1.7.9
|
||||
|
||||
COPY workers/colorize/handler.py .
|
||||
|
||||
CMD ["python3", "-u", "handler.py"]
|
||||
@@ -0,0 +1,33 @@
|
||||
# DDColor Colorization Worker (#8)
|
||||
|
||||
RunPod Serverless worker that colorizes black-and-white images using the
|
||||
ModelScope DDColor pipeline (`damo/cv_ddcolor_image-colorization`).
|
||||
|
||||
## Contract
|
||||
|
||||
Input:
|
||||
```json
|
||||
{ "image": "<base64 (optionally data:...;base64, prefixed)>", "input_size": 512 }
|
||||
```
|
||||
|
||||
Output:
|
||||
```json
|
||||
{ "image": "<base64 PNG, same dimensions as input>" }
|
||||
```
|
||||
|
||||
On error: `{ "error": "..." }`.
|
||||
|
||||
## Deploy
|
||||
|
||||
- Dockerfile path: `/workers/colorize/Dockerfile`
|
||||
- Build context: repo root (`docker build -f workers/colorize/Dockerfile .`)
|
||||
- GPU: medium (RTX 3080 / 3090, ~10-12 GB VRAM)
|
||||
- App env var to set: `RUNPOD_COLORIZE_URL` (point your app at this endpoint)
|
||||
|
||||
## Notes
|
||||
|
||||
- Model weights are cached to the network volume via `MODELSCOPE_CACHE`
|
||||
(`/runpod-volume/modelscope`) when `/runpod-volume` is mounted, so the first
|
||||
request downloads and subsequent cold starts reuse the cached weights.
|
||||
- `input_size` controls the DDColor internal working resolution (default 512).
|
||||
The output is always resized back to the exact input dimensions.
|
||||
@@ -0,0 +1,80 @@
|
||||
_V = "/runpod-volume"
|
||||
import os
|
||||
if os.path.isdir(_V):
|
||||
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
|
||||
os.environ.setdefault("MODELSCOPE_CACHE", os.path.join(_V, "modelscope"))
|
||||
|
||||
import base64
|
||||
import runpod
|
||||
|
||||
_pipeline = None
|
||||
|
||||
|
||||
def _get_pipeline():
|
||||
global _pipeline
|
||||
if _pipeline is None:
|
||||
from modelscope.pipelines import pipeline
|
||||
from modelscope.utils.constant import Tasks
|
||||
_pipeline = pipeline(
|
||||
Tasks.image_colorization,
|
||||
model="damo/cv_ddcolor_image-colorization",
|
||||
)
|
||||
return _pipeline
|
||||
|
||||
|
||||
def _decode_image(b64):
|
||||
if "," in b64:
|
||||
b64 = b64.split(",", 1)[1]
|
||||
return base64.b64decode(b64)
|
||||
|
||||
|
||||
def handler(event):
|
||||
try:
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
inp = (event or {}).get("input") or {}
|
||||
img_b64 = inp.get("image")
|
||||
if not img_b64:
|
||||
return {"error": "ValueError: missing 'image' in input"}
|
||||
|
||||
try:
|
||||
input_size = int(inp.get("input_size", 512))
|
||||
except (TypeError, ValueError):
|
||||
input_size = 512
|
||||
|
||||
raw = _decode_image(img_b64)
|
||||
arr = np.frombuffer(raw, dtype=np.uint8)
|
||||
# Decode to BGR (modelscope DDColor expects BGR numpy array).
|
||||
bgr = cv2.imdecode(arr, cv2.IMREAD_COLOR)
|
||||
if bgr is None:
|
||||
return {"error": "ValueError: could not decode input image"}
|
||||
|
||||
orig_h, orig_w = bgr.shape[:2]
|
||||
|
||||
pipe = _get_pipeline()
|
||||
|
||||
from modelscope.outputs import OutputKeys
|
||||
result = pipe(bgr, input_size=input_size)
|
||||
out_bgr = result[OutputKeys.OUTPUT_IMG]
|
||||
out_bgr = np.asarray(out_bgr)
|
||||
|
||||
# Ensure output matches the input dimensions exactly.
|
||||
if out_bgr.shape[0] != orig_h or out_bgr.shape[1] != orig_w:
|
||||
out_bgr = cv2.resize(
|
||||
out_bgr, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4
|
||||
)
|
||||
|
||||
# BGR -> RGB, then encode PNG. cv2.imencode expects BGR, so re-convert.
|
||||
out_rgb = cv2.cvtColor(out_bgr, cv2.COLOR_BGR2RGB)
|
||||
ok, buf = cv2.imencode(".png", cv2.cvtColor(out_rgb, cv2.COLOR_RGB2BGR))
|
||||
if not ok:
|
||||
return {"error": "RuntimeError: PNG encoding failed"}
|
||||
|
||||
out_b64 = base64.b64encode(buf.tobytes()).decode("utf-8")
|
||||
return {"image": out_b64}
|
||||
except Exception as exc:
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,9 @@
|
||||
# NOTE: runpod is installed LAST in the Dockerfile (after modelscope), so
|
||||
# modelscope's aggressive dependency pins can't break `import runpod` at startup.
|
||||
torch==2.1.2
|
||||
torchvision==0.16.2
|
||||
modelscope==1.9.5
|
||||
opencv-python-headless==4.9.0.80
|
||||
pillow==10.2.0
|
||||
numpy==1.24.4
|
||||
timm==0.9.12
|
||||
@@ -0,0 +1 @@
|
||||
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg=="}}
|
||||
@@ -0,0 +1,26 @@
|
||||
# Combined CPU RunPod worker: background removal (rembg) + audio denoise (ffmpeg).
|
||||
# Slim base, NO CUDA. Build context = repo root.
|
||||
FROM python:3.11-slim
|
||||
|
||||
ENV PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
ffmpeg libgl1 libglib2.0-0 curl && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# RNNoise model for ffmpeg's arnndn (falls back to afftdn if the download fails).
|
||||
RUN mkdir -p /app/models && \
|
||||
(curl -fsSL -o /app/models/rnnoise.rnnn \
|
||||
https://raw.githubusercontent.com/GregorR/rnnoise-models/master/beguiling-drafts-2018-08-30/bd.rnnn || \
|
||||
echo "rnnoise model download skipped; afftdn fallback will be used")
|
||||
|
||||
RUN pip install --upgrade pip
|
||||
COPY workers/cpu-tasks/requirements.txt .
|
||||
RUN pip install -r requirements.txt
|
||||
|
||||
COPY workers/cpu-tasks/handler.py .
|
||||
|
||||
CMD ["python3", "-u", "handler.py"]
|
||||
@@ -0,0 +1,82 @@
|
||||
"""
|
||||
Combined CPU RunPod worker (no GPU, no torch). Routes on input.task:
|
||||
"remove-bg" -> rembg (U^2-Net) background removal -> {"image": <png b64>}
|
||||
"denoise" -> ffmpeg RNNoise/afftdn audio denoise -> {"audio": <wav b64>}
|
||||
One request runs exactly one model (isolated). If `task` is absent it's inferred
|
||||
from the payload (audio -> denoise, image -> remove-bg) for backward-compat.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
|
||||
import runpod
|
||||
|
||||
RNNN_MODEL = "/app/models/rnnoise.rnnn"
|
||||
_sessions = {}
|
||||
|
||||
|
||||
def _decode_b64(s):
|
||||
if isinstance(s, str) and "," in s:
|
||||
s = s.split(",", 1)[1]
|
||||
return base64.b64decode(s)
|
||||
|
||||
|
||||
def _remove_bg(inp):
|
||||
from PIL import Image
|
||||
from rembg import new_session, remove
|
||||
|
||||
image_b64 = inp.get("image")
|
||||
if not image_b64:
|
||||
return {"error": "Missing 'image' (base64)."}
|
||||
name = (inp.get("model_name") or os.environ.get("BG_REMOVE_MODEL", "u2net")).strip() or "u2net"
|
||||
if name not in _sessions:
|
||||
_sessions[name] = new_session(name)
|
||||
img = Image.open(io.BytesIO(_decode_b64(image_b64))).convert("RGBA")
|
||||
out = remove(img, session=_sessions[name])
|
||||
if out.mode != "RGBA":
|
||||
out = out.convert("RGBA")
|
||||
buf = io.BytesIO()
|
||||
out.save(buf, format="PNG")
|
||||
return {"image": base64.b64encode(buf.getvalue()).decode("utf-8")}
|
||||
|
||||
|
||||
def _denoise(inp):
|
||||
audio_b64 = inp.get("audio")
|
||||
if not audio_b64:
|
||||
return {"error": "Missing 'audio' (base64)."}
|
||||
raw = _decode_b64(audio_b64)
|
||||
af = ("arnndn=m=" + RNNN_MODEL) if os.path.isfile(RNNN_MODEL) else "afftdn=nf=-25"
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
in_path = os.path.join(td, "in")
|
||||
out_path = os.path.join(td, "out.wav")
|
||||
with open(in_path, "wb") as f:
|
||||
f.write(raw)
|
||||
cmd = ["ffmpeg", "-y", "-hide_banner", "-nostdin", "-i", in_path,
|
||||
"-af", af, "-ar", "48000", "-ac", "1", "-f", "wav", out_path]
|
||||
proc = subprocess.run(cmd, capture_output=True)
|
||||
if proc.returncode != 0:
|
||||
raise RuntimeError("ffmpeg failed: " + proc.stderr.decode("utf-8", "replace")[-600:])
|
||||
with open(out_path, "rb") as f:
|
||||
out_bytes = f.read()
|
||||
return {"audio": base64.b64encode(out_bytes).decode("ascii")}
|
||||
|
||||
|
||||
def handler(event):
|
||||
inp = (event or {}).get("input") or {}
|
||||
task = (inp.get("task") or "").strip().lower()
|
||||
if not task:
|
||||
task = "denoise" if inp.get("audio") else "remove-bg"
|
||||
try:
|
||||
if task == "remove-bg":
|
||||
return _remove_bg(inp)
|
||||
if task == "denoise":
|
||||
return _denoise(inp)
|
||||
return {"error": f"Unknown task '{task}' for cpu endpoint."}
|
||||
except Exception as exc:
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,6 @@
|
||||
# Combined CPU worker: rembg (bg removal) + ffmpeg denoise. No torch/GPU.
|
||||
runpod==1.7.9
|
||||
rembg==2.0.59
|
||||
onnxruntime==1.17.3
|
||||
numpy<2
|
||||
Pillow==10.4.0
|
||||
@@ -0,0 +1,22 @@
|
||||
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
python3 python3-pip python3-dev \
|
||||
libgl1 libglib2.0-0 \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN ln -sf /usr/bin/python3 /usr/bin/python
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
COPY workers/detect/requirements.txt .
|
||||
RUN pip3 install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Bake in the trained construction-material model and point the worker at it.
|
||||
COPY workers/detect/best.pt .
|
||||
ENV YOLO_MODEL=/app/best.pt
|
||||
|
||||
COPY workers/detect/handler.py .
|
||||
|
||||
CMD ["python3", "-u", "handler.py"]
|
||||
@@ -0,0 +1,37 @@
|
||||
# Construction-Material Detection Worker (YOLOv8)
|
||||
|
||||
RunPod Serverless worker running Ultralytics YOLOv8 object detection for a
|
||||
custom-trained construction-material classifier.
|
||||
|
||||
## Build & Deploy
|
||||
- Dockerfile path: `/workers/detect/Dockerfile`
|
||||
- Build context: repo ROOT (the COPY paths are prefixed with `workers/detect/`).
|
||||
- GPU: light/medium — NVIDIA T4 or A4000 is plenty.
|
||||
|
||||
## Model
|
||||
The trained construction-material model (`best.pt`, 43 MB) is **baked into the
|
||||
image** (`COPY workers/detect/best.pt` → `ENV YOLO_MODEL=/app/best.pt`), so the
|
||||
worker detects your materials out of the box — no volume upload needed.
|
||||
|
||||
- `YOLO_MODEL` — override only if you want to swap models (default `/app/best.pt`).
|
||||
To update the model later, replace `workers/detect/best.pt` and push (RunPod rebuilds).
|
||||
|
||||
## App-side env var
|
||||
Point your application at the deployed endpoint with:
|
||||
- `RUNPOD_YOLO_URL` = your RunPod serverless endpoint URL.
|
||||
|
||||
## Contract
|
||||
Input:
|
||||
```json
|
||||
{ "input": { "image": "<base64>" } }
|
||||
```
|
||||
The base64 may include a `data:image/...;base64,` prefix; it is stripped.
|
||||
|
||||
Output (pixel coordinates):
|
||||
```json
|
||||
{ "detections": [
|
||||
{ "name": "brick", "confidence": 0.94,
|
||||
"box": { "x1": 10.0, "y1": 20.0, "x2": 110.0, "y2": 220.0 } }
|
||||
] }
|
||||
```
|
||||
On error: `{ "error": "<Type>: <message>" }`
|
||||
Binary file not shown.
@@ -0,0 +1,64 @@
|
||||
_V = "/runpod-volume"
|
||||
import os
|
||||
if os.path.isdir(_V):
|
||||
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
|
||||
|
||||
import io
|
||||
import base64
|
||||
import runpod
|
||||
from PIL import Image
|
||||
|
||||
_model = None
|
||||
|
||||
|
||||
def _load_model():
|
||||
global _model
|
||||
if _model is None:
|
||||
from ultralytics import YOLO
|
||||
weights = os.environ.get("YOLO_MODEL", "yolov8n.pt")
|
||||
_model = YOLO(weights)
|
||||
return _model
|
||||
|
||||
|
||||
def handler(event):
|
||||
try:
|
||||
data = (event or {}).get("input") or {}
|
||||
image_b64 = data.get("image")
|
||||
if not image_b64:
|
||||
return {"error": "ValueError: missing 'image' in input"}
|
||||
|
||||
if "," in image_b64:
|
||||
image_b64 = image_b64.split(",", 1)[1]
|
||||
image_bytes = base64.b64decode(image_b64)
|
||||
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
||||
|
||||
# Custom construction models are often less confident than COCO — use a low
|
||||
# detection floor (env YOLO_CONF, default 0.15) so real brick/pipe/steel hits
|
||||
# come through instead of being dropped by ultralytics' default 0.25.
|
||||
conf = float(data.get("conf") or os.environ.get("YOLO_CONF", "0.15"))
|
||||
model = _load_model()
|
||||
results = model.predict(image, conf=conf, verbose=False)
|
||||
|
||||
detections = []
|
||||
for result in results:
|
||||
boxes = getattr(result, "boxes", None)
|
||||
if boxes is None:
|
||||
continue
|
||||
names = result.names if getattr(result, "names", None) else model.names
|
||||
for box in boxes:
|
||||
cls = int(box.cls[0])
|
||||
conf = float(box.conf[0])
|
||||
xyxy = box.xyxy[0].tolist()
|
||||
x1, y1, x2, y2 = [float(v) for v in xyxy]
|
||||
detections.append({
|
||||
"name": str(names[cls]),
|
||||
"confidence": conf,
|
||||
"box": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
|
||||
})
|
||||
|
||||
return {"detections": detections}
|
||||
except Exception as exc:
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,3 @@
|
||||
runpod==1.7.9
|
||||
ultralytics==8.3.40
|
||||
pillow==10.4.0
|
||||
@@ -0,0 +1 @@
|
||||
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg=="}}
|
||||
@@ -0,0 +1,25 @@
|
||||
# Combined GPU diffusion worker: FLUX.2 img2img + SDXL inpaint.
|
||||
# Slim Python base + torch cu124 (torch bundles its own CUDA runtime/cuDNN, so we
|
||||
# skip the heavy nvidia/cuda base — that halves the image so it fits the worker's
|
||||
# container disk). The host NVIDIA driver is mounted by RunPod; torch supplies the
|
||||
# rest. No CUDA-version label = runs on any recent driver. Build ctx = repo root.
|
||||
FROM python:3.11-slim
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
git libgomp1 libgl1 libglib2.0-0 && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
RUN pip install --upgrade pip && \
|
||||
pip install torch --index-url https://download.pytorch.org/whl/cu124
|
||||
COPY workers/gpu-diffusion/requirements.txt .
|
||||
RUN pip install -r requirements.txt
|
||||
|
||||
COPY workers/gpu-diffusion/handler.py .
|
||||
|
||||
CMD ["python3", "-u", "handler.py"]
|
||||
@@ -0,0 +1,148 @@
|
||||
"""
|
||||
Combined GPU diffusion RunPod worker. Routes on input.task:
|
||||
"img2img" -> FLUX.2 [klein] (prompt / colorize / sky) -> {"image": b64}
|
||||
"inpaint" -> SDXL inpainting (inpaint/outpaint/eraser/fill) -> {"image": b64}
|
||||
One request runs exactly one model (isolated). Needs a 24 GB GPU. torch cu124 +
|
||||
diffusers. Both pipelines load lazily and are cached.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
|
||||
_V = "/runpod-volume"
|
||||
if os.path.isdir(_V):
|
||||
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
|
||||
|
||||
import torch # noqa: E402
|
||||
import runpod # noqa: E402
|
||||
from PIL import Image # noqa: E402
|
||||
|
||||
FLUX_MODEL = os.environ.get("FLUX2_MODEL", "black-forest-labs/FLUX.2-klein-4B")
|
||||
FLUX_STEPS = int(os.environ.get("FLUX2_STEPS", "6"))
|
||||
INPAINT_MODEL = os.environ.get("SD_INPAINT_MODEL", "diffusers/stable-diffusion-xl-1.0-inpainting-0.1")
|
||||
MAX_SIZE = int(os.environ.get("MAX_SIZE", "1024"))
|
||||
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
_HAS24 = DEVICE == "cuda" and torch.cuda.get_device_properties(0).total_memory / 1e9 >= 20
|
||||
_flux = None
|
||||
_inpaint = None
|
||||
|
||||
|
||||
def _place(p):
|
||||
if _HAS24:
|
||||
return p.to("cuda")
|
||||
p.enable_model_cpu_offload()
|
||||
return p
|
||||
|
||||
|
||||
def _flux_pipe():
|
||||
global _flux
|
||||
if _flux is None:
|
||||
from diffusers import Flux2KleinPipeline
|
||||
|
||||
p = Flux2KleinPipeline.from_pretrained(FLUX_MODEL, torch_dtype=torch.bfloat16)
|
||||
p.set_progress_bar_config(disable=True)
|
||||
_flux = _place(p)
|
||||
return _flux
|
||||
|
||||
|
||||
def _inpaint_pipe():
|
||||
global _inpaint
|
||||
if _inpaint is None:
|
||||
from diffusers import AutoPipelineForInpainting
|
||||
|
||||
try:
|
||||
p = AutoPipelineForInpainting.from_pretrained(
|
||||
INPAINT_MODEL, torch_dtype=torch.float16, variant="fp16"
|
||||
)
|
||||
except Exception:
|
||||
p = AutoPipelineForInpainting.from_pretrained(INPAINT_MODEL, torch_dtype=torch.float16)
|
||||
p.set_progress_bar_config(disable=True)
|
||||
_inpaint = _place(p)
|
||||
return _inpaint
|
||||
|
||||
|
||||
def _b64_img(data, mode="RGB"):
|
||||
if "," in data:
|
||||
data = data.split(",", 1)[1]
|
||||
return Image.open(io.BytesIO(base64.b64decode(data))).convert(mode)
|
||||
|
||||
|
||||
def _to_b64(img):
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
def _fit(img, mult=16):
|
||||
w, h = img.size
|
||||
s = min(1.0, MAX_SIZE / max(w, h))
|
||||
nw = max(mult, (int(w * s) // mult) * mult)
|
||||
nh = max(mult, (int(h * s) // mult) * mult)
|
||||
return img.resize((nw, nh), Image.LANCZOS)
|
||||
|
||||
|
||||
def _run_flux(inp):
|
||||
image_b64 = inp.get("image")
|
||||
prompt = (inp.get("prompt") or "").strip()
|
||||
if not image_b64:
|
||||
return {"error": "Missing 'image'."}
|
||||
if not prompt:
|
||||
return {"error": "Missing 'prompt'."}
|
||||
image = _fit(_b64_img(image_b64), 16)
|
||||
guidance = 1.0
|
||||
st = inp.get("strength")
|
||||
if st is not None:
|
||||
s = max(0.0, min(1.0, float(st)))
|
||||
guidance = round(1.0 + s * 3.0, 2)
|
||||
seed = inp.get("seed")
|
||||
gen = torch.Generator(device="cpu").manual_seed(int(seed)) if seed is not None else None
|
||||
res = _flux_pipe()(
|
||||
prompt=prompt, image=image, height=image.height, width=image.width,
|
||||
guidance_scale=guidance, num_inference_steps=FLUX_STEPS, generator=gen,
|
||||
).images[0]
|
||||
return {"image": _to_b64(res)}
|
||||
|
||||
|
||||
def _run_inpaint(inp):
|
||||
image_b64 = inp.get("image")
|
||||
mask_b64 = inp.get("mask")
|
||||
if not image_b64:
|
||||
return {"error": "Missing 'image'."}
|
||||
if not mask_b64:
|
||||
return {"error": "Missing 'mask'."}
|
||||
prompt = (inp.get("prompt") or "").strip()
|
||||
negative = (inp.get("negative_prompt") or "").strip() or None
|
||||
image = _fit(_b64_img(image_b64), 8)
|
||||
mask = _b64_img(mask_b64, "L").resize(image.size, Image.NEAREST)
|
||||
steps = int(inp.get("num_inference_steps") or 35)
|
||||
guidance = float(inp.get("guidance_scale") or 7.0)
|
||||
raw = inp.get("strength")
|
||||
strength = 1.0 if not prompt else (max(0.7, min(1.0, float(raw))) if raw is not None else 0.9)
|
||||
seed = inp.get("seed")
|
||||
gen = torch.Generator(device=DEVICE).manual_seed(int(seed)) if seed is not None else None
|
||||
eff = prompt or "clean, seamless, photorealistic background, natural continuation"
|
||||
res = _inpaint_pipe()(
|
||||
prompt=eff, negative_prompt=negative, image=image, mask_image=mask,
|
||||
height=image.height, width=image.width, num_inference_steps=steps,
|
||||
guidance_scale=guidance, strength=strength, generator=gen,
|
||||
).images[0]
|
||||
return {"image": _to_b64(res)}
|
||||
|
||||
|
||||
def handler(event):
|
||||
inp = (event or {}).get("input") or {}
|
||||
task = (inp.get("task") or "").strip().lower()
|
||||
if not task:
|
||||
task = "inpaint" if inp.get("mask") else "img2img"
|
||||
try:
|
||||
if task == "img2img":
|
||||
return _run_flux(inp)
|
||||
if task == "inpaint":
|
||||
return _run_inpaint(inp)
|
||||
return {"error": f"Unknown task '{task}' for diffusion endpoint."}
|
||||
except Exception as exc:
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,9 @@
|
||||
# FLUX.2 (Flux2KleinPipeline, from git) + SDXL inpaint (AutoPipelineForInpainting).
|
||||
diffusers @ git+https://github.com/huggingface/diffusers.git
|
||||
transformers>=4.49.0
|
||||
accelerate>=1.2.0
|
||||
safetensors>=0.4.5
|
||||
sentencepiece>=0.2.0
|
||||
protobuf>=4.25.0
|
||||
Pillow==10.4.0
|
||||
runpod==1.7.9
|
||||
@@ -0,0 +1,30 @@
|
||||
# Combined GPU worker: YOLO detect + Real-ESRGAN upscale + Whisper transcribe.
|
||||
# CUDA 12.1 + torch cu121 (runs on RTX 3090/4090 driver-550 hosts). Build ctx = repo root.
|
||||
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
python3 python3-pip libgl1 libglib2.0-0 && \
|
||||
ln -sf /usr/bin/python3 /usr/bin/python && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
RUN pip install --upgrade pip && \
|
||||
pip install torch==2.1.2 torchvision==0.16.2 --index-url https://download.pytorch.org/whl/cu121
|
||||
COPY workers/gpu-vision/requirements.txt .
|
||||
RUN pip install -r requirements.txt
|
||||
|
||||
# Fix basicsr's removed torchvision import (the #1 Real-ESRGAN crash) — rewrite it.
|
||||
RUN F=$(python3 -c "import basicsr, os; print(os.path.join(os.path.dirname(basicsr.__file__), 'data', 'degradations.py'))") && \
|
||||
sed -i 's/from torchvision.transforms.functional_tensor import rgb_to_grayscale/from torchvision.transforms.functional import rgb_to_grayscale/' "$F"
|
||||
|
||||
# YOLO construction model + handler.
|
||||
COPY workers/gpu-vision/best.pt .
|
||||
ENV YOLO_MODEL=/app/best.pt
|
||||
COPY workers/gpu-vision/handler.py .
|
||||
|
||||
CMD ["python3", "-u", "handler.py"]
|
||||
Binary file not shown.
@@ -0,0 +1,191 @@
|
||||
"""
|
||||
Combined GPU worker: YOLO detect + Real-ESRGAN upscale + Whisper transcribe.
|
||||
Routes on input.task:
|
||||
"detect" -> YOLO (best.pt construction model) -> {"detections": [...]}
|
||||
"upscale" -> Real-ESRGAN (+ optional GFPGAN face) -> {"image": b64}
|
||||
"transcribe" -> Whisper (faster-whisper) -> {"transcript": ..., ...}
|
||||
One request runs exactly one model (isolated). torch cu121 + ultralytics +
|
||||
realesrgan + faster-whisper. Each model loads lazily and is cached.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
import traceback
|
||||
import urllib.request
|
||||
|
||||
_V = "/runpod-volume"
|
||||
if os.path.isdir(_V):
|
||||
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
|
||||
|
||||
import runpod # noqa: E402
|
||||
from PIL import Image # noqa: E402
|
||||
|
||||
_WEIGHTS_DIR = os.path.join(_V, "weights") if os.path.isdir(_V) else "/app/weights"
|
||||
os.makedirs(_WEIGHTS_DIR, exist_ok=True)
|
||||
|
||||
_yolo = None
|
||||
_ups = None
|
||||
_face = None
|
||||
_whisper = None
|
||||
|
||||
_RRDB = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
|
||||
_GFP = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
|
||||
|
||||
|
||||
# ---------- YOLO ----------
|
||||
def _yolo_model():
|
||||
global _yolo
|
||||
if _yolo is None:
|
||||
from ultralytics import YOLO
|
||||
|
||||
_yolo = YOLO(os.environ.get("YOLO_MODEL", "/app/best.pt"))
|
||||
return _yolo
|
||||
|
||||
|
||||
def _detect(inp):
|
||||
b64 = inp.get("image")
|
||||
if not b64:
|
||||
return {"error": "Missing 'image'."}
|
||||
if "," in b64:
|
||||
b64 = b64.split(",", 1)[1]
|
||||
img = Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
|
||||
conf = float(inp.get("conf") or os.environ.get("YOLO_CONF", "0.15"))
|
||||
res = _yolo_model().predict(img, conf=conf, verbose=False)
|
||||
dets = []
|
||||
for r in res:
|
||||
names = r.names
|
||||
for b in (r.boxes or []):
|
||||
cls = int(b.cls[0])
|
||||
x1, y1, x2, y2 = [float(v) for v in b.xyxy[0].tolist()]
|
||||
dets.append({
|
||||
"name": str(names[cls]),
|
||||
"confidence": float(b.conf[0]),
|
||||
"box": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
|
||||
})
|
||||
return {"detections": dets}
|
||||
|
||||
|
||||
# ---------- Real-ESRGAN upscale ----------
|
||||
def _dl(url, dst):
|
||||
if not os.path.isfile(dst):
|
||||
tmp = dst + ".tmp"
|
||||
urllib.request.urlretrieve(url, tmp)
|
||||
os.replace(tmp, dst)
|
||||
return dst
|
||||
|
||||
|
||||
def _upsampler():
|
||||
global _ups
|
||||
if _ups is None:
|
||||
import torch
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||
from realesrgan import RealESRGANer
|
||||
|
||||
mp = _dl(_RRDB, os.path.join(_WEIGHTS_DIR, "RealESRGAN_x4plus.pth"))
|
||||
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
|
||||
_ups = RealESRGANer(
|
||||
scale=4, model_path=mp, model=model, tile=0, tile_pad=10, pre_pad=0,
|
||||
half=torch.cuda.is_available(), gpu_id=None,
|
||||
)
|
||||
return _ups
|
||||
|
||||
|
||||
def _face_enh():
|
||||
global _face
|
||||
if _face is None:
|
||||
from gfpgan import GFPGANer
|
||||
|
||||
gp = _dl(_GFP, os.path.join(_WEIGHTS_DIR, "GFPGANv1.4.pth"))
|
||||
_face = GFPGANer(model_path=gp, upscale=4, arch="clean", channel_multiplier=2, bg_upsampler=_upsampler())
|
||||
return _face
|
||||
|
||||
|
||||
def _upscale(inp):
|
||||
import numpy as np
|
||||
|
||||
b64 = inp.get("image")
|
||||
if not b64:
|
||||
return {"error": "Missing 'image'."}
|
||||
if "," in b64:
|
||||
b64 = b64.split(",", 1)[1]
|
||||
scale = int(inp.get("scale", 4))
|
||||
if scale not in (2, 4):
|
||||
scale = 4
|
||||
face = bool(inp.get("face_enhance", False))
|
||||
pil = Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
|
||||
rgb = np.array(pil)
|
||||
bgr = rgb[:, :, ::-1].copy()
|
||||
if face:
|
||||
_, _, ob = _face_enh().enhance(bgr, has_aligned=False, only_center_face=False, paste_back=True)
|
||||
out = Image.fromarray(ob[:, :, ::-1]).resize((rgb.shape[1] * scale, rgb.shape[0] * scale), Image.LANCZOS)
|
||||
else:
|
||||
ob, _ = _upsampler().enhance(bgr, outscale=scale)
|
||||
out = Image.fromarray(ob[:, :, ::-1])
|
||||
buf = io.BytesIO()
|
||||
out.save(buf, format="PNG")
|
||||
return {"image": base64.b64encode(buf.getvalue()).decode("utf-8")}
|
||||
|
||||
|
||||
# ---------- Whisper transcribe ----------
|
||||
def _whisper_model():
|
||||
global _whisper
|
||||
if _whisper is None:
|
||||
from faster_whisper import WhisperModel
|
||||
|
||||
_whisper = WhisperModel(
|
||||
os.environ.get("WHISPER_MODEL", "medium"),
|
||||
device=os.environ.get("WHISPER_DEVICE", "cuda"),
|
||||
compute_type=os.environ.get("WHISPER_COMPUTE", "float16"),
|
||||
download_root=os.environ.get("HF_HOME"),
|
||||
)
|
||||
return _whisper
|
||||
|
||||
|
||||
def _transcribe(inp):
|
||||
b64 = inp.get("audio")
|
||||
if not b64:
|
||||
return {"error": "Missing 'audio'."}
|
||||
lang = inp.get("language") or None
|
||||
want_ts = bool(inp.get("timestamps", False))
|
||||
data = b64.split(",", 1)[1] if "," in b64 else b64
|
||||
path = None
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile(suffix=".audio", delete=False) as f:
|
||||
f.write(base64.b64decode(data))
|
||||
path = f.name
|
||||
segs, info = _whisper_model().transcribe(path, language=lang, vad_filter=True)
|
||||
parts, out_segs = [], []
|
||||
for s in segs:
|
||||
parts.append(s.text)
|
||||
if want_ts:
|
||||
out_segs.append({"text": s.text.strip(), "start_sec": round(s.start, 2), "end_sec": round(s.end, 2)})
|
||||
out = {"transcript": "".join(parts).strip(), "language": info.language}
|
||||
if want_ts:
|
||||
out["segments"] = out_segs
|
||||
return out
|
||||
finally:
|
||||
if path and os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
def handler(event):
|
||||
inp = (event or {}).get("input") or {}
|
||||
task = (inp.get("task") or "").strip().lower()
|
||||
if not task:
|
||||
task = "transcribe" if inp.get("audio") else ("upscale" if inp.get("scale") else "detect")
|
||||
try:
|
||||
if task == "detect":
|
||||
return _detect(inp)
|
||||
if task == "upscale":
|
||||
return _upscale(inp)
|
||||
if task == "transcribe":
|
||||
return _transcribe(inp)
|
||||
return {"error": f"Unknown task '{task}' for vision endpoint."}
|
||||
except Exception as exc:
|
||||
traceback.print_exc()
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,10 @@
|
||||
# YOLO (ultralytics) + Real-ESRGAN (realesrgan/basicsr/gfpgan) + Whisper (faster-whisper).
|
||||
# torch/torchvision are installed in the Dockerfile first (cu121).
|
||||
runpod==1.7.9
|
||||
ultralytics==8.3.40
|
||||
realesrgan==0.3.0
|
||||
basicsr==1.4.2
|
||||
gfpgan==1.3.8
|
||||
faster-whisper==1.0.3
|
||||
numpy==1.24.4
|
||||
Pillow==10.4.0
|
||||
@@ -0,0 +1,26 @@
|
||||
# RunPod Serverless worker: SDXL inpainting (Magic Eraser / Generative Fill / Outpaint).
|
||||
# CUDA 12.1 base + torch cu121 — the same base the detect/upscale workers use, which
|
||||
# runs on RunPod's RTX 3090/4090 (driver 550) hosts (12.1 <= 12.4, backward-compatible).
|
||||
# Build context = repo root, so COPY paths are workers/inpaint/*.
|
||||
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
python3 python3-pip git && \
|
||||
ln -sf /usr/bin/python3 /usr/bin/python && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# torch (CUDA 12.1 build) FIRST — carries sm_86 (3090) + sm_89 (4090) kernels.
|
||||
RUN pip install --upgrade pip && \
|
||||
pip install torch --index-url https://download.pytorch.org/whl/cu121
|
||||
COPY workers/inpaint/requirements.txt .
|
||||
RUN pip install -r requirements.txt
|
||||
|
||||
COPY workers/inpaint/handler.py .
|
||||
|
||||
CMD ["python3", "-u", "handler.py"]
|
||||
@@ -0,0 +1,30 @@
|
||||
# Inpaint worker — SDXL (`diffusers/stable-diffusion-xl-1.0-inpainting-0.1`)
|
||||
|
||||
One RunPod serverless endpoint that powers **Magic Eraser**, **Generative Fill**,
|
||||
and **Outpaint** in the photo editor. Set its `/runsync` URL as the Vercel env var
|
||||
**`RUNPOD_SD_INPAINT_URL`**.
|
||||
|
||||
## Contract
|
||||
Request `{"input": {...}}` — matches the app's `rpInpaint`:
|
||||
|
||||
| field | required | meaning |
|
||||
|---|---|---|
|
||||
| `image` | yes | base64 image (raw or `data:` URI) |
|
||||
| `mask` | yes | base64 mask — **white = regenerate, black = keep** |
|
||||
| `prompt` | no | `""` = Magic Eraser (clean fill); text = Generative Fill / Outpaint |
|
||||
| `strength` | no | 0..1 (empty prompt forces 1.0 to fully erase) |
|
||||
| `guidance_scale` | no | default 7 |
|
||||
| `num_inference_steps` | no | default 35 |
|
||||
| `negative_prompt` | no | — |
|
||||
| `seed` | no | — |
|
||||
|
||||
Response: `{"image": "<base64 PNG>"}` or `{"error": "..."}`.
|
||||
|
||||
## RunPod setup
|
||||
1. New Serverless endpoint → connect this GitHub repo → **Dockerfile path** `workers/inpaint/Dockerfile` (build context = repo root).
|
||||
2. GPU: **24 GB** (RTX 4090 / 3090). Max workers ≥ 1; optionally a network volume to cache the ~7 GB model so cold starts don't re-download.
|
||||
3. Copy the endpoint's `…/runsync` URL → Vercel `RUNPOD_SD_INPAINT_URL`, then redeploy the web app.
|
||||
|
||||
## Env knobs
|
||||
- `SD_INPAINT_MODEL` (default `diffusers/stable-diffusion-xl-1.0-inpainting-0.1`)
|
||||
- `SD_INPAINT_MAX_SIZE` (default `1024`)
|
||||
@@ -0,0 +1,137 @@
|
||||
"""
|
||||
RunPod Serverless worker: SDXL inpainting.
|
||||
|
||||
Powers the app's Magic Eraser, Generative Fill AND Outpaint — all three go
|
||||
through this one endpoint (RUNPOD_SD_INPAINT_URL). No app change needed; this
|
||||
matches the exact contract the app already sends (rpInpaint):
|
||||
|
||||
{"input": {
|
||||
"image": "<base64>", # required (raw base64 or data: URI)
|
||||
"mask": "<base64>", # required — WHITE = regenerate, BLACK = keep
|
||||
"prompt": "<text>", # optional ("" = Magic Eraser: clean fill)
|
||||
"strength": 0.8, # optional
|
||||
"guidance_scale": 7, # optional
|
||||
"num_inference_steps": 35, # optional
|
||||
"negative_prompt": "...", # optional
|
||||
"seed": 123 # optional
|
||||
}}
|
||||
Returns: {"image": "<base64 PNG>"} (or {"error": "..."}).
|
||||
|
||||
CUDA 12.1 base + torch cu121 (same as the detect/upscale workers) so it runs on
|
||||
RunPod's RTX 3090/4090 (driver 550) hosts.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
|
||||
_V = "/runpod-volume"
|
||||
if os.path.isdir(_V):
|
||||
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
|
||||
|
||||
import torch # noqa: E402
|
||||
import runpod # noqa: E402
|
||||
from PIL import Image # noqa: E402
|
||||
from diffusers import AutoPipelineForInpainting # noqa: E402
|
||||
|
||||
MODEL = os.environ.get("SD_INPAINT_MODEL", "diffusers/stable-diffusion-xl-1.0-inpainting-0.1")
|
||||
MAX_SIZE = int(os.environ.get("SD_INPAINT_MAX_SIZE", "1024"))
|
||||
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
_pipe = None
|
||||
|
||||
|
||||
def _get_pipe():
|
||||
global _pipe
|
||||
if _pipe is None:
|
||||
try:
|
||||
pipe = AutoPipelineForInpainting.from_pretrained(
|
||||
MODEL, torch_dtype=torch.float16, variant="fp16"
|
||||
)
|
||||
except Exception:
|
||||
# Repo may not ship an fp16 variant — fall back to default weights.
|
||||
pipe = AutoPipelineForInpainting.from_pretrained(MODEL, torch_dtype=torch.float16)
|
||||
total_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 if DEVICE == "cuda" else 0
|
||||
if total_gb >= 20:
|
||||
pipe = pipe.to("cuda")
|
||||
else:
|
||||
pipe.enable_model_cpu_offload()
|
||||
pipe.set_progress_bar_config(disable=True)
|
||||
_pipe = pipe
|
||||
return _pipe
|
||||
|
||||
|
||||
def _b64_to_image(data, mode="RGB"):
|
||||
if "," in data:
|
||||
data = data.split(",", 1)[1]
|
||||
return Image.open(io.BytesIO(base64.b64decode(data))).convert(mode)
|
||||
|
||||
|
||||
def _image_to_b64(img):
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
def _fit(img):
|
||||
# SDXL wants multiples of 8; cap the long side at MAX_SIZE.
|
||||
w, h = img.size
|
||||
scale = min(1.0, MAX_SIZE / max(w, h))
|
||||
nw = max(8, (int(w * scale) // 8) * 8)
|
||||
nh = max(8, (int(h * scale) // 8) * 8)
|
||||
return img.resize((nw, nh), Image.LANCZOS)
|
||||
|
||||
|
||||
def handler(event):
|
||||
inp = (event or {}).get("input") or {}
|
||||
image_b64 = inp.get("image")
|
||||
mask_b64 = inp.get("mask")
|
||||
if not image_b64:
|
||||
return {"error": "Missing 'image' (base64)."}
|
||||
if not mask_b64:
|
||||
return {"error": "Missing 'mask' (base64)."}
|
||||
|
||||
prompt = (inp.get("prompt") or "").strip()
|
||||
negative = (inp.get("negative_prompt") or "").strip() or None
|
||||
|
||||
try:
|
||||
image = _fit(_b64_to_image(image_b64, "RGB"))
|
||||
mask = _b64_to_image(mask_b64, "L").resize(image.size, Image.NEAREST)
|
||||
|
||||
steps = int(inp.get("num_inference_steps") or 35)
|
||||
guidance = float(inp.get("guidance_scale") or 7.0)
|
||||
|
||||
# Strength: empty prompt = Magic Eraser -> fully replace (1.0). With a
|
||||
# prompt (Generative Fill / Outpaint) respect the caller but keep a floor
|
||||
# so the masked region actually changes.
|
||||
raw = inp.get("strength")
|
||||
if not prompt:
|
||||
strength = 1.0
|
||||
else:
|
||||
strength = max(0.7, min(1.0, float(raw))) if raw is not None else 0.9
|
||||
|
||||
seed = inp.get("seed")
|
||||
generator = (
|
||||
torch.Generator(device=DEVICE).manual_seed(int(seed)) if seed is not None else None
|
||||
)
|
||||
|
||||
eff_prompt = prompt or "clean, seamless, photorealistic background, natural continuation"
|
||||
|
||||
result = _get_pipe()(
|
||||
prompt=eff_prompt,
|
||||
negative_prompt=negative,
|
||||
image=image,
|
||||
mask_image=mask,
|
||||
height=image.height,
|
||||
width=image.width,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=guidance,
|
||||
strength=strength,
|
||||
generator=generator,
|
||||
).images[0]
|
||||
|
||||
return {"image": _image_to_b64(result)}
|
||||
except Exception as exc:
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,8 @@
|
||||
# SDXL inpainting worker. torch is installed (cu121) in the Dockerfile first.
|
||||
# AutoPipelineForInpainting + SDXL uses two text encoders → transformers.
|
||||
diffusers>=0.30.0
|
||||
transformers>=4.44.0
|
||||
accelerate>=0.34.0
|
||||
safetensors>=0.4.5
|
||||
Pillow==10.4.0
|
||||
runpod==1.7.9
|
||||
@@ -0,0 +1 @@
|
||||
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFlklEQVR42u3VMQEAAAzCMKQjHQ97l0jo0xSAlyIBgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAUgAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAADcDrctaAb6XeXAAAAAASUVORK5CYII=", "mask": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAAAAADRE4smAAADXUlEQVR42u3SQREAAAzCMPybhvc0LJXQS6rXxQIABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAgAACwAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAAsAEAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAAALAAAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAAwAIABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAgAACwAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIBuA8sSO8WbTgaBAAAAAElFTkSuQmCC", "prompt": "a clean concrete wall", "num_inference_steps": 20}}
|
||||
@@ -0,0 +1,33 @@
|
||||
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
python3 python3-pip git ca-certificates \
|
||||
libgl1 libglib2.0-0 && \
|
||||
rm -rf /var/lib/apt/lists/* && \
|
||||
ln -sf /usr/bin/python3 /usr/bin/python
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Clone the research repo (inference code + configs). Pinning is recommended:
|
||||
# replace HEAD with a specific commit once verified.
|
||||
RUN git clone --depth 1 https://github.com/AlanSavio25/DeepSingleImageCalibration.git \
|
||||
/opt/DeepSingleImageCalibration
|
||||
|
||||
# Install the CUDA 12.1 torch stack + repo deps.
|
||||
COPY workers/tilt/requirements.txt .
|
||||
RUN pip install --index-url https://download.pytorch.org/whl/cu121 \
|
||||
torch==2.1.2 torchvision==0.16.2 && \
|
||||
pip install -r requirements.txt
|
||||
|
||||
# Best-effort install of the repo's own requirements if present (non-fatal).
|
||||
RUN if [ -f /opt/DeepSingleImageCalibration/requirements.txt ]; then \
|
||||
pip install -r /opt/DeepSingleImageCalibration/requirements.txt || true ; \
|
||||
fi
|
||||
|
||||
COPY workers/tilt/handler.py .
|
||||
|
||||
CMD ["python3", "-u", "handler.py"]
|
||||
@@ -0,0 +1,35 @@
|
||||
# workers/tilt — Single-Image Camera Calibration (roll / pitch / FOV)
|
||||
|
||||
Predicts camera **roll**, **pitch**, and **vertical field-of-view** from one image, based on
|
||||
[AlanSavio25/DeepSingleImageCalibration](https://github.com/AlanSavio25/DeepSingleImageCalibration).
|
||||
|
||||
- **Dockerfile:** `/workers/tilt/Dockerfile` (build context = repo root)
|
||||
- **GPU:** light — NVIDIA **T4** (16 GB) is plenty; CPU also works, slower.
|
||||
- **App env var to set on the endpoint:** `RUNPOD_TILT_URL` — the deployed RunPod endpoint URL the app calls.
|
||||
|
||||
## Contract
|
||||
Input:
|
||||
```json
|
||||
{ "input": { "image": "<base64>" } }
|
||||
```
|
||||
`image` may include a `data:image/...;base64,` prefix (stripped automatically).
|
||||
|
||||
Output:
|
||||
```json
|
||||
{ "roll_degrees": 0.0, "pitch_degrees": 0.0, "fov_degrees": 0.0 }
|
||||
```
|
||||
On error: `{ "error": "..." }`.
|
||||
|
||||
## Weights (required — research repo, not on pip)
|
||||
The pretrained checkpoint is **not** bundled. Download the repo's released weights and place
|
||||
them on the network volume, then point the worker at them:
|
||||
|
||||
- Default path: `/runpod-volume/tilt/weights.tar`
|
||||
- Override with env var `WEIGHTS_PATH`.
|
||||
|
||||
The backbone loads into the volume-cached `TORCH_HOME`/`HF_HOME` on first request.
|
||||
|
||||
## Tunable env vars (no rebuild needed)
|
||||
`WEIGHTS_PATH`, `CALIB_BACKBONE` (densenet161|densenet121), `CALIB_NUM_BINS`,
|
||||
`CALIB_INPUT_SIZE`, `CALIB_ROLL_MIN/MAX`, `CALIB_VFOV_MIN/MAX`, `CALIB_RHO_MIN/MAX`,
|
||||
`CALIB_DIRECT_PITCH` (=1 if the net outputs pitch directly), `CALIB_PITCH_MIN/MAX`.
|
||||
@@ -0,0 +1,206 @@
|
||||
# --- network volume cache setup (must run before importing heavy libs) ---
|
||||
_V = "/runpod-volume"
|
||||
import os
|
||||
if os.path.isdir(_V):
|
||||
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
|
||||
os.environ.setdefault("TORCH_HOME", os.path.join(_V, "torch"))
|
||||
|
||||
import io
|
||||
import sys
|
||||
import base64
|
||||
import math
|
||||
|
||||
import runpod
|
||||
|
||||
# The DeepSingleImageCalibration repo is cloned into /opt/DeepSingleImageCalibration
|
||||
# by the Dockerfile. Make it importable.
|
||||
_REPO = os.environ.get("CALIB_REPO", "/opt/DeepSingleImageCalibration")
|
||||
if _REPO not in sys.path:
|
||||
sys.path.insert(0, _REPO)
|
||||
|
||||
_MODEL = None # cached (model, device)
|
||||
_TF = None # cached torchvision transform
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Bin / range configuration.
|
||||
#
|
||||
# NOTE (RESEARCH REPO -- VERIFY THESE): AlanSavio25/DeepSingleImageCalibration
|
||||
# predicts each parameter as a classification over discretised bins, then takes
|
||||
# the soft-expected value over the bin centres. The exact number of bins, the
|
||||
# min/max of each range, the head ORDER/NAMES, and whether the network emits
|
||||
# `pitch` directly or a horizon-offset `rho` all come from the repo's training
|
||||
# config + the single-image-prediction notebook. The values below follow the
|
||||
# "A Perceptual Measure for Deep Single Image Camera Calibration" formulation
|
||||
# and are the most likely defaults, but MUST be confirmed against the checkpoint.
|
||||
# Override any of them via env vars without rebuilding.
|
||||
# ---------------------------------------------------------------------------
|
||||
NUM_BINS = int(os.environ.get("CALIB_NUM_BINS", "256"))
|
||||
ROLL_MIN = float(os.environ.get("CALIB_ROLL_MIN", "-45.0"))
|
||||
ROLL_MAX = float(os.environ.get("CALIB_ROLL_MAX", "45.0"))
|
||||
VFOV_MIN = float(os.environ.get("CALIB_VFOV_MIN", "20.0"))
|
||||
VFOV_MAX = float(os.environ.get("CALIB_VFOV_MAX", "105.0"))
|
||||
# rho = signed vertical horizon offset as a fraction of image height.
|
||||
RHO_MIN = float(os.environ.get("CALIB_RHO_MIN", "-1.5"))
|
||||
RHO_MAX = float(os.environ.get("CALIB_RHO_MAX", "1.5"))
|
||||
# If the network already emits pitch (degrees) instead of rho, set this to "1".
|
||||
CALIB_DIRECT_PITCH = os.environ.get("CALIB_DIRECT_PITCH", "0") == "1"
|
||||
PITCH_MIN = float(os.environ.get("CALIB_PITCH_MIN", "-45.0"))
|
||||
PITCH_MAX = float(os.environ.get("CALIB_PITCH_MAX", "45.0"))
|
||||
|
||||
INPUT_SIZE = int(os.environ.get("CALIB_INPUT_SIZE", "320"))
|
||||
BACKBONE = os.environ.get("CALIB_BACKBONE", "densenet161") # repo default backbone
|
||||
|
||||
# Where the pretrained checkpoint lives. Download the repo's released weights
|
||||
# onto the network volume and point this at them.
|
||||
WEIGHTS_PATH = os.environ.get(
|
||||
"WEIGHTS_PATH",
|
||||
os.path.join(_V, "tilt", "weights.tar") if os.path.isdir(_V) else "/opt/weights.tar",
|
||||
)
|
||||
|
||||
|
||||
def _build_model():
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
# DenseNet backbone with one classification head per predicted parameter.
|
||||
# This mirrors the repo architecture (ImageNet-pretrained backbone whose
|
||||
# classifier is replaced by per-parameter heads over NUM_BINS bins).
|
||||
if BACKBONE == "densenet161":
|
||||
base = torchvision.models.densenet161(weights=None)
|
||||
feat_dim = 2208
|
||||
elif BACKBONE == "densenet121":
|
||||
base = torchvision.models.densenet121(weights=None)
|
||||
feat_dim = 1024
|
||||
else:
|
||||
raise RuntimeError(f"Unsupported CALIB_BACKBONE={BACKBONE!r}")
|
||||
|
||||
class CalibNet(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.features = base.features
|
||||
self.head_roll = nn.Linear(feat_dim, NUM_BINS)
|
||||
self.head_pitch = nn.Linear(feat_dim, NUM_BINS) # 'rho' or pitch head
|
||||
self.head_vfov = nn.Linear(feat_dim, NUM_BINS)
|
||||
self.head_k1 = nn.Linear(feat_dim, NUM_BINS) # distortion (unused in output)
|
||||
|
||||
def forward(self, x):
|
||||
f = self.features(x)
|
||||
f = nn.functional.relu(f, inplace=True)
|
||||
f = nn.functional.adaptive_avg_pool2d(f, (1, 1)).flatten(1)
|
||||
return {
|
||||
"roll": self.head_roll(f),
|
||||
"pitch": self.head_pitch(f),
|
||||
"vfov": self.head_vfov(f),
|
||||
"k1": self.head_k1(f),
|
||||
}
|
||||
|
||||
model = CalibNet()
|
||||
|
||||
# Load pretrained weights. Research checkpoints commonly store the state dict
|
||||
# under a "model" or "state_dict" key; fall back to the raw object.
|
||||
if not os.path.isfile(WEIGHTS_PATH):
|
||||
raise RuntimeError(
|
||||
f"checkpoint not found at {WEIGHTS_PATH!r}; place the repo's pretrained "
|
||||
f"weights on the network volume or set WEIGHTS_PATH"
|
||||
)
|
||||
ckpt = torch.load(WEIGHTS_PATH, map_location="cpu")
|
||||
if isinstance(ckpt, dict):
|
||||
state = ckpt.get("model", ckpt.get("state_dict", ckpt))
|
||||
else:
|
||||
state = ckpt
|
||||
if isinstance(state, dict):
|
||||
# strip common prefixes
|
||||
cleaned = {}
|
||||
for k, v in state.items():
|
||||
nk = k
|
||||
for pref in ("module.", "model.", "net."):
|
||||
if nk.startswith(pref):
|
||||
nk = nk[len(pref):]
|
||||
cleaned[nk] = v
|
||||
# strict=False: head naming in the checkpoint may differ from ours; the
|
||||
# backbone (bulk of the weights) will still load. See notes.
|
||||
model.load_state_dict(cleaned, strict=False)
|
||||
|
||||
model.eval().to(device)
|
||||
return model, device
|
||||
|
||||
|
||||
def _get_transform():
|
||||
global _TF
|
||||
if _TF is None:
|
||||
import torchvision.transforms as T
|
||||
_TF = T.Compose([
|
||||
T.Resize((INPUT_SIZE, INPUT_SIZE)),
|
||||
T.ToTensor(),
|
||||
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
])
|
||||
return _TF
|
||||
|
||||
|
||||
def _load_model():
|
||||
global _MODEL
|
||||
if _MODEL is None:
|
||||
_MODEL = _build_model()
|
||||
return _MODEL
|
||||
|
||||
|
||||
def _expected_value(logits, vmin, vmax):
|
||||
"""Soft-argmax: softmax over bins, dot with evenly-spaced bin centres."""
|
||||
import torch
|
||||
probs = torch.softmax(logits, dim=-1)
|
||||
n = probs.shape[-1]
|
||||
centers = torch.linspace(vmin, vmax, n, device=probs.device, dtype=probs.dtype)
|
||||
return float((probs * centers).sum(dim=-1).item())
|
||||
|
||||
|
||||
def _decode_image(b64):
|
||||
from PIL import Image
|
||||
if "," in b64:
|
||||
b64 = b64.split(",", 1)[1]
|
||||
raw = base64.b64decode(b64)
|
||||
return Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
|
||||
|
||||
def handler(event):
|
||||
try:
|
||||
data = (event or {}).get("input") or {}
|
||||
b64 = data.get("image")
|
||||
if not b64:
|
||||
return {"error": "missing required 'image' (base64) field"}
|
||||
|
||||
import torch
|
||||
|
||||
img = _decode_image(b64)
|
||||
model, device = _load_model()
|
||||
tf = _get_transform()
|
||||
x = tf(img).unsqueeze(0).to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
out = model(x)
|
||||
|
||||
roll = _expected_value(out["roll"], ROLL_MIN, ROLL_MAX)
|
||||
vfov = _expected_value(out["vfov"], VFOV_MIN, VFOV_MAX)
|
||||
|
||||
if CALIB_DIRECT_PITCH:
|
||||
pitch = _expected_value(out["pitch"], PITCH_MIN, PITCH_MAX)
|
||||
else:
|
||||
# Head predicts rho = signed vertical horizon offset as a fraction of
|
||||
# image height. Convert to pitch via the pinhole relation:
|
||||
# tan(pitch) = 2 * rho * tan(vfov/2)
|
||||
rho = _expected_value(out["pitch"], RHO_MIN, RHO_MAX)
|
||||
vfov_rad = math.radians(vfov)
|
||||
pitch = math.degrees(math.atan(2.0 * rho * math.tan(vfov_rad / 2.0)))
|
||||
|
||||
return {
|
||||
"roll_degrees": float(roll),
|
||||
"pitch_degrees": float(pitch),
|
||||
"fov_degrees": float(vfov),
|
||||
}
|
||||
except Exception as exc:
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,9 @@
|
||||
runpod==1.7.9
|
||||
numpy==1.26.4
|
||||
Pillow==10.4.0
|
||||
opencv-python-headless==4.10.0.84
|
||||
omegaconf==2.3.0
|
||||
tqdm==4.66.5
|
||||
matplotlib==3.8.4
|
||||
h5py==3.11.0
|
||||
scipy==1.13.1
|
||||
@@ -0,0 +1 @@
|
||||
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg=="}}
|
||||
@@ -0,0 +1,34 @@
|
||||
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV PIP_NO_CACHE_DIR=1
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
python3 python3-pip python3-dev \
|
||||
libgl1 libglib2.0-0 wget ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& ln -sf /usr/bin/python3 /usr/bin/python
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Newer CUDA-12.1 torch so the compiled kernels cover current GPUs (fixes
|
||||
# "no kernel image is available for execution on the device"). The basicsr
|
||||
# functional_tensor import (removed in torchvision>=0.16) is handled by the
|
||||
# sed-patch below instead of pinning an old torchvision.
|
||||
RUN pip install --upgrade pip && \
|
||||
pip install torch==2.1.2 torchvision==0.16.2 \
|
||||
--index-url https://download.pytorch.org/whl/cu121
|
||||
|
||||
COPY workers/upscale/requirements.txt .
|
||||
RUN pip install -r requirements.txt
|
||||
|
||||
# Robustly fix basicsr's removed torchvision import (the #1 Real-ESRGAN crash).
|
||||
# Works whether or not the torch/torchvision pin held — rewrites the bad import.
|
||||
RUN F=$(python3 -c 'import basicsr,os;print(os.path.join(os.path.dirname(basicsr.__file__),"data","degradations.py"))') && \
|
||||
sed -i 's/from torchvision.transforms.functional_tensor import rgb_to_grayscale/from torchvision.transforms.functional import rgb_to_grayscale/' "$F" && \
|
||||
echo "patched basicsr degradations.py"
|
||||
|
||||
COPY workers/upscale/handler.py .
|
||||
|
||||
CMD ["python3","-u","handler.py"]
|
||||
@@ -0,0 +1,35 @@
|
||||
# Real-ESRGAN Upscale Worker (#7)
|
||||
|
||||
RunPod Serverless worker for Real-ESRGAN image super-resolution with optional GFPGAN face enhancement.
|
||||
|
||||
- **Dockerfile path:** `/workers/upscale/Dockerfile`
|
||||
- **Build context:** repository ROOT (COPY paths are prefixed with `workers/upscale/`).
|
||||
- **GPU:** medium (RTX 3090 / RTX 4090, ~24 GB VRAM).
|
||||
- **App env var to set:** `RUNPOD_UPSCALE_URL` (point your app at this endpoint's RunPod URL).
|
||||
|
||||
## Contract
|
||||
|
||||
Input:
|
||||
```json
|
||||
{"image": "<base64 PNG/JPEG, optional data: prefix>", "scale": 2 or 4 (default 4), "face_enhance": false}
|
||||
```
|
||||
Output (success):
|
||||
```json
|
||||
{"image": "<base64 PNG>"}
|
||||
```
|
||||
Output (error):
|
||||
```json
|
||||
{"error": "SomeError: message"}
|
||||
```
|
||||
|
||||
## Model weights
|
||||
|
||||
Downloaded on first request and cached to the network volume when mounted:
|
||||
- `RealESRGAN_x4plus.pth` (RRDBNet x4) -> `/runpod-volume/weights/`
|
||||
- `GFPGANv1.4.pth` (only when `face_enhance=true`) -> `/runpod-volume/weights/`
|
||||
|
||||
Attach a network volume mounted at `/runpod-volume` so weights (and `HF_HOME`) persist between cold starts. Without a volume, weights download to `/app/weights` per container.
|
||||
|
||||
## Notes
|
||||
|
||||
The base model is x4; `scale=2` is produced via RealESRGANer's `outscale` parameter. For `face_enhance`, GFPGAN upscales at its fixed factor and the result is resized to the requested scale.
|
||||
@@ -0,0 +1,113 @@
|
||||
_V = "/runpod-volume"
|
||||
import os
|
||||
if os.path.isdir(_V):
|
||||
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
|
||||
|
||||
import io
|
||||
import base64
|
||||
import traceback
|
||||
import urllib.request
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
_WEIGHTS_DIR = os.path.join(_V, "weights") if os.path.isdir(_V) else "/app/weights"
|
||||
os.makedirs(_WEIGHTS_DIR, exist_ok=True)
|
||||
|
||||
_RRDB_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
|
||||
_GFPGAN_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
|
||||
|
||||
_UPSAMPLER = None
|
||||
_FACE_ENHANCER = None
|
||||
|
||||
|
||||
def _download(url, dst):
|
||||
if not os.path.isfile(dst):
|
||||
tmp = dst + ".tmp"
|
||||
urllib.request.urlretrieve(url, tmp)
|
||||
os.replace(tmp, dst)
|
||||
return dst
|
||||
|
||||
|
||||
def _get_upsampler():
|
||||
global _UPSAMPLER
|
||||
if _UPSAMPLER is None:
|
||||
import torch
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||
from realesrgan import RealESRGANer
|
||||
|
||||
model_path = _download(_RRDB_URL, os.path.join(_WEIGHTS_DIR, "RealESRGAN_x4plus.pth"))
|
||||
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=model_path,
|
||||
model=model,
|
||||
tile=0,
|
||||
tile_pad=10,
|
||||
pre_pad=0,
|
||||
half=torch.cuda.is_available(),
|
||||
gpu_id=None,
|
||||
)
|
||||
return _UPSAMPLER
|
||||
|
||||
|
||||
def _get_face_enhancer():
|
||||
global _FACE_ENHANCER
|
||||
if _FACE_ENHANCER is None:
|
||||
from gfpgan import GFPGANer
|
||||
|
||||
gfpgan_path = _download(_GFPGAN_URL, os.path.join(_WEIGHTS_DIR, "GFPGANv1.4.pth"))
|
||||
_FACE_ENHANCER = GFPGANer(
|
||||
model_path=gfpgan_path,
|
||||
upscale=4,
|
||||
arch="clean",
|
||||
channel_multiplier=2,
|
||||
bg_upsampler=_get_upsampler(),
|
||||
)
|
||||
return _FACE_ENHANCER
|
||||
|
||||
|
||||
def handler(event):
|
||||
try:
|
||||
data = (event or {}).get("input") or {}
|
||||
b64 = data.get("image")
|
||||
if not b64:
|
||||
return {"error": "ValueError: 'image' is required"}
|
||||
if "," in b64:
|
||||
b64 = b64.split(",", 1)[1]
|
||||
|
||||
scale = int(data.get("scale", 4))
|
||||
if scale not in (2, 4):
|
||||
scale = 4
|
||||
face_enhance = bool(data.get("face_enhance", False))
|
||||
|
||||
raw = base64.b64decode(b64)
|
||||
pil = Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
rgb = np.array(pil)
|
||||
bgr = rgb[:, :, ::-1].copy() # RealESRGANer expects BGR (cv2 convention)
|
||||
|
||||
if face_enhance:
|
||||
enhancer = _get_face_enhancer()
|
||||
_, _, out_bgr = enhancer.enhance(
|
||||
bgr, has_aligned=False, only_center_face=False, paste_back=True
|
||||
)
|
||||
# GFPGANer upscales by its fixed factor; resize to the requested scale.
|
||||
target = (rgb.shape[1] * scale, rgb.shape[0] * scale)
|
||||
out_rgb = out_bgr[:, :, ::-1]
|
||||
out_pil = Image.fromarray(out_rgb).resize(target, Image.LANCZOS)
|
||||
else:
|
||||
upsampler = _get_upsampler()
|
||||
out_bgr, _ = upsampler.enhance(bgr, outscale=scale)
|
||||
out_rgb = out_bgr[:, :, ::-1]
|
||||
out_pil = Image.fromarray(out_rgb)
|
||||
|
||||
buf = io.BytesIO()
|
||||
out_pil.save(buf, format="PNG")
|
||||
return {"image": base64.b64encode(buf.getvalue()).decode("utf-8")}
|
||||
except Exception as exc:
|
||||
traceback.print_exc()
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
|
||||
import runpod
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,6 @@
|
||||
runpod==1.7.9
|
||||
realesrgan==0.3.0
|
||||
basicsr==1.4.2
|
||||
gfpgan==1.3.8
|
||||
numpy==1.24.4
|
||||
pillow==10.4.0
|
||||
@@ -0,0 +1 @@
|
||||
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg==", "scale": 4}}
|
||||
@@ -0,0 +1,20 @@
|
||||
# RunPod Serverless worker: Speech-to-Text (faster-whisper).
|
||||
# Lightweight — no torch/NeMo; ctranslate2 uses the CUDA + cuDNN in the base image.
|
||||
# Paths are relative to the REPO ROOT (RunPod's build context).
|
||||
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
python3 python3-pip ffmpeg && \
|
||||
ln -sf /usr/bin/python3 /usr/bin/python && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
COPY workers/voice-to-text/requirements.txt .
|
||||
RUN pip3 install --no-cache-dir -r requirements.txt
|
||||
COPY workers/voice-to-text/handler.py .
|
||||
|
||||
CMD ["python3", "-u", "handler.py"]
|
||||
@@ -0,0 +1,48 @@
|
||||
# Speech-to-Text worker — Whisper (RunPod Serverless)
|
||||
|
||||
Multilingual voice-to-text via **faster-whisper** — lets a worker **speak**
|
||||
instead of typing annotation text, in both the image and video editors. 99
|
||||
languages incl. Hindi/regional + English. Chosen over Parakeet/NeMo because it
|
||||
has no heavy dependencies (no NeMo, no torch), so it builds cleanly.
|
||||
|
||||
## Contract
|
||||
|
||||
Request → `POST /runsync`:
|
||||
```json
|
||||
{ "input": { "audio": "<base64 audio>", "language": "en", "timestamps": false } }
|
||||
```
|
||||
- `audio`: base64 wav/mp3/m4a (any sample rate — resampled internally)
|
||||
- `language`: optional (e.g. `"hi"`, `"en"`); omit to auto-detect
|
||||
- `timestamps`: optional — per-segment start/end
|
||||
|
||||
Response → `{ "output": { "transcript": "...", "language": "en" } }`
|
||||
(with `segments: [{text, start_sec, end_sec}]` if `timestamps: true`), or `{ "output": { "error": "..." } }`.
|
||||
|
||||
## Deploy (new endpoint — one per model)
|
||||
|
||||
1. **RunPod → Serverless → New Endpoint → Deploy from a GitHub repository** → `Sumit-Pluto/photo_gallery`.
|
||||
2. **Dockerfile path:** `/workers/voice-to-text/Dockerfile`
|
||||
3. Config:
|
||||
- **GPU:** light — **16 GB (T4 / A4000)** is plenty (`large-v3` ≈ 5 GB VRAM).
|
||||
- **Network volume:** attach your weights volume (caches the model).
|
||||
- **Workers:** Max 1, Active 0, FlashBoot on. **Execution timeout** ~300s (first-run download).
|
||||
4. **Warm up once** (downloads the model, ~1.5 GB):
|
||||
```bash
|
||||
curl -X POST https://api.runpod.ai/v2/<NEW_ID>/runsync \
|
||||
-H "Content-Type: application/json" -H "Authorization: Bearer <API_KEY>" \
|
||||
-d '{"input":{"audio":"<base64-wav>","language":"en"}}'
|
||||
```
|
||||
5. In **Vercel**, set `RUNPOD_STT_URL` = `https://api.runpod.ai/v2/<NEW_ID>/runsync` (server-only).
|
||||
|
||||
## Tuning (env vars)
|
||||
|
||||
| Var | Default | Purpose |
|
||||
|---|---|---|
|
||||
| `WHISPER_MODEL` | `large-v3` | accuracy vs speed — `medium` / `small` are faster/lighter |
|
||||
| `WHISPER_COMPUTE` | `float16` | `int8_float16` uses less VRAM |
|
||||
|
||||
## App integration (next step, not yet built)
|
||||
|
||||
Mic-record button in the editor's text/annotation tool → `/api/ai/transcribe`
|
||||
(proxies here) → inserts the transcript. Client `rpTranscribe()` already exists in
|
||||
`apps/web/src/lib/runpod/endpoints.ts`.
|
||||
@@ -0,0 +1,91 @@
|
||||
"""
|
||||
RunPod Serverless worker: Speech-to-Text (Whisper via faster-whisper).
|
||||
|
||||
Chosen over Parakeet/NeMo because it has NO heavy dependencies (no NeMo, no
|
||||
torch — just ctranslate2), so it builds cleanly, AND it's multilingual (99
|
||||
languages incl. Hindi/regional + English) — a better fit for a mixed crew.
|
||||
|
||||
Lets a worker speak instead of typing annotation text (image + video editors).
|
||||
|
||||
Matches the app's transcribe contract:
|
||||
{"input": {
|
||||
"audio": "<base64 audio>", # required (wav/mp3/m4a; any sample rate — resampled internally)
|
||||
"language": "en", # optional (e.g. "hi", "en"); omit = auto-detect
|
||||
"timestamps": false # optional — return per-segment start/end
|
||||
}}
|
||||
Returns: {"transcript": "...", "language": "en", "segments"?: [{text,start_sec,end_sec}]}
|
||||
(or {"error": "..."}).
|
||||
|
||||
Tuning (env): WHISPER_MODEL (default large-v3; "medium"/"small" are faster),
|
||||
WHISPER_COMPUTE (default float16; "int8_float16" for less VRAM).
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
_VOLUME = "/runpod-volume"
|
||||
if os.path.isdir(_VOLUME):
|
||||
os.environ.setdefault("HF_HOME", os.path.join(_VOLUME, "huggingface"))
|
||||
|
||||
import runpod # noqa: E402
|
||||
from faster_whisper import WhisperModel # noqa: E402
|
||||
|
||||
MODEL_SIZE = os.environ.get("WHISPER_MODEL", "large-v3")
|
||||
DEVICE = os.environ.get("WHISPER_DEVICE", "cuda")
|
||||
COMPUTE = os.environ.get("WHISPER_COMPUTE", "float16")
|
||||
|
||||
_model = None
|
||||
|
||||
|
||||
def _get_model():
|
||||
global _model
|
||||
if _model is None:
|
||||
_model = WhisperModel(
|
||||
MODEL_SIZE,
|
||||
device=DEVICE,
|
||||
compute_type=COMPUTE,
|
||||
download_root=os.environ.get("HF_HOME"),
|
||||
)
|
||||
return _model
|
||||
|
||||
|
||||
def handler(event):
|
||||
inp = (event or {}).get("input") or {}
|
||||
audio_b64 = inp.get("audio")
|
||||
if not audio_b64:
|
||||
return {"error": "Missing 'audio' (base64)."}
|
||||
|
||||
language = inp.get("language") or None
|
||||
want_ts = bool(inp.get("timestamps", False))
|
||||
|
||||
path = None
|
||||
try:
|
||||
data = audio_b64.split(",", 1)[1] if "," in audio_b64 else audio_b64
|
||||
with tempfile.NamedTemporaryFile(suffix=".audio", delete=False) as f:
|
||||
f.write(base64.b64decode(data))
|
||||
path = f.name
|
||||
|
||||
segments, info = _get_model().transcribe(path, language=language, vad_filter=True)
|
||||
|
||||
parts = []
|
||||
segs = []
|
||||
for s in segments:
|
||||
parts.append(s.text)
|
||||
if want_ts:
|
||||
segs.append(
|
||||
{"text": s.text.strip(), "start_sec": round(s.start, 2), "end_sec": round(s.end, 2)}
|
||||
)
|
||||
|
||||
out = {"transcript": "".join(parts).strip(), "language": info.language}
|
||||
if want_ts:
|
||||
out["segments"] = segs
|
||||
return out
|
||||
except Exception as exc:
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
finally:
|
||||
if path and os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,2 @@
|
||||
runpod==1.7.9
|
||||
faster-whisper==1.0.3
|
||||
File diff suppressed because one or more lines are too long
Reference in New Issue
Block a user