Skip to content

Commit d13952a

Browse files
jddunnjddunn
committed
feat(api): add upscaleImage and variateImage for image post-processing
Add two more provider-agnostic image API functions: - upscaleImage() — super-resolution with scale factor (2x/4x) or target dimensions. Supported by Stability (conservative upscale), A1111 (extra-single-image), and Replicate (real-esrgan). - variateImage() — generate visual variations of a source image. Native support via OpenAI (variations endpoint), with automatic img2img fallback for Stability, A1111, and Replicate using variance-to-strength mapping. Add 10 tests covering 2x/4x upscale, width/height targets, unsupported provider errors, Replicate real-esrgan, OpenAI variations, Stability img2img fallback, local SD img2img fallback, and default variance.
1 parent f5bc775 commit d13952a

4 files changed

Lines changed: 808 additions & 0 deletions

File tree

‎src/api/__tests__/upscaleImage.test.ts‎

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
/**
2+
* @file upscaleImage.test.ts
3+
* Tests for the high-level upscaleImage API covering 2x and 4x upscaling,
4+
* target dimensions, and unsupported-provider errors.
5+
*
6+
* All tests mock `globalThis.fetch` — no real API calls are made.
7+
*/
8+
import { afterEach, describe, expect, it, vi } from 'vitest';
9+
10+
import { upscaleImage } from '../upscaleImage.js';
11+
import { ImageUpscaleNotSupportedError } from '../../core/images/ImageOperationError.js';
12+
13+
// ---------------------------------------------------------------------------
14+
// Helpers
15+
// ---------------------------------------------------------------------------
16+
17+
/** A minimal 1x1 PNG as a Buffer. */
18+
const TINY_PNG = Buffer.from(
19+
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==',
20+
'base64',
21+
);
22+
23+
// ---------------------------------------------------------------------------
24+
// Tests
25+
// ---------------------------------------------------------------------------
26+
27+
describe('upscaleImage', () => {
28+
afterEach(() => {
29+
vi.restoreAllMocks();
30+
});
31+
32+
it('upscales 2x via the Stability provider', async () => {
33+
vi.spyOn(globalThis, 'fetch').mockResolvedValue(
34+
new Response(
35+
JSON.stringify({
36+
image: 'dXBzY2FsZWQ=',
37+
seed: 1,
38+
finish_reason: 'SUCCESS',
39+
}),
40+
{ status: 200, headers: { 'Content-Type': 'application/json' } },
41+
),
42+
);
43+
44+
const result = await upscaleImage({
45+
model: 'stability:stable-image-core',
46+
image: TINY_PNG,
47+
scale: 2,
48+
apiKey: 'stab-key',
49+
});
50+
51+
const [url] = vi.mocked(globalThis.fetch).mock.calls[0];
52+
expect(String(url)).toContain('/v2beta/stable-image/upscale/conservative');
53+
54+
expect(result.provider).toBe('stability');
55+
expect(result.image).toMatchObject({
56+
mimeType: 'image/png',
57+
base64: 'dXBzY2FsZWQ=',
58+
});
59+
});
60+
61+
it('upscales 4x via the Stability provider', async () => {
62+
vi.spyOn(globalThis, 'fetch').mockResolvedValue(
63+
new Response(
64+
JSON.stringify({
65+
image: 'Mng=',
66+
seed: 2,
67+
finish_reason: 'SUCCESS',
68+
}),
69+
{ status: 200, headers: { 'Content-Type': 'application/json' } },
70+
),
71+
);
72+
73+
const result = await upscaleImage({
74+
model: 'stability:stable-image-core',
75+
image: TINY_PNG,
76+
scale: 4,
77+
apiKey: 'stab-key',
78+
});
79+
80+
// Verify the target width was sent (4 * 512 = 2048).
81+
const [, requestInit] = vi.mocked(globalThis.fetch).mock.calls[0];
82+
const formData = requestInit?.body as FormData;
83+
expect(formData.get('width')).toBe('2048');
84+
85+
expect(result.image.base64).toBe('Mng=');
86+
});
87+
88+
it('accepts explicit width/height target dimensions', async () => {
89+
vi.spyOn(globalThis, 'fetch').mockResolvedValue(
90+
new Response(
91+
JSON.stringify({
92+
image: 'd2lkdGg=',
93+
finish_reason: 'SUCCESS',
94+
}),
95+
{ status: 200, headers: { 'Content-Type': 'application/json' } },
96+
),
97+
);
98+
99+
await upscaleImage({
100+
model: 'stability:stable-image-core',
101+
image: TINY_PNG,
102+
width: 3840,
103+
height: 2160,
104+
apiKey: 'stab-key',
105+
});
106+
107+
const [, requestInit] = vi.mocked(globalThis.fetch).mock.calls[0];
108+
const formData = requestInit?.body as FormData;
109+
expect(formData.get('width')).toBe('3840');
110+
expect(formData.get('height')).toBe('2160');
111+
});
112+
113+
it('throws ImageUpscaleNotSupportedError for providers without upscaleImage', async () => {
114+
// OpenAI does not implement upscaleImage.
115+
vi.spyOn(globalThis, 'fetch').mockResolvedValue(
116+
new Response('{}', { status: 200, headers: { 'Content-Type': 'application/json' } }),
117+
);
118+
119+
await expect(
120+
upscaleImage({
121+
model: 'openai:gpt-image-1',
122+
image: TINY_PNG,
123+
scale: 2,
124+
apiKey: 'test-key',
125+
}),
126+
).rejects.toThrow(ImageUpscaleNotSupportedError);
127+
});
128+
129+
it('upscales via the Stable Diffusion local provider extras endpoint', async () => {
130+
vi.spyOn(globalThis, 'fetch').mockImplementation(async (url: string | URL | Request) => {
131+
const urlStr = typeof url === 'string' ? url : url instanceof URL ? url.toString() : (url as Request).url;
132+
133+
if (urlStr.includes('/sdapi/v1/sd-models')) {
134+
return new Response(
135+
JSON.stringify([{ model_name: 'sd-v1-5', title: 'sd-v1-5', filename: 'sd.safetensors' }]),
136+
{ status: 200, headers: { 'Content-Type': 'application/json' } },
137+
);
138+
}
139+
if (urlStr.includes('/sdapi/v1/extra-single-image')) {
140+
return new Response(
141+
JSON.stringify({ image: 'ZXh0cmFz', html_info: '' }),
142+
{ status: 200, headers: { 'Content-Type': 'application/json' } },
143+
);
144+
}
145+
return new Response('Not Found', { status: 404 });
146+
});
147+
148+
const result = await upscaleImage({
149+
model: 'stable-diffusion-local:sd-v1-5',
150+
image: TINY_PNG,
151+
scale: 4,
152+
baseUrl: 'http://localhost:7860',
153+
});
154+
155+
expect(result.provider).toBe('stable-diffusion-local');
156+
expect(result.image.base64).toBe('ZXh0cmFz');
157+
158+
// Verify extras endpoint was called with correct scale.
159+
const calls = vi.mocked(globalThis.fetch).mock.calls;
160+
const extrasCall = calls.find((c) => String(c[0]).includes('/sdapi/v1/extra-single-image'));
161+
expect(extrasCall).toBeDefined();
162+
163+
const body = JSON.parse(String((extrasCall![1] as RequestInit).body));
164+
expect(body.upscaling_resize).toBe(4);
165+
});
166+
167+
it('upscales via the Replicate provider with real-esrgan', async () => {
168+
vi.spyOn(globalThis, 'fetch').mockResolvedValue(
169+
new Response(
170+
JSON.stringify({
171+
status: 'succeeded',
172+
output: ['https://replicate.delivery/upscaled.png'],
173+
}),
174+
{ status: 200, headers: { 'Content-Type': 'application/json' } },
175+
),
176+
);
177+
178+
const result = await upscaleImage({
179+
model: 'replicate:nightmareai/real-esrgan',
180+
image: TINY_PNG,
181+
scale: 4,
182+
apiKey: 'replicate-token',
183+
});
184+
185+
const [, requestInit] = vi.mocked(globalThis.fetch).mock.calls[0];
186+
const body = JSON.parse(String(requestInit?.body));
187+
expect(body.input.scale).toBe(4);
188+
189+
expect(result.provider).toBe('replicate');
190+
expect(result.image.url).toBe('https://replicate.delivery/upscaled.png');
191+
});
192+
});

