I’m not quite sure what you are asking :). I assume it works. The code is a bit verbose/complex because you made the model configurable, how you want to do this is up to you.
When it comes to autoencoder, one has to be careful that the “unflatten” operation of the decoder is indeed the inverse of the “flatten” operation of the decoder. It’s not just about reconstructing the shape but also preserving the “semantic organization” of that data (cf. this example). Just from looking at your code, I cannot tell for sure. Most likely it’s fine.