Initial commit

.gitignore

@@ -0,0 +1,3 @@
+checkpoints
+test.py
+model.py

config.json

@@ -0,0 +1,13 @@
+{
+  "architectures": [
+    "Emuru"
+  ],
+  "model_type": "emuru",
+  "slices_per_query": 1,
+  "t5_config": "google-t5/t5-large",
+  "tokenizer_config": "google/byt5-small",
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.2",
+  "vae_channels": 1,
+  "vae_config": "blowing-up-groundhogs/emuru_vae"
+}

configuration_emuru.py

@@ -0,0 +1,18 @@
+from transformers import PretrainedConfig
+
+class EmuruConfig(PretrainedConfig):
+    model_type = "emuru"  # Unique identifier for your model
+
+    def __init__(self, 
+                 t5_config='google-t5/t5-large', 
+                 vae_config='blowing-up-groundhogs/emuru_vae',
+                 tokenizer_config='google/byt5-small',
+                 slices_per_query=1,
+                 vae_channels=1,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.t5_config = t5_config
+        self.vae_config = vae_config
+        self.tokenizer_config = tokenizer_config
+        self.slices_per_query = slices_per_query
+        self.vae_channels = vae_channels

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b1c1d034a3c8b158d4f656bdccf4498816c8499300da30b764fa9f290f33f7f
+size 2876698952

modeling_emuru.py

@@ -0,0 +1,63 @@
+# modeling_emuru.py
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, T5ForConditionalGeneration, T5Config, AutoTokenizer
+from configuration_emuru import EmuruConfig
+from diffusers import AutoencoderKL
+from einops.layers.torch import Rearrange
+from einops import rearrange, repeat
+
+class Emuru(PreTrainedModel):
+    config_class = EmuruConfig  # Link to your configuration
+
+    def __init__(self, config):
+        super().__init__(config)
+        # Initialize the tokenizer (if you want it as part of your model)
+        self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_config)
+
+        # Load T5 using the provided filename from config
+        t5_config = T5Config.from_pretrained(config.t5_config)
+        t5_config.vocab_size = len(self.tokenizer)
+        self.T5 = T5ForConditionalGeneration(t5_config)
+        self.T5.lm_head = nn.Identity()
+        self.sos = nn.Embedding(1, t5_config.d_model)
+
+        vae_latent_size = 8 * config.vae_channels * config.slices_per_query
+        self.query_emb = nn.Linear(vae_latent_size, t5_config.d_model)
+        self.t5_to_vae = nn.Linear(t5_config.d_model, vae_latent_size, bias=False)
+
+        self.padding_token = nn.Parameter(torch.empty((1, vae_latent_size)), requires_grad=False)
+        self.padding_token_threshold = nn.Parameter(torch.empty(1), requires_grad=False)
+
+        # Load VAE
+        self.vae = AutoencoderKL.from_pretrained(config.vae_config)
+        self.set_training(self.vae, False)
+
+        # Define the rearrange layers
+        self.query_rearrange = Rearrange('b c h (w q) -> b w (q c h)', q=config.slices_per_query)
+        self.z_rearrange = Rearrange('b w (q c h) -> b c h (w q)', c=config.vae_channels, q=config.slices_per_query)
+        self.z_rearrange_eval = Rearrange('w b (q c h) -> b c h (w q)', c=config.vae_channels, q=config.slices_per_query)
+
+        # Define your loss functions
+        self.mse_criterion = nn.MSELoss()
+
+        # Initialize weights following Hugging Face conventions (if needed)
+        self.init_weights()
+
+    def set_training(self, model, training):
+        model.train() if training else model.eval()
+        for param in model.parameters():
+            param.requires_grad = training
+
+    # --- Implement the rest of your methods ---
+    # For example, _img_encode, forward, generate, etc.
+    # You can largely port your existing code here, making sure that:
+    #  - The forward method returns a dictionary with your losses and outputs.
+    #  - You use the Hugging Face methods for saving/loading weights.
+    
+    def forward(self, text=None, img=None, input_ids=None, attention_mask=None, length=None, noise=0):
+        # Your forward implementation (port over from your original code)
+        # Make sure to call self._img_encode(img, noise) and use self.T5, etc.
+        ...
+        
+    # Add other methods (forward_recurrent, generate, etc.) as needed.