Skip to content

Commit ca6cdc7

Browse files
author
Thiago Crepaldi
authored
Enable PyTorch's FakeTensorMode for EulerDiscreteScheduler scheduler (huggingface#7151)
* Enable FakeTensorMode for EulerDiscreteScheduler scheduler PyTorch's FakeTensorMode does not support `.numpy()` or `numpy.array()` calls. This PR replaces `sigmas` numpy tensor by a PyTorch tensor equivalent Repro ```python with torch._subclasses.FakeTensorMode() as fake_mode, ONNXTorchPatcher(): fake_model = DiffusionPipeline.from_pretrained(model_name, low_cpu_mem_usage=False) ``` that otherwise would fail with `RuntimeError: .numpy() is not supported for tensor subclasses.` * Address comments
1 parent f4977ab commit ca6cdc7

File tree

1 file changed

+1
-3
lines changed

1 file changed

+1
-3
lines changed

src/diffusers/schedulers/scheduling_euler_discrete.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -216,10 +216,8 @@ def __init__(
216216
# FP16 smallest positive subnormal works well here
217217
self.alphas_cumprod[-1] = 2**-24
218218

219-
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
219+
sigmas = (((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5).flip(0)
220220
timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=float)[::-1].copy()
221-
222-
sigmas = torch.from_numpy(sigmas[::-1].copy()).to(dtype=torch.float32)
223221
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
224222

225223
# setable values

0 commit comments

Comments
 (0)