Implementation:Ollama Ollama Imagegen Flux2 Scheduler
| Knowledge Sources | |
|---|---|
| Domains | Image Generation, Diffusion Models |
| Last Updated | 2025-02-15 00:00 GMT |
Overview
Implements the Flow-Match Euler discrete scheduler for FLUX.2 Klein with dynamic time shifting and empirical mu calculation.
Description
The scheduler.go file provides the FlowMatchScheduler for FLUX.2 Klein's denoising process. It generates evenly spaced sigmas from 1.0 to 0.0 with configurable time shifting: fixed shift (sigma * shift / (1 + (shift-1)*sigma)) or dynamic exponential/linear shifting with a learned mu parameter. The Step function performs Euler integration (x_{t-dt} = x_t + dt * v_t) with float32 upcasting for precision matching diffusers behavior. CalculateShift computes the empirical mu from image sequence length using piecewise linear fits (a1*seqLen + b1 for small images, a2*seqLen + b2 for large) with interpolation between step counts.
Usage
Used by the FLUX.2 pipeline to manage timestep scheduling, noise initialization, and denoising step computation.
Code Reference
Source Location
- Repository: Ollama
- File: x/imagegen/models/flux2/scheduler.go
- Lines: 1-149
Signature
type SchedulerConfig struct {
NumTrainTimesteps int32 `json:"num_train_timesteps"`
Shift float32 `json:"shift"`
UseDynamicShifting bool `json:"use_dynamic_shifting"`
TimeShiftType string `json:"time_shift_type"`
}
type FlowMatchScheduler struct {
Config *SchedulerConfig
Timesteps []float32
Sigmas []float32
NumSteps int
}
func NewFlowMatchScheduler(cfg *SchedulerConfig) *FlowMatchScheduler
func (s *FlowMatchScheduler) SetTimestepsWithMu(numSteps int, mu float32)
func (s *FlowMatchScheduler) Step(modelOutput, sample *mlx.Array, timestepIdx int) *mlx.Array
func (s *FlowMatchScheduler) InitNoise(shape []int32, seed int64) *mlx.Array
func CalculateShift(imgSeqLen int32, numSteps int) float32
Import
import "github.com/ollama/ollama/x/imagegen/models/flux2"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| numSteps | int | Yes | Number of inference denoising steps |
| mu | float32 | No | Dynamic time shift parameter (0 for fixed shifting) |
| modelOutput | *mlx.Array | Yes | Predicted velocity from transformer |
| sample | *mlx.Array | Yes | Current noisy sample |
Outputs
| Name | Type | Description |
|---|---|---|
| *mlx.Array | *mlx.Array | Denoised sample after one Euler step |
Usage Examples
scheduler := flux2.NewFlowMatchScheduler(flux2.DefaultSchedulerConfig())
mu := flux2.CalculateShift(imgSeqLen, 4)
scheduler.SetTimestepsWithMu(4, mu)
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)
}