Add image generation support (#271)

Author: jmorganca

README.md

@@ -162,6 +162,9 @@ ollama.generate(request)
   - `logprobs` `<boolean>`: (Optional) Return log probabilities for tokens. Requires model support.
   - `top_logprobs` `<number>`: (Optional) Number of top log probabilities to return per token when `logprobs` is enabled.
   - `keep_alive` `<string | number>`: (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` `<number>`: (Optional, Experimental) Width of the generated image in pixels. For image generation models only.
+  - `height` `<number>`: (Optional, Experimental) Height of the generated image in pixels. For image generation models only.
+  - `steps` `<number>`: (Optional, Experimental) Number of diffusion steps. For image generation models only.
   - `options` `<Options>`: (Optional) Options to configure the runtime.
 - Returns: `<GenerateResponse>`
 

examples/README.md

@@ -8,3 +8,10 @@ To run the examples run:
 ```sh
 npx tsx <folder-name>/<file-name>.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

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)

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<Options>
 }
 
@@ -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 {

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)
+  })
+})