‎src/api/__tests__/variateImage.test.ts‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
/**
2+
* @file variateImage.test.ts
3+
* Tests for the high-level variateImage API covering native variation
4+
* support (OpenAI), img2img-based fallback, and variance parameter mapping.
5+
*
6+
* All tests mock `globalThis.fetch` — no real API calls are made.
7+
*/
8+
import { afterEach, describe, expect, it, vi } from 'vitest';
9+
10+
import { variateImage } from '../variateImage.js';
11+
12+
// ---------------------------------------------------------------------------
13+
// Helpers
14+
// ---------------------------------------------------------------------------
15+
16+
/** A minimal 1x1 PNG as a Buffer. */
17+
const TINY_PNG = Buffer.from(
18+
'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==',
19+
'base64',
20+
);
21+
22+
// ---------------------------------------------------------------------------
23+
// Tests
24+
// ---------------------------------------------------------------------------
25+
26+
describe('variateImage', () => {
27+
afterEach(() => {
28+
vi.restoreAllMocks();
29+
});
30+
31+
it('creates N variations via the OpenAI variations endpoint', async () => {
32+
vi.spyOn(globalThis, 'fetch').mockResolvedValue(
33+
new Response(
34+
JSON.stringify({
35+
created: 400,
36+
data: [
37+
{ b64_json: 'dmFyMQ==' },
38+
{ b64_json: 'dmFyMg==' },
39+
{ b64_json: 'dmFyMw==' },
40+
],
41+
}),
42+
{ status: 200, headers: { 'Content-Type': 'application/json' } },
43+
),
44+
);
45+
46+
const result = await variateImage({
47+
model: 'openai:dall-e-2',
48+
image: TINY_PNG,
49+
n: 3,
50+
apiKey: 'test-key',
51+
});
52+
53+
// Verify the variations endpoint was used (not edits or generations).
54+
const [url] = vi.mocked(globalThis.fetch).mock.calls[0];
55+
expect(String(url)).toContain('/images/variations');
56+
57+
expect(result.provider).toBe('openai');
58+
expect(result.images).toHaveLength(3);
59+
expect(result.images[0].base64).toBe('dmFyMQ==');
60+
expect(result.images[1].base64).toBe('dmFyMg==');
61+
expect(result.images[2].base64).toBe('dmFyMw==');
62+
});
63+
64+
it('maps the variance parameter to strength for Stability img2img fallback', async () => {
65+
// Stability does not have a native variateImage method, so the high-level
66+
// API should fall back to editImage with strength = variance.
67+
vi.spyOn(globalThis, 'fetch').mockResolvedValue(
68+
new Response(
69+
JSON.stringify({
70+
image: 'dmFyaWF0aW9u',
71+
seed: 55,
72+
finish_reason: 'SUCCESS',
73+
}),
74+
{ status: 200, headers: { 'Content-Type': 'application/json' } },
75+
),
76+
);
77+
78+
const result = await variateImage({
79+
model: 'stability:sd3-medium',
80+
image: TINY_PNG,
81+
variance: 0.3,
82+
apiKey: 'stab-key',
83+
});
84+
85+
// The fallback uses the editImage path which calls the SD3 endpoint.
86+
const [url, requestInit] = vi.mocked(globalThis.fetch).mock.calls[0];
87+
expect(String(url)).toContain('/v2beta/stable-image/generate/sd3');
88+
89+
const formData = requestInit?.body as FormData;
90+
// Variance should be forwarded as strength.
91+
expect(formData.get('strength')).toBe('0.3');
92+
93+
expect(result.provider).toBe('stability');
94+
expect(result.images).toHaveLength(1);
95+
});
96+
97+
it('falls back to img2img for the Stable Diffusion local provider', async () => {
98+
vi.spyOn(globalThis, 'fetch').mockImplementation(async (url: string | URL | Request) => {
99+
const urlStr = typeof url === 'string' ? url : url instanceof URL ? url.toString() : (url as Request).url;
100+
101+
if (urlStr.includes('/sdapi/v1/sd-models')) {
102+
return new Response(
103+
JSON.stringify([{ model_name: 'sd-v1-5', title: 'sd-v1-5', filename: 'sd.safetensors' }]),
104+
{ status: 200, headers: { 'Content-Type': 'application/json' } },
105+
);
106+
}
107+
if (urlStr.includes('/sdapi/v1/img2img')) {
108+
return new Response(
109+
JSON.stringify({
110+
images: ['dmFyMQ==', 'dmFyMg=='],
111+
parameters: {},
112+
info: '{}',
113+
}),
114+
{ status: 200, headers: { 'Content-Type': 'application/json' } },
115+
);
116+
}
117+
return new Response('Not Found', { status: 404 });
118+
});
119+
120+
const result = await variateImage({
121+
model: 'stable-diffusion-local:sd-v1-5',
122+
image: TINY_PNG,
123+
n: 2,
124+
variance: 0.4,
125+
baseUrl: 'http://localhost:7860',
126+
});
127+
128+
expect(result.provider).toBe('stable-diffusion-local');
129+
expect(result.images).toHaveLength(2);
130+
131+
// Verify the img2img endpoint was used with low denoising_strength.
132+
const calls = vi.mocked(globalThis.fetch).mock.calls;
133+
const img2imgCall = calls.find((c) => String(c[0]).includes('/sdapi/v1/img2img'));
134+
expect(img2imgCall).toBeDefined();
135+
136+
const body = JSON.parse(String((img2imgCall![1] as RequestInit).body));
137+
// Variance 0.4 maps to denoising_strength 0.4.
138+
expect(body.denoising_strength).toBe(0.4);
139+
});
140+
141+
it('uses default variance of 0.5 when not specified', async () => {
142+
vi.spyOn(globalThis, 'fetch').mockResolvedValue(
143+
new Response(
144+
JSON.stringify({
145+
image: 'ZGVmYXVsdA==',
146+
seed: 1,
147+
finish_reason: 'SUCCESS',
148+
}),
149+
{ status: 200, headers: { 'Content-Type': 'application/json' } },
150+
),
151+
);
152+
153+
await variateImage({
154+
model: 'stability:sd3-medium',
155+
image: TINY_PNG,
156+
apiKey: 'stab-key',
157+
});
158+
159+
const [, requestInit] = vi.mocked(globalThis.fetch).mock.calls[0];
160+
const formData = requestInit?.body as FormData;
161+
// Default variance 0.5 forwarded as strength.
162+
expect(formData.get('strength')).toBe('0.5');
163+
});
164+
});

0 commit comments

Comments
 (0)