• Thiago Crepaldi's avatar
    Enable PyTorch's FakeTensorMode for EulerDiscreteScheduler scheduler (#7151) · ca6cdc77
    Thiago Crepaldi authored
    * 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
    ca6cdc77
scheduling_euler_discrete.py 24.7 KB