WIP: Second stage training

[isamu-isozaki]

This is based on this issue: #100
Here are the tasks:

• In train_muse.py's prepare_inputs_and_labels function, interpolate pixel values to 256x256 and get tokens using f16 vqgan for low resolutions and 512x512 and f8 for high resolutions. We can use precomputed embeddings here
• the TransformerLayers for low resolution. For this, I'm planning to be an adapter with a MaskGitTransformer or any transformer but without the encoder_hidden_states.
• the MaxVitTransformerLayers for high resolution(in different pr)
• and projection layer and concatenating layer between the low res+text embeddings
• Test out on imagenet

[isamu-isozaki]

I think for the basic idea I'm roughly done(except projection/concatenating). Will test it out in a bit

[isamu-isozaki]

I'm testing on jewels on branch 2nd_stage_testing https://wandb.ai/dsbuddy27/muse/runs/34i7fqg4?workspace=user-isamu. It seems to be able to run however I think I'm messing up the generation for the pre-trained low-res transformer. So I will fix that, test with adafactor/increasing batch size and the try scaling up to laion a

[isamu-isozaki]

Ok! The Generation part is fixed! I'll try running overnight on imagenet and see how it does
https://wandb.ai/dsbuddy27/muse/runs/v0z3bvd7?workspace=user-isamu

[isamu-isozaki]

Ok. I'm having trouble getting the sbatch job running on jewels so will try moving to the hpc cluster. In the mean time, will update this branch so it's up to date!