Implementation:Ollama Ollama Imagegen ZImage

Knowledge Sources	Ollama
Domains	Image Generation, Diffusion Models
Last Updated	2025-02-15 00:00 GMT

Overview

Implements the Z-Image diffusion pipeline for text-to-image generation, supporting CFG (classifier-free guidance), TeaCache, and tiled VAE decoding.

Description

The zimage.go file is the top-level orchestrator for the Z-Image pipeline, coordinating text encoding, latent noise initialization, iterative denoising with flow-match scheduling, and VAE decoding. The Model struct loads text_encoder (Qwen3), transformer, and VAE components from Ollama's manifest storage. GenerateFromConfig handles the full pipeline with options for negative prompts (CFG), TeaCache for accelerated inference (skipping recomputation on similar timesteps), fused QKV projections, and dynamic time-shift scheduling. It supports tiled VAE decoding for large images and progress callbacks for UI integration.

Usage

Used as the default image generation pipeline for Z-Image and FluxPipeline model types via Ollama's image generation API.

Code Reference

Source Location

• Repository: Ollama
• File: x/imagegen/models/zimage/zimage.go
• Lines: 1-490

Signature

type GenerateConfig struct {
	Prompt            string
	NegativePrompt    string
	CFGScale          float32
	Width             int32
	Height            int32
	Steps             int
	Seed              int64
	CapturePath       string
	TeaCache          bool
	TeaCacheThreshold float32
	FusedQKV          bool
}

type Model struct {
	ModelName       string
	Tokenizer       *tokenizer.Tokenizer
	TextEncoder     *qwen3.TextEncoder
	Transformer     *Transformer
	VAE             *AutoencoderKL
	SchedulerConfig *FlowMatchSchedulerConfig
}

func (m *Model) Load(modelName string) error
func (m *Model) GenerateImage(ctx context.Context, prompt string, width, height int32, steps int, seed int64, progress func(step, total int)) (*mlx.Array, error)
func (m *Model) GenerateFromConfig(ctx context.Context, cfg *GenerateConfig) (*mlx.Array, error)

Import

import "github.com/ollama/ollama/x/imagegen/models/zimage"

I/O Contract

Inputs

Name	Type	Required	Description
modelName	string	Yes	Ollama model name for manifest resolution
cfg	*GenerateConfig	Yes	Full generation configuration

Outputs

Name	Type	Description
*mlx.Array	*mlx.Array	Generated image [B, C, H, W] with values in [0, 1]
error	error	Error from loading or generation

Usage Examples

m := &zimage.Model{}
if err := m.Load("zimage-turbo:latest"); err != nil {
    return err
}

img, err := m.GenerateFromConfig(ctx, &zimage.GenerateConfig{
    Prompt:         "a cat in a garden",
    NegativePrompt: "blurry, low quality",
    CFGScale:       4.0,
    Width:          1024,
    Height:         1024,
    Steps:          8,
    Seed:           42,
    TeaCache:       true,
})

Related Pages

• Principle:Ollama_Ollama_ImageGeneration