Implementation:Ollama Ollama Imagegen ZImage Scheduler

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

Overview

Implements the Flow Match Euler discrete scheduler for Z-Image, managing timestep generation, noise initialization, and denoising steps.

Description

The scheduler.go file provides FlowMatchEulerScheduler for Z-Image's denoising process. It generates linearly spaced timesteps from 1.0 (full noise) to 0.0 (clean) with optional dynamic time shifting via exponential mu. The Step function performs Euler integration (x_next = x + dt * velocity where dt = sigma_next - sigma). Additional utilities include ScaleSample (identity for flow matching), AddNoise for img2img (x_t = (1-t)*x_0 + t*noise), and InitNoise for bfloat16 random noise generation. GetLatentShape computes the latent dimensions with 8x VAE downscaling. The scheduler matches the diffusers library behavior for compatibility.

Usage

Used by the Z-Image pipeline to manage the denoising loop, from initial noise to final denoised latents.

Code Reference

Source Location

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

Signature

type FlowMatchSchedulerConfig struct {
	NumTrainTimesteps  int32   `json:"num_train_timesteps"`
	Shift              float32 `json:"shift"`
	UseDynamicShifting bool    `json:"use_dynamic_shifting"`
}

type FlowMatchEulerScheduler struct {
	Config    *FlowMatchSchedulerConfig
	Timesteps []float32
	Sigmas    []float32
	NumSteps  int
}

func NewFlowMatchEulerScheduler(cfg *FlowMatchSchedulerConfig) *FlowMatchEulerScheduler
func (s *FlowMatchEulerScheduler) SetTimesteps(numSteps int)
func (s *FlowMatchEulerScheduler) Step(modelOutput, sample *mlx.Array, timestepIdx int) *mlx.Array
func (s *FlowMatchEulerScheduler) AddNoise(cleanSample, noise *mlx.Array, timestepIdx int) *mlx.Array
func (s *FlowMatchEulerScheduler) InitNoise(shape []int32, seed int64) *mlx.Array
func GetLatentShape(batchSize, height, width, latentChannels, patchSize int32) []int32

Import

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

I/O Contract

Inputs

Name	Type	Required	Description
numSteps	int	Yes	Number of denoising steps
modelOutput	*mlx.Array	Yes	Predicted velocity from transformer
sample	*mlx.Array	Yes	Current noisy sample
timestepIdx	int	Yes	Index into timestep schedule

Outputs

Name	Type	Description
*mlx.Array	*mlx.Array	Denoised sample after Euler step

Usage Examples

scheduler := zimage.NewFlowMatchEulerScheduler(zimage.DefaultFlowMatchSchedulerConfig())
scheduler.SetTimesteps(8)

noise := scheduler.InitNoise(latentShape, 42)
for i := 0; i < scheduler.NumSteps; i++ {
    velocity := transformer.Forward(noise, scheduler.GetTimestep(i))
    noise = scheduler.Step(velocity, noise, i)
}

Related Pages

• Principle:Ollama_Ollama_ImageGeneration