File size: 2,490 Bytes
27fe41c 82fb56d 27fe41c 82fb56d 27fe41c 82fb56d 27fe41c 82fb56d 27fe41c 82fb56d 27fe41c 82fb56d 27fe41c 82fb56d 27fe41c 82fb56d 27fe41c 82fb56d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | from __future__ import annotations
import torch
import torch.nn.functional as F
class ResidualTimestepCache:
"""Training-free residual/velocity cache for multi-step flow inference.
Reuse is guarded by state cosine similarity, shape equality and a maximum timestep gap.
It is intended only for >=4-step students; 1-2 step students should bypass it entirely.
"""
def __init__(self,threshold=.985,max_reuse=2,max_timestep_gap=300.0):
self.threshold=threshold;self.max_reuse=max_reuse;self.max_timestep_gap=max_timestep_gap;self.reset()
def reset(self):self.prev_state=None;self.prev_output=None;self.prev_t=None;self.reuse_count=0
@staticmethod
def _feat(x):return x.feats if hasattr(x,'feats') else x
@staticmethod
def _t_scalar(t):
if t is None:return None
if torch.is_tensor(t):return float(t.detach().float().mean().cpu())
return float(t)
def similarity(self,x):
if self.prev_state is None:return -1.
a=self._feat(x).float();b=self._feat(self.prev_state).float()
if a.shape!=b.shape:return -1.
a=a.flatten(1);b=b.flatten(1);return float(F.cosine_similarity(a,b,dim=-1).mean())
def get(self,x,t=None):
sim=self.similarity(x);ts=self._t_scalar(t);time_ok=True
if self.prev_t is not None and ts is not None:time_ok=abs(ts-self.prev_t)<=self.max_timestep_gap
if time_ok and sim>=self.threshold and self.reuse_count<self.max_reuse:
self.reuse_count+=1;return self.prev_output,{'hit':True,'similarity':sim,'t_gap':None if ts is None or self.prev_t is None else abs(ts-self.prev_t)}
return None,{'hit':False,'similarity':sim,'t_gap':None if ts is None or self.prev_t is None else abs(ts-self.prev_t)}
def put(self,x,out,t=None):
self.prev_state=x;self.prev_output=out;self.prev_t=self._t_scalar(t);self.reuse_count=0
class CachedFlowModel:
def __init__(self,model,threshold=.985,max_reuse=2,max_timestep_gap=300.0):
self.model=model;self.cache=ResidualTimestepCache(threshold,max_reuse,max_timestep_gap);self.hits=0;self.calls=0
def __call__(self,x,xs,t,cond,*args,**kwargs):
self.calls+=1;cached,meta=self.cache.get(x,t)
if cached is not None:self.hits+=1;return cached
out=self.model(x,xs,t,cond,*args,**kwargs);self.cache.put(x,out,t);return out
@property
def hit_rate(self):return self.hits/max(1,self.calls)
def reset(self):self.cache.reset();self.hits=0;self.calls=0
|