make input_images universal

Author: ProgrammerIn-wonderland

src/backend/drivers/ai-image/inputImage.ts

@@ -0,0 +1,102 @@
+/*
+ * Copyright (C) 2024-present Puter Technologies Inc.
+ *
+ * This file is part of Puter.
+ *
+ * Puter is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+/**
+ * Shared helpers for `input_images` (image-to-image) handling across image
+ * providers. `input_images` is the canonical, cross-provider field; an entry
+ * may be a public URL, a data-URI, or raw base64. Providers whose upstream
+ * API needs base64 use these helpers to normalize URLs server-side (via the
+ * SSRF-guarded `secureFetch`); providers that accept URLs natively (Replicate,
+ * xAI) pass them through untouched.
+ */
+
+import { HttpError } from '../../core/http/HttpError.js';
+import { secureFetch } from '../../util/secureHttp.js';
+import type { IGenerateParams } from './types.js';
+
+export function isHttpUrl(s: string): boolean {
+    return s.startsWith('http://') || s.startsWith('https://');
+}
+
+/**
+ * Resolve the single input image for providers that only support one.
+ * Throws 400 if `input_images` carries more than one entry. Returns the
+ * chosen image string (URL / data-URI / raw base64) or undefined.
+ */
+export function resolveSingleInputImage(
+    params: Pick<IGenerateParams, 'input_image' | 'input_images'>,
+    providerLabel: string,
+): string | undefined {
+    const imgs = params.input_images;
+    if (imgs && imgs.length > 1) {
+        throw new HttpError(
+            400,
+            `${providerLabel} supports only a single input image; pass one image via input_image or a single-element input_images.`,
+            { legacyCode: 'bad_request' },
+        );
+    }
+    return params.input_image ?? imgs?.[0];
+}
+
+const DATA_URI_PATTERN = /^data:([^;,]+)?(?:;base64)?,(.*)$/s;
+
+/** Parse a `data:<mime>;base64,<payload>` URI into raw base64 + mime. */
+export function parseDataUri(
+    s: string,
+): { base64: string; mime: string } | null {
+    const m = DATA_URI_PATTERN.exec(s);
+    if (!m) return null;
+    return { base64: m[2] ?? '', mime: m[1] ?? 'image/png' };
+}
+
+/** Fetch an http(s) image and return raw base64 + mime (SSRF-guarded). */
+export async function fetchImageAsBase64(
+    url: string,
+): Promise<{ base64: string; mime: string }> {
+    const res = await secureFetch(url);
+    if (!res.ok) {
+        throw new HttpError(
+            400,
+            `Failed to fetch input image (status ${res.status})`,
+            { legacyCode: 'bad_request' },
+        );
+    }
+    const buffer = Buffer.from(await res.arrayBuffer());
+    const mime =
+        res.headers.get('content-type')?.split(';')[0]?.trim() || 'image/png';
+    return { base64: buffer.toString('base64'), mime };
+}
+
+/**
+ * Normalize any input-image string to a base64 data-URI:
+ *   • http(s) URL  → fetched via secureFetch
+ *   • data-URI     → returned as-is
+ *   • raw base64   → wrapped with `mimeHint` (default image/png)
+ */
+export async function toBase64DataUri(
+    img: string,
+    mimeHint?: string,
+): Promise<string> {
+    if (img.startsWith('data:')) return img;
+    if (isHttpUrl(img)) {
+        const { base64, mime } = await fetchImageAsBase64(img);
+        return `data:${mime};base64,${base64}`;
+    }
+    return `data:${mimeHint ?? 'image/png'};base64,${img}`;
+}

src/backend/drivers/ai-image/providers/cloudflare/CloudflareImageProvider.test.ts

@@ -47,6 +47,16 @@ import { withTestActor } from '../../../integrationTestUtil.js';
 import { CLOUDFLARE_IMAGE_GENERATION_MODELS } from './models.js';
 import { CloudflareImageProvider } from './CloudflareImageProvider.js';
 
+// Stub the URL→base64 fetch so URL inputs stay offline; keep the rest real.
+const { fetchImageAsBase64Mock } = vi.hoisted(() => ({
+    fetchImageAsBase64Mock: vi.fn(),
+}));
+
+vi.mock('../../inputImage.js', async (orig) => ({
+    ...(await orig<typeof import('../../inputImage.js')>()),
+    fetchImageAsBase64: fetchImageAsBase64Mock,
+}));
+
 // ── Test harness ────────────────────────────────────────────────────
 
 let server: PuterServer;
@@ -81,6 +91,7 @@ const makeProvider = (
     );
 
 beforeEach(() => {
+    fetchImageAsBase64Mock.mockReset();
     fetchSpy = vi.spyOn(globalThis, 'fetch') as MockInstance<typeof fetch>;
     hasCreditsSpy = vi.spyOn(server.services.metering, 'hasEnoughCredits');
     batchIncrementUsagesSpy = vi.spyOn(
@@ -470,3 +481,70 @@ describe('CloudflareImageProvider.generate cost components', () => {
         );
     });
 });
