From 687d650175789cfe7b087cef7a88e8991d7a6856 Mon Sep 17 00:00:00 2001 From: wangshuai6 Date: Fri, 11 Apr 2025 19:14:51 +0800 Subject: [PATCH] app demo --- src/diffusion/stateful_flow_matching/sharing_sampling.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/diffusion/stateful_flow_matching/sharing_sampling.py b/src/diffusion/stateful_flow_matching/sharing_sampling.py index f372028..9a73d2c 100644 --- a/src/diffusion/stateful_flow_matching/sharing_sampling.py +++ b/src/diffusion/stateful_flow_matching/sharing_sampling.py @@ -56,7 +56,7 @@ class EulerSampler(BaseSampler): logger.warning("current sampler is ODE sampler, but w_scheduler is enabled") # init recompute - self.num_recompute_timesteps = int(self.num_steps / self.state_refresh_rate) + self.recompute_timesteps = list(range(self.num_steps)) def sharing_dp(self, net, noise, condition, uncondition): @@ -143,6 +143,7 @@ class EulerSampler(BaseSampler): return x, pooled_state_list def __call__(self, net, noise, condition, uncondition): + self.num_recompute_timesteps = int(self.num_steps / self.state_refresh_rate) if len(self.recompute_timesteps) != self.num_recompute_timesteps: self.recompute_timesteps = self.sharing_dp(net, noise, condition, uncondition) denoised, _ = self._impl_sampling(net, noise, condition, uncondition)