omerbt / TokenFlow

Official Pytorch Implementation for "TokenFlow: Consistent Diffusion Features for Consistent Video Editing" presenting "TokenFlow" (ICLR 2024)
https://diffusion-tokenflow.github.io
MIT License
1.52k stars 134 forks source link

batching pivots allows processing bigger/longer sequences #32

Open eps696 opened 9 months ago

eps696 commented 9 months ago

for those who (like me) wanted to apply this exciting technique for longer videos: i've integrated this method into my SD repo https://github.com/eps696/SDfu and added there batches for pivots with offloading them onto CPU. this allowed to process e.g. 300 frames in 960x540 res on 3090 (24gb). as i renamed some variables to my convenience, my code is not directly copypastable into this repo, yet i hope it's readable enough to apply here. the solution is also pretty clumsy, as i had very little idea about that attention stuff and just tried to debug OOMs..