Skip to content

Commit f029e84

Browse files
committed
Slight performance improvement to FlowMatchHeun
1 parent e42fdb8 commit f029e84

File tree

1 file changed

+4
-5
lines changed

1 file changed

+4
-5
lines changed

src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -266,14 +266,13 @@ def step(
266266

267267
gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0
268268

269-
noise = randn_tensor(
270-
model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator
271-
)
272-
273-
eps = noise * s_noise
274269
sigma_hat = sigma * (gamma + 1)
275270

276271
if gamma > 0:
272+
noise = randn_tensor(
273+
model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator
274+
)
275+
eps = noise * s_noise
277276
sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5
278277

279278
if self.state_in_first_order:

0 commit comments

Comments
 (0)