Implementation:Ollama Ollama Imagegen LLM
| Knowledge Sources | |
|---|---|
| Domains | Image Generation, LLM Inference |
| Last Updated | 2025-02-15 00:00 GMT |
Overview
Provides LLM text generation support within the imagegen server, enabling MLX-based language model inference alongside image generation.
Description
The llm.go file implements the TextModel interface and Decoder for autoregressive text generation within the imagegen HTTP server. It defines prefill (chunked processing with memory management) and step (single token decode) methods, manages a dedicated MLX generation stream, and handles token sampling. The llmState holds the loaded model, and generation is serialized via a mutex to ensure safe MLX usage. This enables the same server process to handle both LLM and image generation requests.
Usage
Used when the imagegen server is loaded in LLM mode to handle text completion requests via the /completion endpoint.
Code Reference
Source Location
- Repository: Ollama
- File: x/imagegen/llm.go
- Lines: 1-420
Signature
type TextModel interface {
Forward(tokens *mlx.Array, caches []cache.Cache) *mlx.Array
NewCache(maxSeqLen int32) []cache.Cache
Tokenizer() *tokenizer.Tokenizer
VocabSize() int32
MaxContextLength() int32
NumLayers() int
}
type Decoder struct {
model TextModel
caches []cache.Cache
vocabSize int32
temp float32
token *mlx.Array
oldCacheState []*mlx.Array
}
func NewDecoder(m TextModel, temp float32) *Decoder
func (d *Decoder) prefill(inputIDs []int32) int
func (d *Decoder) step() int32
Import
import "github.com/ollama/ollama/x/imagegen"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| m | TextModel | Yes | Loaded text model implementing forward pass |
| temp | float32 | Yes | Sampling temperature for generation |
Outputs
| Name | Type | Description |
|---|---|---|
| *Decoder | *Decoder | Decoder managing prefill and decode for LLM generation |
Usage Examples
decoder := NewDecoder(textModel, 0.7)
numPrefilled := decoder.prefill(inputTokens)
nextToken := decoder.step()