Skip to content

Commit 973fae6

Browse files
TimPietruskyRunPodTimPietruskyRunPod
authored
feat(image): add Z-Image Turbo support (#56)
* feat(image): add z-image-turbo support Add endpoint mapping, model-specific size/aspect validation, and docs. * chore: fix lint in z-image tests Rename unused parameter to satisfy eslint rules.
1 parent 871a402 commit 973fae6

File tree

6 files changed

+282
-6
lines changed

6 files changed

+282
-6
lines changed

‎.changeset/add-z-image-turbo.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@runpod/ai-sdk-provider': minor
3+
---
4+
5+
Add support for the Tongyi-MAI Z-Image Turbo image model with validated sizes and aspect ratios.

‎README.md‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -291,6 +291,7 @@ Check out our [examples](https://github.com/runpod/examples/tree/main/ai-sdk/get
291291
| `qwen/qwen-image` | t2i | up to 4096x4096 | 1:1, 4:3, 3:4 |
292292
| `qwen/qwen-image-edit` | edit | up to 4096x4096 | 1:1, 4:3, 3:4 |
293293
| `qwen/qwen-image-edit-2511` | edit | up to 1536x1536 | 1:1, 4:3, 3:4 |
294+
| `tongyi-mai/z-image-turbo` | t2i | up to 1536x1536 | 1:1, 4:3, 3:4, 3:2, 2:3, 16:9, 9:16 |
294295
| `black-forest-labs/flux-1-schnell` | t2i | up to 2048x2048 | 1:1, 4:3, 3:4 |
295296
| `black-forest-labs/flux-1-dev` | t2i | up to 2048x2048 | 1:1, 4:3, 3:4 |
296297
| `black-forest-labs/flux-1-kontext-dev` | edit | up to 2048x2048 | 1:1, 4:3, 3:4 |
@@ -435,6 +436,14 @@ const { image } = await generateImage({
435436
});
436437
```
437438

439+
#### Tongyi-MAI (Z-Image Turbo)
440+
441+
Supported model: `tongyi-mai/z-image-turbo`
442+
443+
- Supported sizes (validated by provider): 512x512, 768x768, 1024x1024, 1280x1280, 1536x1536, 512x768, 768x512, 1024x768, 768x1024, 1328x1328, 1472x1140, 1140x1472, 768x432, 1024x576, 1280x720, 1536x864, 432x768, 576x1024, 720x1280, 864x1536
444+
- Supported `aspectRatio` values: 1:1, 4:3, 3:4, 3:2, 2:3, 16:9, 9:16 (maps to sizes above; use `size` for exact dimensions)
445+
- Additional parameters: `strength`, `output_format`, `enable_safety_checker`, `seed`
446+
438447
## Speech Models
439448

440449
Generate speech using the AI SDK's `generateSpeech` and `runpod.speech(...)`:

‎src/runpod-image-model.test.ts‎

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -349,6 +349,162 @@ describe('RunpodImageModel', () => {
349349
});
350350
});
351351

352+
describe('Z-Image Turbo size validation', () => {
353+
let zImageModel: RunpodImageModel;
354+
355+
beforeEach(() => {
356+
zImageModel = new RunpodImageModel('tongyi-mai/z-image-turbo', {
357+
provider: 'runpod',
358+
baseURL: 'https://api.runpod.ai/v2/z-image-turbo',
359+
headers: () => ({ Authorization: 'Bearer test-key' }),
360+
fetch: mockFetch,
361+
});
362+
});
363+
364+
it('should accept supported sizes', async () => {
365+
const supportedSizes = [
366+
'1328x1328',
367+
'1472x1140',
368+
'1140x1472',
369+
'512x512',
370+
'768x768',
371+
'1024x1024',
372+
'1280x1280',
373+
'1536x1536',
374+
'512x768',
375+
'768x512',
376+
'1024x768',
377+
'768x1024',
378+
'768x432',
379+
'1024x576',
380+
'1280x720',
381+
'1536x864',
382+
'432x768',
383+
'576x1024',
384+
'720x1280',
385+
'864x1536',
386+
];
387+
388+
mockFetch.mockImplementation(async (input: any, _init?: any) => {
389+
const url = typeof input === 'string' ? input : input?.url;
390+
391+
if (url?.includes('/runsync')) {
392+
return new Response(
393+
JSON.stringify({
394+
id: 'test',
395+
status: 'COMPLETED',
396+
output: { result: 'https://cdn.test/z-image.png' },
397+
}),
398+
{ headers: { 'content-type': 'application/json' } }
399+
);
400+
}
401+
402+
if (url === 'https://cdn.test/z-image.png') {
403+
return new Response(new Uint8Array([1, 2, 3]), {
404+
headers: { 'content-type': 'image/png' },
405+
});
406+
}
407+
408+
throw new Error(`Unexpected fetch url: ${String(url)}`);
409+
});
410+
411+
for (const size of supportedSizes) {
412+
await expect(
413+
zImageModel.doGenerate({
414+
prompt: 'Test',
415+
n: 1,
416+
size,
417+
aspectRatio: undefined,
418+
seed: undefined,
419+
providerOptions: {},
420+
headers: {},
421+
abortSignal: undefined,
422+
})
423+
).resolves.toBeDefined();
424+
}
425+
});
426+
427+
it('should reject unsupported sizes', async () => {
428+
await expect(
429+
zImageModel.doGenerate({
430+
prompt: 'Test',
431+
n: 1,
432+
size: '2048x2048',
433+
aspectRatio: undefined,
434+
seed: undefined,
435+
providerOptions: {},
436+
headers: {},
437+
abortSignal: undefined,
438+
})
439+
).rejects.toThrow(InvalidArgumentError);
440+
});
441+
442+
it('should map supported aspect ratios to sizes', async () => {
443+
const aspectRatioToSize = {
444+
'1:1': '1328*1328',
445+
'4:3': '1472*1140',
446+
'3:4': '1140*1472',
447+
'3:2': '768*512',
448+
'2:3': '512*768',
449+
'16:9': '1280*720',
450+
'9:16': '720*1280',
451+
};
452+
453+
for (const [aspectRatio, expectedSize] of Object.entries(
454+
aspectRatioToSize
455+
)) {
456+
let capturedBody: any;
457+
458+
mockFetch.mockImplementationOnce(async (_input: any, init?: any) => {
459+
capturedBody = JSON.parse(init?.body ?? '{}');
460+
return new Response(
461+
JSON.stringify({
462+
id: 'test',
463+
status: 'COMPLETED',
464+
output: { result: 'https://cdn.test/z-image.png' },
465+
}),
466+
{ headers: { 'content-type': 'application/json' } }
467+
);
468+
});
469+
mockFetch.mockImplementationOnce(() =>
470+
Promise.resolve(
471+
new Response(new Uint8Array([1, 2, 3]), {
472+
headers: { 'content-type': 'image/png' },
473+
})
474+
)
475+
);
476+
477+
await zImageModel.doGenerate({
478+
prompt: 'Test',
479+
n: 1,
480+
size: undefined,
481+
aspectRatio,
482+
seed: undefined,
483+
providerOptions: {},
484+
headers: {},
485+
abortSignal: undefined,
486+
});
487+
488+
expect(capturedBody?.input?.size).toBe(expectedSize);
489+
}
490+
});
491+
492+
it('should reject unsupported aspect ratios', async () => {
493+
await expect(
494+
zImageModel.doGenerate({
495+
prompt: 'Test',
496+
n: 1,
497+
size: undefined,
498+
aspectRatio: '21:9',
499+
seed: undefined,
500+
providerOptions: {},
501+
headers: {},
502+
abortSignal: undefined,
503+
})
504+
).rejects.toThrow(InvalidArgumentError);
505+
});
506+
});
507+
352508
describe('parameter conversion', () => {
353509
it('should build correct payload for Qwen models', () => {
354510
const qwenModel = new RunpodImageModel('qwen/qwen-image', {
@@ -375,6 +531,32 @@ describe('RunpodImageModel', () => {
375531
});
376532
});
377533

534+
it('should build correct payload for Z-Image Turbo', () => {
535+
const zImageModel = new RunpodImageModel('tongyi-mai/z-image-turbo', {
536+
provider: 'runpod',
537+
baseURL: 'https://api.runpod.ai/v2/z-image-turbo',
538+
headers: () => ({ Authorization: 'Bearer test-key' }),
539+
fetch: mockFetch,
540+
});
541+
542+
const payload = (zImageModel as any).buildInputPayload(
543+
'Test prompt',
544+
'1024*1024',
545+
42,
546+
{ strength: 0.8, output_format: 'png', enable_safety_checker: true }
547+
);
548+
549+
expect(payload).toMatchObject({
550+
prompt: 'Test prompt',
551+
size: '1024*1024',
552+
seed: 42,
553+
strength: 0.8,
554+
output_format: 'png',
555+
enable_safety_checker: true,
556+
});
557+
expect(payload).not.toHaveProperty('negative_prompt');
558+
});
559+
378560
it('should build correct payload for Flux standard models', () => {
379561
const fluxModel = new RunpodImageModel(
380562
'black-forest-labs/flux-1-schnell',

‎src/runpod-image-model.ts‎

Lines changed: 81 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ const SUPPORTED_ASPECT_RATIOS: Record<string, string> = {
3434
'3:4': '1140*1472', // ✅ Native support
3535
};
3636

37-
// Runpod supported sizes (validated working sizes)
37+
// Runpod supported sizes (validated working sizes for most models)
3838
const SUPPORTED_SIZES = new Set([
3939
// Native aspect ratio sizes
4040
'1328*1328', // 1:1
@@ -54,6 +54,50 @@ const SUPPORTED_SIZES = new Set([
5454
'768*1024',
5555
]);
5656

57+
// Z-Image Turbo supported sizes (validated working sizes)
58+
const Z_IMAGE_TURBO_SUPPORTED_SIZES = new Set([
59+
'1328*1328', // 1:1
60+
'1472*1140', // 4:3
61+
'1140*1472', // 3:4
62+
'512*512',
63+
'768*768',
64+
'1024*1024',
65+
'1280*1280',
66+
'1536*1536',
67+
'512*768',
68+
'768*512',
69+
'1024*768',
70+
'768*1024',
71+
'768*432',
72+
'1024*576',
73+
'1280*720',
74+
'1536*864',
75+
'432*768',
76+
'576*1024',
77+
'720*1280',
78+
'864*1536',
79+
]);
80+
81+
const Z_IMAGE_TURBO_ASPECT_RATIOS: Record<string, string> = {
82+
'1:1': '1328*1328',
83+
'4:3': '1472*1140',
84+
'3:4': '1140*1472',
85+
'3:2': '768*512',
86+
'2:3': '512*768',
87+
'16:9': '1280*720',
88+
'9:16': '720*1280',
89+
};
90+
91+
const MODEL_SUPPORTED_SIZES: Record<string, Set<string>> = {
92+
'tongyi-mai/z-image-turbo': Z_IMAGE_TURBO_SUPPORTED_SIZES,
93+
'z-image-turbo': Z_IMAGE_TURBO_SUPPORTED_SIZES, // alias, not advertised
94+
};
95+
96+
const MODEL_SUPPORTED_ASPECT_RATIOS: Record<string, Record<string, string>> = {
97+
'tongyi-mai/z-image-turbo': Z_IMAGE_TURBO_ASPECT_RATIOS,
98+
'z-image-turbo': Z_IMAGE_TURBO_ASPECT_RATIOS, // alias, not advertised
99+
};
100+
57101
// WAN 2.6 specific aspect ratio to size mappings
58102
// Total pixels must be between 768*768 (589,824) and 1280*1280 (1,638,400)
59103
// Aspect ratio must be between 1:4 and 4:1
@@ -193,11 +237,13 @@ export class RunpodImageModel implements ImageModelV3 {
193237
const runpodSizeCandidate = size.replace('x', '*');
194238

195239
// Validate size is supported
196-
if (!SUPPORTED_SIZES.has(runpodSizeCandidate)) {
240+
const supportedSizes =
241+
MODEL_SUPPORTED_SIZES[this.modelId] ?? SUPPORTED_SIZES;
242+
if (!supportedSizes.has(runpodSizeCandidate)) {
197243
throw new InvalidArgumentError({
198244
argument: 'size',
199245
message: `Size ${size} is not supported by Runpod. Supported sizes: ${Array.from(
200-
SUPPORTED_SIZES
246+
supportedSizes
201247
)
202248
.map((s) => s.replace('*', 'x'))
203249
.join(', ')}`,
@@ -207,15 +253,19 @@ export class RunpodImageModel implements ImageModelV3 {
207253
runpodSize = runpodSizeCandidate;
208254
} else if (aspectRatio) {
209255
// Validate aspect ratio is supported
210-
if (!SUPPORTED_ASPECT_RATIOS[aspectRatio]) {
256+
const supportedAspectRatios =
257+
MODEL_SUPPORTED_ASPECT_RATIOS[this.modelId] ?? SUPPORTED_ASPECT_RATIOS;
258+
if (!supportedAspectRatios[aspectRatio]) {
211259
throw new InvalidArgumentError({
212260
argument: 'aspectRatio',
213-
message: `Aspect ratio ${aspectRatio} is not supported by Runpod. Supported aspect ratios: ${Object.keys(SUPPORTED_ASPECT_RATIOS).join(', ')}`,
261+
message: `Aspect ratio ${aspectRatio} is not supported by Runpod. Supported aspect ratios: ${Object.keys(
262+
supportedAspectRatios
263+
).join(', ')}`,
214264
});
215265
}
216266

217267
// Use supported aspect ratio mapping
218-
runpodSize = SUPPORTED_ASPECT_RATIOS[aspectRatio];
268+
runpodSize = supportedAspectRatios[aspectRatio];
219269
} else {
220270
// Default to square format
221271
runpodSize = '1328*1328';
@@ -662,6 +712,31 @@ export class RunpodImageModel implements ImageModelV3 {
662712
return qwenEdit2511Payload;
663713
}
664714

715+
// Check if this is Tongyi Z-Image Turbo (t2i)
716+
const isZImageTurbo =
717+
this.modelId === 'tongyi-mai/z-image-turbo' ||
718+
this.modelId === 'z-image-turbo';
719+
if (isZImageTurbo) {
720+
const zImageTurboPayload: Record<string, unknown> = {
721+
prompt,
722+
size: runpodSize,
723+
seed: seed ?? -1,
724+
output_format: (runpodOptions?.output_format as string) ?? 'png',
725+
enable_safety_checker: runpodOptions?.enable_safety_checker ?? true,
726+
...runpodOptions,
727+
};
728+
729+
if (standardizedImages && standardizedImages.length > 0) {
730+
if (standardizedImages.length === 1) {
731+
zImageTurboPayload.image = standardizedImages[0];
732+
} else {
733+
zImageTurboPayload.images = standardizedImages;
734+
}
735+
}
736+
737+
return zImageTurboPayload;
738+
}
739+
665740
// Check if this is an Alibaba Wan model
666741
const isWanModel = this.modelId.includes('wan-2');
667742
if (isWanModel) {

‎src/runpod-image-options.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ export type RunpodImageModelId =
1010
| 'black-forest-labs/flux-1-dev'
1111
// Alibaba Wan 2.6 (t2i)
1212
| 'alibaba/wan-2.6'
13+
// Tongyi Z-Image Turbo (t2i)
14+
| 'tongyi-mai/z-image-turbo'
1315
// Nano Banana (edit only)
1416
| 'google/nano-banana-edit'
1517
| 'nano-banana-edit'; // backwards compatibility

‎src/runpod-provider.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,9 @@ const IMAGE_MODEL_ID_TO_ENDPOINT_URL: Record<string, string> = {
122122
'https://api.runpod.ai/v2/black-forest-labs-flux-1-dev',
123123
// Alibaba Wan 2.6 (t2i)
124124
'alibaba/wan-2.6': 'https://api.runpod.ai/v2/wan-2-6-t2i',
125+
// Tongyi Z-Image Turbo (t2i)
126+
'tongyi-mai/z-image-turbo': 'https://api.runpod.ai/v2/z-image-turbo',
127+
'z-image-turbo': 'https://api.runpod.ai/v2/z-image-turbo', // alias, not advertised
125128
// Nano Banana (edit only)
126129
'google/nano-banana-edit': 'https://api.runpod.ai/v2/nano-banana-edit',
127130
'nano-banana-edit': 'https://api.runpod.ai/v2/nano-banana-edit', // backwards compatibility

0 commit comments

Comments
 (0)