Initial commit

config.json

@@ -1,10 +1,11 @@
 {
+  "_name_or_path": "blowing-up-groundhogs/emuru",
   "architectures": [
     "Emuru"
   ],
   "auto_map": {
-    "AutoConfig": "configuration_emuru.EmuruConfig",
-    "AutoModel": "modeling_emuru.Emuru"
+    "AutoConfig": "blowing-up-groundhogs/emuru--configuration_emuru.EmuruConfig",
+    "AutoModel": "blowing-up-groundhogs/emuru--modeling_emuru.Emuru"
   },
   "model_type": "emuru",
   "slices_per_query": 1,

model.safetensors

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

modeling_emuru.py

@@ -23,7 +23,7 @@ class Emuru(PreTrainedModel):
         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.vae_to_t5 = 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)
@@ -36,7 +36,6 @@ class Emuru(PreTrainedModel):
         # 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()
@@ -55,9 +54,39 @@ class Emuru(PreTrainedModel):
     #  - 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)
+    def forward(self, img=None, input_ids=None, attention_mask=None, noise=0, **kwargs):
+        decoder_inputs_embeds, z_sequence, z = self._img_encode(img, noise)
+
+        output = self.T5(input_ids, attention_mask=attention_mask, decoder_inputs_embeds=decoder_inputs_embeds)
+        vae_latent = self.t5_to_vae(output.logits[:, :-1])
+        pred_latent = self.z_rearrange(vae_latent)
+
+        mse_loss = self.mse_criterion(vae_latent, z_sequence)
+        return mse_loss, pred_latent, z
+        
+    def generate(self, text=None, img=None, max_length=128, noise=0):
+        # Your generate 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.
+
+    def _img_encode(self, img, noise=0):
+        posterior = self.vae.encode(img.float())
+        z = posterior.latent_dist.sample()
+        z_sequence = self.query_rearrange(z)
+
+        noise_sequence = z_sequence
+        if noise > 0:
+            noise_sequence = z_sequence + torch.randn_like(z_sequence) * noise
+
+        decoder_inputs_embeds = self.query_emb(noise_sequence)
+        sos = repeat(self.sos.weight, '1 d -> b 1 d', b=decoder_inputs_embeds.size(0))
+        decoder_inputs_embeds = torch.cat([sos, decoder_inputs_embeds], dim=1)
+        return decoder_inputs_embeds, z_sequence, z
+
+    def compute_padding_token(self):
+        # Your compute_padding_token implementation (port over from your original code)
+        ...
+
+    def compute_padding_token_threshold(self):
+        # Your compute_padding_token_threshold implementation (port over from your original code)
+        ...