+
+// ── input_images (canonical image-to-image field) ──────────────────
+
+describe('CloudflareImageProvider.generate input_images', () => {
+    const klein9bWith = (extra: Record<string, unknown>) => {
+        const provider = makeProvider();
+        fetchSpy.mockResolvedValueOnce(
+            new Response(Buffer.from([1, 2, 3]).buffer, {
+                status: 200,
+                headers: { 'content-type': 'image/png' },
+            }),
+        );
+        return withTestActor(() =>
+            provider.generate({
+                model: '@cf/black-forest-labs/flux-2-klein-9b',
+                prompt: 'edit it',
+                ratio: { w: 2000, h: 1000 },
+                ...extra,
+            } as never),
+        );
+    };
+
+    const hasInputCostLine = () => {
+        const [, entries] = batchIncrementUsagesSpy.mock.calls[0]!;
+        return (entries as Array<{ usageType: string }>).some((e) =>
+            e.usageType.endsWith(':input_image_mp'),
+        );
+    };
+
+    it('maps a base64/data-URI input_images entry to the input image (cost line appears)', async () => {
+        await klein9bWith({ input_images: ['data:image/png;base64,AAAA'] });
+        expect(hasInputCostLine()).toBe(true);
+    });
+
+    it('maps a singular input_image to the input image', async () => {
+        await klein9bWith({ input_image: 'data:image/png;base64,AAAA' });
+        expect(hasInputCostLine()).toBe(true);
+    });
+
+    it('fetches an http(s) URL input via secureFetch and uses it as the input image', async () => {
+        fetchImageAsBase64Mock.mockResolvedValueOnce({
+            base64: 'AAAA',
+            mime: 'image/png',
+        });
+        await klein9bWith({ input_images: ['https://example.com/in.png'] });
+        expect(fetchImageAsBase64Mock).toHaveBeenCalledWith(
+            'https://example.com/in.png',
+        );
+        expect(hasInputCostLine()).toBe(true);
+    });
+
+    it('throws 400 when more than one input image is supplied (before any fetch)', async () => {
+        const provider = makeProvider();
+        await expect(
+            withTestActor(() =>
+                provider.generate({
+                    model: '@cf/black-forest-labs/flux-2-klein-9b',
+                    prompt: 'edit it',
+                    ratio: { w: 1024, h: 1024 },
+                    input_images: ['data:image/png;base64,AAAA', 'data:image/png;base64,BBBB'],
+                } as never),
+            ),
+        ).rejects.toMatchObject({ statusCode: 400 });
+        expect(fetchSpy).not.toHaveBeenCalled();
+        expect(fetchImageAsBase64Mock).not.toHaveBeenCalled();
+    });
+});

src/backend/drivers/ai-image/providers/cloudflare/CloudflareImageProvider.ts

@@ -29,6 +29,11 @@ import {
     CLOUDFLARE_IMAGE_GENERATION_MODELS,
     CloudflareImageModel,
 } from './models.js';
+import {
+    fetchImageAsBase64,
+    isHttpUrl,
+    resolveSingleInputImage,
+} from '../../inputImage.js';
 
 type CloudflareGenerateParams = IGenerateParams & {
     steps?: number;
@@ -101,6 +106,16 @@ export class CloudflareImageProvider implements IImageProvider {
             });
         }
 
+        // Canonical `input_images`/`input_image` → Cloudflare's `image` field.
+        // Cloudflare accepts a single input image; a URL is fetched to base64
+        // server-side (SSRF-guarded) since the API has no URL field.
+        const singleInput = resolveSingleInputImage(options, 'Cloudflare');
+        if (singleInput) {
+            options.image ??= isHttpUrl(singleInput)
+                ? (await fetchImageAsBase64(singleInput)).base64
+                : singleInput;
+        }
+
         const steps = this.#resolveSteps(selectedModel, options);
         const costComponents = this.#estimateCost(selectedModel, ratio, steps, {
             hasInputImage:

src/backend/drivers/ai-image/providers/gemini/GeminiImageProvider.test.ts

@@ -71,6 +71,16 @@ vi.mock('@google/genai', () => {
     return { GoogleGenAI };
 });
 
+// Stub the URL→data-URI normalizer so URL inputs stay offline; keep the rest real.
+const { toBase64DataUriMock } = vi.hoisted(() => ({
+    toBase64DataUriMock: vi.fn(),
+}));
+
+vi.mock('../../inputImage.js', async (orig) => ({
+    ...(await orig<typeof import('../../inputImage.js')>()),
+    toBase64DataUri: toBase64DataUriMock,
+}));
+
 // ── Test harness ────────────────────────────────────────────────────
 
 let server: PuterServer;
