| 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 |
|
|
