Unverified Commit ca6cdc77 authored by Thiago Crepaldi's avatar Thiago Crepaldi Committed by GitHub
Browse files

Enable PyTorch's FakeTensorMode for EulerDiscreteScheduler scheduler (#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
parent f4977abc
......@@ -216,10 +216,8 @@ class EulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
# FP16 smallest positive subnormal works well here
self.alphas_cumprod[-1] = 2**-24
sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
sigmas = (((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5).flip(0)
timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=float)[::-1].copy()
sigmas = torch.from_numpy(sigmas[::-1].copy()).to(dtype=torch.float32)
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
# setable values
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment