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/__init__.py
View File
acestep/schedulers/scheduling_flow_match_euler_discrete.py
+14 -6
View File
@@ -72,7 +72,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
base_image_seq_len: Optional[int] = 256,
max_image_seq_len: Optional[int] = 4096,
):
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
@@ -146,7 +148,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
# self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
if self.begin_index is None:
step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
step_indices = [
self.index_for_timestep(t, schedule_timesteps) for t in timestep
]
elif self.step_index is not None:
# add_noise is called after first denoising step (for inpainting)
step_indices = [self.step_index] * timestep.shape[0]
@@ -186,12 +190,16 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
"""
if self.config.use_dynamic_shifting and mu is None:
raise ValueError(" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`")
raise ValueError(
" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
)
if sigmas is None:
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
@@ -243,7 +251,7 @@ class FlowMatchEulerDiscreteScheduler(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[FlowMatchEulerDiscreteSchedulerOutput, Tuple]:
"""
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
@@ -360,7 +368,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
# ## --
# ## channel mean 2
# dx = (sigma_next - sigma) * model_output
# dx = (sigma_next - sigma) * model_output
# m = dx.mean(dim=(2, 3), keepdim=True)
# # print(m.shape) # torch.Size([1, 16, 1, 1])
# dx_ = (dx - m) * omega + m
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