I am working on something similar on one of my projects. I use grid for that instead of flex. Also grid-template-rows: masonry, which does not yet have a good support but I added few fallbacks just in case. I have not put much time into this, but here I made a project so you can take a look.
.jpeg/public)