87 lines
3.1 KiB
TypeScript
87 lines
3.1 KiB
TypeScript
/**
|
|
* 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.
|
|
*
|
|
* PROVENANCE: a faithful port of
|
|
* `advance-photo-gallery-web-sdk/apps/web/src/lib/runpod/base64.ts`.
|
|
* The `pickOutputImage` normalization is intentionally identical so both apps
|
|
* tolerate the same set of endpoint response shapes.
|
|
*
|
|
* SERVER-ONLY.
|
|
*/
|
|
|
|
/** "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)
|
|
*/
|
|
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;
|
|
}
|