Merge pull request #1 from ahmadazim/VAE

Working DiffuGene implementation (needs to be reorganized/cleaned) with UNet backbone.

New homogenized AE structure

VQ-VAE strategy implemented.

Implemented chromosome-specific VQ-VAE models (TODO: integrate better with pipeline).

Changing to a VQVAE strategy.

decoding scripts tweaks.

minor big fix

Fixed generation bugs (imputation, fam matching, etc). Added decoding functionality for generated samples.

Pipeline updates to generation, encoding, and decoders.

Generation script updates, added several visualizations + Notebook updates.

testing new VAE strategy

VAE architecture tweaking. Added per-PC MSE logging during training, richer pos embed, MHSA at bottleneck, MBConv style blocks.

VAE architecture and training adjustments: 1. norm + activation swaps, 2. FiLM identity init, 3. updated LR schedule, 4. residual 1D blocks, 5. KL cyclical schedule + free-bits.

Added CFG, tweaks to treaining scripts.