Adding jax maskgit weights to muse

[isamu-isozaki]

@patil-suraj Hi! I cleaned up the script a bit and made this pr. I need to double check everything works but this is the main idea I have for now

[isamu-isozaki]

My plan is

• Confirm the maskgit weights are properly loaded/compatible with jax
• Test that inference script works properly
• Test full pipeline

And I should be done.

[isamu-isozaki]

One note, I updated the huggingface_hub version since I noticed that on 0.4.0, some of the functions fail to import

[isamu-isozaki]

Let me fix one layer norm in a bit

[isamu-isozaki]

fixed!