@@ -94,6 +104,7 @@ const makeProvider = () =>
 beforeEach(() => {
     generateContentMock.mockReset();
     generateImagesMock.mockReset();
+    toBase64DataUriMock.mockReset();
     googleAICtor.mockReset();
     hasCreditsSpy = vi.spyOn(server.services.metering, 'hasEnoughCredits');
     incrementUsageSpy = vi.spyOn(server.services.metering, 'incrementUsage');
@@ -272,6 +283,31 @@ describe('GeminiImageProvider.generate Flash path (generateContent)', () => {
         expect(result).toBe('data:image/png;base64,BASE64IMG');
     });
 
+    it('fetches an http(s) URL input and sends it as an inlineData part', async () => {
+        const provider = makeProvider();
+        generateContentMock.mockResolvedValueOnce(inlineImageResponse);
+        toBase64DataUriMock.mockResolvedValueOnce(
+            'data:image/png;base64,URLBYTES',
+        );
+
+        await withTestActor(() =>
+            provider.generate({
+                model: 'gemini-2.5-flash-image',
+                prompt: 'add a hat',
+                input_images: ['https://example.com/in.png'],
+            }),
+        );
+
+        expect(toBase64DataUriMock).toHaveBeenCalledWith(
+            'https://example.com/in.png',
+            undefined,
+        );
+        const sent = generateContentMock.mock.calls[0]![0];
+        expect(sent.contents).toContainEqual({
+            inlineData: { mimeType: 'image/png', data: 'URLBYTES' },
+        });
+    });
+
     it('throws 400 when the SDK returns no inline image data', async () => {
         const provider = makeProvider();
         generateContentMock.mockResolvedValueOnce({

src/backend/drivers/ai-image/providers/gemini/GeminiImageProvider.ts

@@ -31,6 +31,7 @@ import type {
     IImageModel,
     IImageProvider,
 } from '../../types.js';
+import { isHttpUrl, toBase64DataUri } from '../../inputImage.js';
 import { HttpError } from '@heyputer/backend/src/core/http/HttpError.js';
 
 const MIME_SIGNATURES: Record<string, string> = {
@@ -107,6 +108,18 @@ export class GeminiImageProvider implements IImageProvider {
             input_images = [input_image];
         }
 
+        // Resolve any http(s) URL inputs to base64 data-URIs server-side
+        // (SSRF-guarded) so the rest of the flow only deals with inline data.
+        if (input_images?.length) {
+            input_images = await Promise.all(
+                input_images.map((img) =>
+                    isHttpUrl(img)
+                        ? toBase64DataUri(img, input_image_mime_type)
+                        : img,
+                ),
+            );
+        }
+
         // Validate input images have detectable MIME types
         if (input_images?.length) {
             for (const img of input_images) {

src/backend/drivers/ai-image/providers/openai/OpenAiImageProvider.test.ts

@@ -78,6 +78,16 @@ vi.mock('openai', () => {
     };
 });
 
+// Stub the URL→base64 fetch so URL inputs stay offline; keep the rest real.
+const { fetchImageAsBase64Mock } = vi.hoisted(() => ({
+    fetchImageAsBase64Mock: vi.fn(),
+}));
+
+vi.mock('../../inputImage.js', async (orig) => ({
+    ...(await orig<typeof import('../../inputImage.js')>()),
+    fetchImageAsBase64: fetchImageAsBase64Mock,
+}));
+
 // ── Test harness ────────────────────────────────────────────────────
 
 let server: PuterServer;
@@ -100,6 +110,7 @@ const makeProvider = () =>
 beforeEach(() => {
     generateMock.mockReset();
     editMock.mockReset();
+    fetchImageAsBase64Mock.mockReset();
     openAICtor.mockReset();
     hasCreditsSpy = vi.spyOn(server.services.metering, 'hasEnoughCredits');
     batchIncrementUsagesSpy = vi.spyOn(
@@ -326,6 +337,32 @@ describe('OpenAiImageProvider.generate input_images (edit endpoint)', () => {
         expect(Array.isArray(sent.image)).toBe(false);
         expect((sent.image as { __file?: boolean }).__file).toBe(true);
     });
+
+    it('fetches an http(s) URL input and sends the bytes to images.edit', async () => {
+        const provider = makeProvider();
+        editMock.mockResolvedValueOnce(editResponse);
+        fetchImageAsBase64Mock.mockResolvedValueOnce({
+            base64: 'iVBORw0KGgo=',
+            mime: 'image/png',
+        });
+
+        await withTestActor(() =>
+            provider.generate({
+                model: 'gpt-image-1',
+                prompt: 'add a hat',
+                ratio: { w: 1024, h: 1024 },
+                input_images: ['https://example.com/in.png'],
+            }),
+        );
+
+        expect(fetchImageAsBase64Mock).toHaveBeenCalledWith(
+            'https://example.com/in.png',
+        );
+        expect(generateMock).not.toHaveBeenCalled();
+        expect(editMock).toHaveBeenCalledTimes(1);
+        const sent = editMock.mock.calls[0]![0];
+        expect((sent.image as { __file?: boolean }).__file).toBe(true);
+    });
 });
 
 // ── gpt-image (token-priced) request shape & metering ──────────────