From 2b80e37555a1d3ca72f4563c0201e27fe981fae6 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sat, 12 Oct 2024 10:45:01 +0300 Subject: [PATCH] handle both tensor and numpy input in init_sigmas_for_each_stage from upstream: https://github.com/jy0205/Pyramid-Flow/commit/bf10e7340cd1e3467c7b88053bc0f413eb1a6073#diff-8e1c08bf8ab5d587ee253dccfb631b170202adf0664272785a9b377975ffb8a5R145 --- diffusion_schedulers/scheduling_flow_matching.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/diffusion_schedulers/scheduling_flow_matching.py b/diffusion_schedulers/scheduling_flow_matching.py index 3345304..b8c1507 100644 --- a/diffusion_schedulers/scheduling_flow_matching.py +++ b/diffusion_schedulers/scheduling_flow_matching.py @@ -142,7 +142,8 @@ class PyramidFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): timesteps = np.linspace( timestep_max, timestep_min, training_steps + 1, ) - self.timesteps_per_stage[i_s] = torch.from_numpy(timesteps[:-1]) + self.timesteps_per_stage[i_s] = timesteps[:-1] if isinstance(timesteps, torch.Tensor) else torch.from_numpy(timesteps[:-1]) + stage_sigmas = np.linspace( 1, 0, training_steps + 1, )