app demo

app.py

@@ -70,7 +70,7 @@ class Pipeline:
         self.diffusion_sampler = diffusion_sampler
         self.resolution = resolution
         self.classlabels2ids = classlabels2ids
-
+    @spaces.GPU
     @torch.no_grad()
     @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
     def __call__(self, y, num_images, seed, num_steps, guidance, state_refresh_rate, guidance_interval_min, guidance_interval_max, timeshift):

src/diffusion/stateful_flow_matching/sharing_sampling.py

@@ -56,7 +56,6 @@ 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 +142,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)