mrleo1nid/ACE-Step
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

work on pip package

Browse Source
This commit is contained in:
mrfakename
2025-05-06 18:59:32 -07:00
parent 54da683d36
commit a5746eaab6
36 changed files with 2928 additions and 1251 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acestep/schedulers/scheduling_flow_match_heun_discrete.py
+23 -8
View File
@@ -67,7 +67,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
num_train_timesteps: int = 1000,
shift: float = 1.0,
):
timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
timesteps = np.linspace(
1, num_train_timesteps, num_train_timesteps, dtype=np.float32
)[::-1].copy()
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
sigmas = timesteps / num_train_timesteps
@@ -137,7 +139,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
def _sigma_to_t(self, sigma):
return sigma * self.config.num_train_timesteps
def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
def set_timesteps(
self, num_inference_steps: int, device: Union[str, torch.device] = None
):
"""
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -150,7 +154,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
self.num_inference_steps = num_inference_steps
timesteps = np.linspace(
self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
self._sigma_to_t(self.sigma_max),
self._sigma_to_t(self.sigma_min),
num_inference_steps,
)
sigmas = timesteps / self.config.num_train_timesteps
@@ -162,7 +168,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
self.timesteps = timesteps.to(device=device)
sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]])
self.sigmas = torch.cat(
[sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]]
)
# empty dt and derivative
self.prev_derivative = None
@@ -208,7 +216,7 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
s_noise: float = 1.0,
generator: Optional[torch.Generator] = None,
return_dict: bool = True,
omega: Union[float, np.array] = 0.0
omega: Union[float, np.array] = 0.0,
) -> Union[FlowMatchHeunDiscreteSchedulerOutput, Tuple]:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
@@ -257,7 +265,7 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
self.omega_bef_rescale = omega
omega = logistic_function(omega, k=0.1)
self.omega_aft_rescale = omega
if (
isinstance(timestep, int)
or isinstance(timestep, torch.IntTensor)
@@ -285,13 +293,20 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
sigma = self.sigmas[self.step_index - 1]
sigma_next = self.sigmas[self.step_index]
gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0
gamma = (
min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1)
if s_tmin <= sigma <= s_tmax
else 0.0
)
sigma_hat = sigma * (gamma + 1)
if gamma > 0:
noise = randn_tensor(
model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator
model_output.shape,
dtype=model_output.dtype,
device=model_output.device,
generator=generator,
)
eps = noise * s_noise
sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5