Currently it's not possible to run ComfyUI with python -O because asserts only run when optimization is disabled and there are a couple places where ComfyUI actually depends on assert raising an exception for correct handling.

It's possible there are places I missed when I looking through the source code for similar structures but with these changes I can verify that ACE-Steps, Chroma Radiance and SDXL work when the Python interpreter is running with optimization enabled.

In the process of fixing that problem, I also noticed that BatchedBrownianTree doesn't correctly handle the case where w0 is passed in kwargs. I don't think this happens in practice anywhere but if w0 does get passed, it'll be included in the kwargs passed to torchsde.BrownianTree which doesn't support that argument. The simple fix is to just pop it with a default instead.

If necessary I can revert the changes that aren't strictly necessary: This pull also includes a bit of cleanup/refactoring for BatchedBrownianTree. The current implementation will move t0 and t1 to CPU (in cpu tree mode) for every tree which isn't necessary. It also always creates an empty w0 tensor whether it's needed or not. Because of the bug mentioned above, if it ever wasn't needed we would crash so not much real-world impact since apparently no one has ever tried passing w0 to the BatchedBrownianTree constructor.