Logic/DDT
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix bugs(admas timedeltas)

Browse Source
This commit is contained in:
wangshuai6
2025-05-20 12:30:40 +08:00
parent 3693640ca3
commit 99d92c94e7
2 changed files with 20 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/diffusion/flow_matching/adam_sampling.py
View File

@@ -41,23 +41,25 @@ class AdamLMSampler(BaseSampler):
self, self,
order: int = 2, order: int = 2,
timeshift: float = 1.0, timeshift: float = 1.0,
guidance_interval_min: float = 0.0,
guidance_interval_max: float = 1.0,
lms_transform_fn: Callable = nop, lms_transform_fn: Callable = nop,
w_scheduler: BaseScheduler = None,
step_fn: Callable = ode_step_fn, step_fn: Callable = ode_step_fn,
*args, *args,
**kwargs **kwargs
): ):
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
self.step_fn = step_fn self.step_fn = step_fn
self.w_scheduler = w_scheduler
assert self.scheduler is not None assert self.scheduler is not None
assert self.w_scheduler is not None or self.step_fn in [ode_step_fn, ] assert self.step_fn in [ode_step_fn, ]
self.order = order self.order = order
self.lms_transform_fn = lms_transform_fn self.lms_transform_fn = lms_transform_fn
timesteps = torch.linspace(0.0, 1 - self.last_step, self.num_steps) timesteps = torch.linspace(0.0, 1 - self.last_step, self.num_steps)
timesteps = torch.cat([timesteps, torch.tensor([1.0])], dim=0) timesteps = torch.cat([timesteps, torch.tensor([1.0])], dim=0)
self.guidance_interval_min = guidance_interval_min
self.guidance_interval_max = guidance_interval_max
self.timesteps = shift_respace_fn(timesteps, timeshift) self.timesteps = shift_respace_fn(timesteps, timeshift)
self.timedeltas = self.timesteps[1:] - self.timesteps[:-1] self.timedeltas = self.timesteps[1:] - self.timesteps[:-1]
self._reparameterize_coeffs() self._reparameterize_coeffs()
@@ -93,7 +95,11 @@ class AdamLMSampler(BaseSampler):
cfg_x = torch.cat([x, x], dim=0) cfg_x = torch.cat([x, x], dim=0)
cfg_t = t_cur.repeat(2) cfg_t = t_cur.repeat(2)
out = net(cfg_x, cfg_t, cfg_condition) out = net(cfg_x, cfg_t, cfg_condition)
out = self.guidance_fn(out, self.guidance) if t_cur[0] > self.guidance_interval_min and t_cur[0] < self.guidance_interval_max:
guidance = self.guidance
out = self.guidance_fn(out, guidance)
else:
out = self.guidance_fn(out, 1.0)
pred_trajectory.append(out) pred_trajectory.append(out)
out = torch.zeros_like(out) out = torch.zeros_like(out)
order = len(self.solver_coeffs[i]) order = len(self.solver_coeffs[i])

14
src/diffusion/stateful_flow_matching/adam_sampling.py
View File

@@ -41,22 +41,24 @@ class AdamLMSampler(BaseSampler):
self, self,
order: int = 2, order: int = 2,
timeshift: float = 1.0, timeshift: float = 1.0,
guidance_interval_min: float = 0.0,
guidance_interval_max: float = 1.0,
state_refresh_rate: int = 1, state_refresh_rate: int = 1,
lms_transform_fn: Callable = nop, lms_transform_fn: Callable = nop,
w_scheduler: BaseScheduler = None,
step_fn: Callable = ode_step_fn, step_fn: Callable = ode_step_fn,
*args, *args,
**kwargs **kwargs
): ):
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
self.step_fn = step_fn self.step_fn = step_fn
self.w_scheduler = w_scheduler
self.state_refresh_rate = state_refresh_rate self.state_refresh_rate = state_refresh_rate
assert self.scheduler is not None assert self.scheduler is not None
assert self.w_scheduler is not None or self.step_fn in [ode_step_fn, ] assert self.step_fn in [ode_step_fn, ]
self.order = order self.order = order
self.lms_transform_fn = lms_transform_fn self.lms_transform_fn = lms_transform_fn
self.guidance_interval_min = guidance_interval_min
self.guidance_interval_max = guidance_interval_max
timesteps = torch.linspace(0.0, 1 - self.last_step, self.num_steps) timesteps = torch.linspace(0.0, 1 - self.last_step, self.num_steps)
timesteps = torch.cat([timesteps, torch.tensor([1.0])], dim=0) timesteps = torch.cat([timesteps, torch.tensor([1.0])], dim=0)
@@ -98,7 +100,11 @@ class AdamLMSampler(BaseSampler):
if i % self.state_refresh_rate == 0: if i % self.state_refresh_rate == 0:
state = None state = None
out, state = net(cfg_x, cfg_t, cfg_condition, state) out, state = net(cfg_x, cfg_t, cfg_condition, state)
out = self.guidance_fn(out, self.guidance) if t_cur[0] > self.guidance_interval_min and t_cur[0] < self.guidance_interval_max:
guidance = self.guidance
out = self.guidance_fn(out, guidance)
else:
out = self.guidance_fn(out, 1.0)
pred_trajectory.append(out) pred_trajectory.append(out)
out = torch.zeros_like(out) out = torch.zeros_like(out)
order = len(self.solver_coeffs[i]) order = len(self.solver_coeffs[i])