This takes me back to 2019 when me and my team won HackInOut, a hackathon held in Bangalore, India by building Rosaline, a generative drum machine, that took in lead melodies (guitar tracks, piano samples, etc) and generated a full fledged drum track.
Used GANs & a lot of weird hacky optimizations, can't even find the code for it anymore, but it was a great time :)
https://devfolio.co/projects/rosaline
This is all that's left of it
Looks cool! It's a much more difficult problem than it looks - great that you got something working.