Implementation:Ollama Ollama Imagegen Flux2

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

Overview

Implements the FLUX.2 Klein diffusion pipeline for text-to-image and image-to-image generation on MLX.

Description

The flux2.go file implements the full FLUX.2 Klein model pipeline, a 4B parameter distilled diffusion transformer supporting sub-second inference. The Model struct orchestrates loading of the tokenizer, Qwen3 text encoder, Flux2 transformer, and VAE from Ollama's manifest storage. GenerateFromConfig handles the complete generation flow: text encoding with multi-layer pooling (layers 8, 17, 26), latent noise initialization, flow-match scheduling with dynamic time shifting, iterative denoising, and VAE decoding with optional tiling for large images. It supports both text-to-image and image-to-image editing with reference image conditioning.

Usage

Used as the primary pipeline when generating images with FLUX.2 Klein models via Ollama's image generation API.

Code Reference

Source Location

• Repository: Ollama
• File: x/imagegen/models/flux2/flux2.go
• Lines: 1-553

Signature

type GenerateConfig struct {
	Prompt        string
	Width         int32
	Height        int32
	Steps         int
	GuidanceScale float32
	Seed          int64
	Progress      func(step, totalSteps int)
	InputImages   []image.Image
}

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

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

Import

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

I/O Contract

Inputs

Name	Type	Required	Description
modelName	string	Yes	Ollama model name for manifest lookup
cfg	*GenerateConfig	Yes	Generation configuration with prompt, dimensions, steps

Outputs

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

Usage Examples

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

img, err := m.GenerateFromConfig(ctx, &flux2.GenerateConfig{
    Prompt: "a sunset over mountains",
    Width:  1024,
    Height: 1024,
    Steps:  4,
    Seed:   42,
})

Related Pages

• Principle:Ollama_Ollama_ImageGeneration