Affected by 
Documentation
¶
Overview ¶
Package cache provides caching mechanisms for diffusion model inference.
Index ¶
- type Cache
- type KVCache
- type RotatingKVCache
- type StepCache
- func (c *StepCache) Arrays() []*mlx.Array
- func (c *StepCache) Free()
- func (c *StepCache) Get(layer int) *mlx.Array
- func (c *StepCache) Get2(layer int) *mlx.Array
- func (c *StepCache) GetConstant() *mlx.Array
- func (c *StepCache) NumLayers() int
- func (c *StepCache) Set(layer int, arr *mlx.Array)
- func (c *StepCache) Set2(layer int, arr *mlx.Array)
- func (c *StepCache) SetConstant(arr *mlx.Array)
- func (c *StepCache) ShouldRefresh(step, interval int) bool
- type TeaCache
- func (tc *TeaCache) Arrays() []*mlx.Array
- func (tc *TeaCache) Free()
- func (tc *TeaCache) GetCFGCached() (pos, neg *mlx.Array)
- func (tc *TeaCache) GetCached() *mlx.Array
- func (tc *TeaCache) HasCFGCache() bool
- func (tc *TeaCache) ShouldCompute(step int, timestep float32) bool
- func (tc *TeaCache) Stats() (hits, misses int)
- func (tc *TeaCache) UpdateCFGCache(posOutput, negOutput *mlx.Array, timestep float32)
- func (tc *TeaCache) UpdateCache(output *mlx.Array, timestep float32)
- type TeaCacheConfig
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
This section is empty.
Types ¶
type KVCache ¶
type KVCache struct {
// contains filtered or unexported fields
}
func NewKVCache ¶
func NewKVCache() *KVCache
type RotatingKVCache ¶
type RotatingKVCache struct {
// contains filtered or unexported fields
}
RotatingKVCache implements sliding window attention with bounded memory
func NewRotatingKVCache ¶
func NewRotatingKVCache(maxSize int) *RotatingKVCache
func (*RotatingKVCache) Len ¶
func (c *RotatingKVCache) Len() int
func (*RotatingKVCache) Offset ¶
func (c *RotatingKVCache) Offset() int
func (*RotatingKVCache) Reset ¶
func (c *RotatingKVCache) Reset()
Reset clears the cache state for a new generation session
func (*RotatingKVCache) State ¶
func (c *RotatingKVCache) State() []*mlx.Array
type StepCache ¶
type StepCache struct {
// contains filtered or unexported fields
}
StepCache caches layer outputs across diffusion denoising steps. Based on DeepCache (CVPR 2024) and Learning-to-Cache (NeurIPS 2024): shallow layers change little between consecutive steps, so we can cache their outputs and skip recomputation on non-refresh steps.
Supports both single-stream and dual-stream architectures:
- Single-stream: use Get/Set for the single output per layer
- Dual-stream: use Get/Set for stream 1 (imgH), Get2/Set2 for stream 2 (txtH)
Usage (single-stream):
cache := NewStepCache(15) // cache first 15 layers
for step := 0; step < numSteps; step++ {
refresh := cache.ShouldRefresh(step, 3) // refresh every 3 steps
for i, layer := range layers {
if i < 15 && !refresh && cache.Get(i) != nil {
output = cache.Get(i) // reuse cached
} else {
output = layer.Forward(input)
if i < 15 && refresh {
cache.Set(i, output)
}
}
}
}
cache.Free() // cleanup when done
Usage (dual-stream):
cache := NewStepCache(15)
for step := 0; step < numSteps; step++ {
refresh := cache.ShouldRefresh(step, 3)
for i, layer := range layers {
if i < 15 && !refresh && cache.Get(i) != nil {
imgH, txtH = cache.Get(i), cache.Get2(i)
} else {
imgH, txtH = layer.Forward(imgH, txtH, ...)
if i < 15 && refresh {
cache.Set(i, imgH)
cache.Set2(i, txtH)
}
}
}
}
func NewStepCache ¶
NewStepCache creates a cache for the given number of layers.
func (*StepCache) Free ¶
func (c *StepCache) Free()
Free releases all cached arrays. Call when generation completes.
func (*StepCache) Get2 ¶
Get2 returns the cached output for a layer (stream 2), or nil if not cached. Used for dual-stream architectures.
func (*StepCache) GetConstant ¶
GetConstant returns the cached constant value.
func (*StepCache) Set2 ¶
Set2 stores a layer output (stream 2), freeing any previous value. Used for dual-stream architectures.
func (*StepCache) SetConstant ¶
SetConstant stores a constant value, freeing any previous value.
func (*StepCache) ShouldRefresh ¶
ShouldRefresh returns true if the cache should be refreshed at this step. Refresh happens on step 0, interval, 2*interval, etc.
type TeaCache ¶
type TeaCache struct {
// contains filtered or unexported fields
}
TeaCache implements Timestep Embedding Aware Caching for diffusion models. It caches the transformer output and reuses it when timestep values are similar between consecutive steps.
For CFG (classifier-free guidance), it caches pos and neg predictions separately and always computes CFG fresh to avoid error amplification.
Reference: "Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model" https://github.com/ali-vilab/TeaCache
func NewTeaCache ¶
func NewTeaCache(cfg *TeaCacheConfig) *TeaCache
NewTeaCache creates a new TeaCache instance.
func (*TeaCache) GetCFGCached ¶
GetCFGCached returns cached pos and neg outputs for CFG mode.
func (*TeaCache) HasCFGCache ¶
HasCFGCache returns true if CFG cache is available.
func (*TeaCache) ShouldCompute ¶
ShouldCompute determines if we should compute the full forward pass or reuse the cached output based on timestep similarity.
Algorithm: 1. First step always computes 2. Subsequent steps compare |currTimestep - prevTimestep| * rescaleFactor 3. If accumulated difference > threshold, compute new output 4. Otherwise, reuse cached output
func (*TeaCache) UpdateCFGCache ¶
UpdateCFGCache stores pos and neg outputs separately for CFG mode. This allows CFG to be computed fresh each step, avoiding error amplification.
type TeaCacheConfig ¶
type TeaCacheConfig struct {
// Threshold for recomputation. Lower = more cache hits, potential quality loss.
// Recommended: 0.05-0.15 for image models
Threshold float32
// Rescale factor to adjust timestep embedding differences.
// Model-specific, typically 1.0-2.0
RescaleFactor float32
// SkipEarlySteps: number of early steps to always compute (never cache).
// Set to 2-3 for CFG mode to preserve structure. 0 = no skipping.
SkipEarlySteps int
}
TeaCacheConfig holds configuration for TeaCache.
func DefaultTeaCacheConfig ¶
func DefaultTeaCacheConfig() *TeaCacheConfig
DefaultTeaCacheConfig returns default configuration for TeaCache.
Source Files