diff --git a/README.md b/README.md index a2062bd..2c44137 100644 --- a/README.md +++ b/README.md @@ -162,6 +162,9 @@ ollama.generate(request) - `logprobs` ``: (Optional) Return log probabilities for tokens. Requires model support. - `top_logprobs` ``: (Optional) Number of top log probabilities to return per token when `logprobs` is enabled. - `keep_alive` ``: (Optional) How long to keep the model loaded. A number (seconds) or a string with a duration unit suffix ("300ms", "1.5h", "2h45m", etc.) + - `width` ``: (Optional, Experimental) Width of the generated image in pixels. For image generation models only. + - `height` ``: (Optional, Experimental) Height of the generated image in pixels. For image generation models only. + - `steps` ``: (Optional, Experimental) Number of diffusion steps. For image generation models only. - `options` ``: (Optional) Options to configure the runtime. - Returns: `` diff --git a/examples/README.md b/examples/README.md index 66c21c7..cc09c8e 100644 --- a/examples/README.md +++ b/examples/README.md @@ -8,3 +8,10 @@ To run the examples run: ```sh npx tsx /.ts ``` + +### Image Generation (Experimental) + +> **Note:** Image generation is experimental and currently only available on macOS. + +- [image-generation/image-generation.ts](image-generation/image-generation.ts) +- [image-generation/image-generation-stream.ts](image-generation/image-generation-stream.ts) - Streamed progress diff --git a/examples/image-generation/image-generation.ts b/examples/image-generation/image-generation.ts new file mode 100644 index 0000000..c6ed867 --- /dev/null +++ b/examples/image-generation/image-generation.ts @@ -0,0 +1,29 @@ +// Image generation is experimental and currently only available on macOS + +import ollama from 'ollama' +import { writeFileSync } from 'fs' + +async function main() { + const prompt = 'a sunset over mountains' + console.log(`Prompt: ${prompt}`) + + const response = await ollama.generate({ + model: 'x/z-image-turbo', + prompt, + stream: true, + }) + + for await (const part of response) { + if (part.image) { + // Final response contains the image + const imageBuffer = Buffer.from(part.image, 'base64') + writeFileSync('output.png', imageBuffer) + console.log('\nImage saved to output.png') + } else if (part.total) { + // Progress update + process.stdout.write(`\rProgress: ${part.completed}/${part.total}`) + } + } +} + +main().catch(console.error) diff --git a/src/interfaces.ts b/src/interfaces.ts index 2b5527b..45c6b74 100644 --- a/src/interfaces.ts +++ b/src/interfaces.ts @@ -61,6 +61,11 @@ export interface GenerateRequest { logprobs?: boolean top_logprobs?: number + // Experimental image generation parameters + width?: number + height?: number + steps?: number + options?: Partial } @@ -191,7 +196,7 @@ export interface Logprob extends TokenLogprob { export interface GenerateResponse { model: string created_at: Date - response: string + response?: string thinking?: string done: boolean done_reason: string @@ -203,6 +208,11 @@ export interface GenerateResponse { eval_count: number eval_duration: number logprobs?: Logprob[] + + // Image generation response fields + image?: string // Base64-encoded generated image data + completed?: number // Number of completed steps (for streaming progress) + total?: number // Total number of steps (for streaming progress) } export interface ChatResponse { diff --git a/test/browser.test.ts b/test/browser.test.ts index 0cf0e61..e5d505e 100644 --- a/test/browser.test.ts +++ b/test/browser.test.ts @@ -56,3 +56,89 @@ describe('Ollama logprob request fields', () => { ) }) }) + +describe('Ollama image generation request fields', () => { + it('forwards image generation parameters in generate requests', async () => { + const client = new Ollama() + const spy = vi + .spyOn(client as any, 'processStreamableRequest') + .mockResolvedValue({} as GenerateResponse) + + await client.generate({ + model: 'dummy-image', + prompt: 'a sunset over mountains', + width: 1024, + height: 768, + steps: 20, + }) + + expect(spy).toHaveBeenCalledWith( + 'generate', + expect.objectContaining({ + model: 'dummy-image', + prompt: 'a sunset over mountains', + width: 1024, + height: 768, + steps: 20, + }), + ) + }) + + it('handles image generation response with image field', async () => { + const mockResponse: GenerateResponse = { + model: 'dummy-image', + created_at: new Date(), + done: true, + done_reason: 'stop', + context: [], + total_duration: 1000, + load_duration: 100, + prompt_eval_count: 10, + prompt_eval_duration: 50, + eval_count: 0, + eval_duration: 0, + image: 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==', + } + + const client = new Ollama() + vi.spyOn(client as any, 'processStreamableRequest').mockResolvedValue(mockResponse) + + const response = await client.generate({ + model: 'dummy-image', + prompt: 'a sunset', + }) + + expect(response.image).toBeDefined() + expect(response.done).toBe(true) + }) + + it('handles streaming progress fields for image generation', async () => { + const mockResponse: GenerateResponse = { + model: 'dummy-image', + created_at: new Date(), + done: false, + done_reason: '', + context: [], + total_duration: 0, + load_duration: 0, + prompt_eval_count: 0, + prompt_eval_duration: 0, + eval_count: 0, + eval_duration: 0, + completed: 5, + total: 20, + } + + const client = new Ollama() + vi.spyOn(client as any, 'processStreamableRequest').mockResolvedValue(mockResponse) + + const response = await client.generate({ + model: 'dummy-image', + prompt: 'a sunset', + }) + + expect(response.completed).toBe(5) + expect(response.total).toBe(20) + expect(response.done).toBe(false) + }) +})