Implementation:Ollama Ollama Imagegen Cache TeaCache

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

Overview

Implements Timestep Embedding Aware Caching (TeaCache) for diffusion models, reusing transformer outputs when consecutive timesteps are similar.

Description

The teacache.go file implements the TeaCache algorithm based on "Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model". It caches the full transformer output and determines reuse based on the absolute difference between consecutive timestep values scaled by a model-specific rescale factor. The accumulated difference is tracked to detect drift over multiple cached steps. When the difference exceeds a configurable threshold, the cache is invalidated and a new forward pass is computed. TeaCache supports both non-CFG (single output) and CFG (separate pos/neg outputs) modes, where CFG is always recomputed fresh from cached components to avoid error amplification. Early steps can be forced to always compute for structural quality.

Usage

Used by the Z-Image pipeline when --teacache is enabled, reducing denoising computation by 30-50% with minimal quality loss.

Code Reference

Source Location

• Repository: Ollama
• File: x/imagegen/cache/teacache.go
• Lines: 1-197

Signature

type TeaCache struct {
	cachedOutput    *mlx.Array
	cachedPosOutput *mlx.Array
	cachedNegOutput *mlx.Array
	prevTimestep    float32
	accumulatedDiff float32
	threshold       float32
	rescaleFactor   float32
	skipEarlySteps  int
	cacheHits       int
	cacheMisses     int
}

type TeaCacheConfig struct {
	Threshold      float32
	RescaleFactor  float32
	SkipEarlySteps int
}

func NewTeaCache(cfg *TeaCacheConfig) *TeaCache
func (tc *TeaCache) ShouldCompute(step int, timestep float32) bool
func (tc *TeaCache) UpdateCache(output *mlx.Array, timestep float32)
func (tc *TeaCache) UpdateCFGCache(posOutput, negOutput *mlx.Array, timestep float32)
func (tc *TeaCache) GetCached() *mlx.Array
func (tc *TeaCache) GetCFGCached() (pos, neg *mlx.Array)
func (tc *TeaCache) Stats() (hits, misses int)
func (tc *TeaCache) Free()

Import

import "github.com/ollama/ollama/x/imagegen/cache"

I/O Contract

Inputs

Name	Type	Required	Description
cfg	*TeaCacheConfig	No	Configuration (nil uses defaults: threshold=0.1, rescale=1.0)
step	int	Yes	Current denoising step
timestep	float32	Yes	Current timestep value

Outputs

Name	Type	Description
bool	bool	Whether full computation is needed (true) or cache can be used (false)

Usage Examples

tc := cache.NewTeaCache(&cache.TeaCacheConfig{
    Threshold:      0.1,
    RescaleFactor:  1.0,
    SkipEarlySteps: 2,
})
defer tc.Free()

for step, timestep := range scheduler.Timesteps {
    if tc.ShouldCompute(step, timestep) {
        output := transformer.Forward(latents, timestep)
        tc.UpdateCache(output, timestep)
    } else {
        output = tc.GetCached()
    }
}

Related Pages

• Principle:Ollama_Ollama_ImageGeneration