Merge pull request #145 from NianJiuZst/feat/sdk-image-download

feat: add image download and save support to ImageSDK

Author: RyanLee-Dev

src/sdk/image/index.ts

@@ -1,11 +1,27 @@
+import { existsSync, mkdirSync, writeFileSync } from 'node:fs';
+import { join, resolve, dirname } from 'node:path';
 import { Client } from "../client";
 import { imageEndpoint } from "../../client/endpoints";
 import { ImageRequest, ImageResponse } from "../../types/api";
 import { ModelPartial } from "../types";
 import { SDKError } from "../../errors/base";
 import { ExitCode } from "../../errors/codes";
+import { downloadFile } from "../../files/download";
 import { toMerged } from 'es-toolkit/object';
 
+export interface ImageSaveOptions {
+  /** Save to exact file path (single image only). */
+  out?: string;
+  /** Save images to directory (default: "."). */
+  outDir?: string;
+  /** Filename prefix (default: "image"). */
+  prefix?: string;
+  /** Response format used by generate() — "url" or "base64" (default: "url"). */
+  responseFormat?: 'url' | 'base64';
+  /** Suppress progress output. */
+  quiet?: boolean;
+}
+
 export class ImageSDK extends Client {
   async generate(request: ModelPartial<ImageRequest>): Promise<ImageResponse> {
     const body = this.validateParams(request);
@@ -18,6 +34,69 @@ export class ImageSDK extends Client {
     });
   }
 
+  /**
+   * Download and save images from a `generate()` response to disk.
+   *
+   * Handles both `"url"` (CDN download) and `"base64"` response formats
+   * and creates intermediate directories as needed. Returns the absolute
+   * paths of all saved files.
+   */
+  async save(response: ImageResponse, options: ImageSaveOptions = {}): Promise<string[]> {
+    const fmt = options.responseFormat || 'url';
+    const quiet = options.quiet ?? true;
+    const saved: string[] = [];
+
+    if (options.out) {
+      // Single-image, exact path
+      const count = (fmt === 'base64'
+        ? (response.data.image_base64 || []).length
+        : (response.data.image_urls || []).length);
+
+      if (count > 1) {
+        throw new SDKError(
+          'Cannot use `out` with multiple images. Use `outDir` instead.',
+          ExitCode.USAGE,
+        );
+      }
+
+      const destPath = resolve(options.out);
+      const dir = dirname(destPath);
+      if (!existsSync(dir)) mkdirSync(dir, { recursive: true });
+
+      if (fmt === 'base64') {
+        const image = (response.data.image_base64 || [])[0];
+        if (image) writeFileSync(destPath, image, 'base64');
+      } else {
+        const imageUrl = (response.data.image_urls || [])[0];
+        if (imageUrl) await downloadFile(imageUrl, destPath, { quiet });
+      }
+      saved.push(destPath);
+    } else {
+      // Multi-image, numbered filenames in a directory
+      const outDir = resolve(options.outDir || '.');
+      if (!existsSync(outDir)) mkdirSync(outDir, { recursive: true });
+      const prefix = options.prefix || 'image';
+
+      if (fmt === 'base64') {
+        const images = response.data.image_base64 || [];
+        for (let i = 0; i < images.length; i++) {
+          const destPath = join(outDir, `${prefix}_${String(i + 1).padStart(3, '0')}.jpg`);
+          writeFileSync(destPath, images[i]!, 'base64');
+          saved.push(destPath);
+        }
+      } else {
+        const imageUrls = response.data.image_urls || [];
+        for (let i = 0; i < imageUrls.length; i++) {
+          const destPath = join(outDir, `${prefix}_${String(i + 1).padStart(3, '0')}.jpg`);
+          await downloadFile(imageUrls[i]!, destPath, { quiet });
+          saved.push(destPath);
+        }
+      }
+    }
+
+    return saved;
+  }
+
   private validateParams(params: Partial<ImageRequest>): ImageRequest {
     const { width, height, aspect_ratio } = params;
 

test/sdk/image.test.ts

@@ -1,6 +1,35 @@
 import { describe, it, expect, afterEach } from 'bun:test';
 import { createMockServer, jsonResponse, type MockServer } from '../helpers/mock-server';
 import { MiniMaxSDK } from '../../src/sdk';
+import { ImageSDK, ImageSaveOptions } from '../../src/sdk/image';
+import { existsSync, unlinkSync, rmdirSync } from 'node:fs';
+import { join } from 'node:path';
+import { tmpdir } from 'node:os';
+import type { ImageResponse } from '../../src/types/api';
+
+function makeBase64Response(images: string[]): ImageResponse {
+  return {
+    base_resp: { status_code: 0, status_msg: 'ok' },
+    data: {
+      image_base64: images,
+      task_id: 'task-456',
+      success_count: images.length,
+      failed_count: 0,
+    },
+  };
+}
+
+function makeUrlResponse(urls: string[]): ImageResponse {
+  return {
+    base_resp: { status_code: 0, status_msg: 'ok' },
+    data: {
+      image_urls: urls,
+      task_id: 'task-123',
+      success_count: urls.length,
+      failed_count: 0,
+    },
+  };
+}
 
 describe('MiniMaxSDK.image', () => {
   let server: MockServer;
@@ -37,3 +66,58 @@ describe('MiniMaxSDK.image', () => {
     expect(result.data.task_id).toBe('img-123');
   });
 });
+
+describe('ImageSDK.save', () => {
+  const sdk = new ImageSDK({ apiKey: 'sk-test', region: 'global' });
+
+  it('saves a single base64 image with `out` option', async () => {
+    const tmpFile = join(tmpdir(), `mmx-sdk-save-${Date.now()}.jpg`);
+    const response = makeBase64Response(['/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAA==']);
+
+    const saved = await sdk.save(response, { out: tmpFile, responseFormat: 'base64' });
+
+    expect(saved).toEqual([tmpFile]);
+    expect(existsSync(tmpFile)).toBe(true);
+    unlinkSync(tmpFile);
+  });
+
+  it('saves multiple base64 images to a directory', async () => {
+    const tmpDir = join(tmpdir(), `mmx-sdk-imgs-${Date.now()}`);
+    const response = makeBase64Response([
+      '/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAA==',
+      '/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAA==',
+    ]);
+
+    const saved = await sdk.save(response, {
+      outDir: tmpDir,
+      prefix: 'test',
+      responseFormat: 'base64',
+    });
+
+    expect(saved.length).toBe(2);
+    expect(saved[0]).toContain('test_001.jpg');
+    expect(saved[1]).toContain('test_002.jpg');
+    saved.forEach(p => expect(existsSync(p)).toBe(true));
+    saved.forEach(p => unlinkSync(p));
+    rmdirSync(tmpDir);
+  });
+
+  it('throws when `out` is used with multiple images', async () => {
+    const response = makeUrlResponse([
+      'https://example.com/img1.jpg',
+      'https://example.com/img2.jpg',
+    ]);
+    await expect(
+      sdk.save(response, { out: '/tmp/single.jpg', responseFormat: 'url' }),
+    ).rejects.toThrow('Cannot use `out` with multiple images');
+  });
+
+  it('uses current directory by default', async () => {
+    const response = makeBase64Response(['/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAA==']);
+    const saved = await sdk.save(response, { responseFormat: 'base64' });
+    expect(saved.length).toBe(1);
+    expect(saved[0]).toContain('image_001.jpg');
+    expect(existsSync(saved[0]!)).toBe(true);
+    unlinkSync(saved[0]!);
+  });
+});