search

replicate / cog-flux

Cog inference for flux models
https://replicate.com/black-forest-labs/flux-dev
Apache License 2.0
272 stars 28 forks source link

These are really slow #35

Closed doctorpangloss closed 1 week ago

doctorpangloss commented 1 week ago

https://github.com/replicate/cog-flux/blob/4c52f930e536bbe7d891c507d570667c803c0765/flux/modules/layers.py#L232

Replacing the einops rearrange with vanilla torch will improve speed.

technillogue commented 1 week ago

how do you know? in my testing torch.compile seems to unroll it into the equivalent vanilla operations and dropping einops did not make a different

doctorpangloss commented 1 week ago

that's good, I think this is more of an issue on Windows where torch.compile will not do that