rkstgr / jukebox_diffusion Goto Github PK

Motivation:
Music continuation (equivalent to outpainting for images) is not trivially possible with normal unconditional diffusion models. Workarounds like RePaint exist but require 10x more samples for generation.

Idea:
Instead of diffusing the whole sequence (music) to the same diffusion step, we generalise the idea to every temporal step. Thus, there can be regions that have no noise applied and should be learned to contain context information by the model

Implementation:
In the unconditional setting the model gets the sequence and diffusion step t like model(x, t), where t is (B, 1), such that every batch is diffused equally. Now t should have shape (B, S, 1), such that every batch B and sequence pos has its own diffusion step. The implementation should be backwards compatible, such that t with (B,1) is treated as before.