Initial commit

.gitignore

@@ -1,3 +1,5 @@
 checkpoints
 test.py
-model.py
+model.py
+sample.png
+visual_prompting.py

modeling_emuru.py

@@ -2,7 +2,8 @@
 import torch
 import torch.nn as nn
 from transformers import PreTrainedModel, T5ForConditionalGeneration, T5Config, AutoTokenizer
-from .configuration_emuru import EmuruConfig
+from configuration_emuru import EmuruConfig
+# from .configuration_emuru import EmuruConfig
 from diffusers import AutoencoderKL
 from einops.layers.torch import Rearrange
 from einops import rearrange, repeat
@@ -43,6 +44,7 @@ class Emuru(PreTrainedModel):
         # 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():
@@ -53,7 +55,8 @@ class Emuru(PreTrainedModel):
     # 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, img=None, input_ids=None, attention_mask=None, noise=0, **kwargs):
         decoder_inputs_embeds, z_sequence, z = self._img_encode(img, noise)
 
@@ -63,11 +66,98 @@ class Emuru(PreTrainedModel):
 
         mse_loss = self.mse_criterion(vae_latent, z_sequence)
         return mse_loss, pred_latent, z
+
+
+    def old_generate(self, text=None, img=None, z_sequence=None, input_ids=None, max_new_tokens=256,
+                 stopping_criteria='latent', stopping_after=10, stopping_errors=1):
+        assert text is not None or input_ids is not None, 'Either text or input_ids must be provided'
+        assert img is not None or z_sequence is not None, 'Either img or z_sequence must be provided'
+
+        if input_ids is None:
+            input_ids = self.tokenizer(text, return_tensors='pt', padding=True).input_ids
+            input_ids = input_ids.to(next(self.T5.parameters()).device)
         
-    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.
-        ...
+        if z_sequence is None:
+            _, z_sequence, _ = self._img_encode(img)
+        z_sequence = [z_sequence]
+
+        sos = repeat(self.sos.weight, '1 d -> b 1 d', b=input_ids.size(0))
+        for _ in range(max_new_tokens):
+            if len(z_sequence) == 0:
+                decoder_inputs_embeds = sos
+            else:
+                decoder_inputs_embeds = self.vae_to_t5(torch.cat(z_sequence, dim=1))
+                decoder_inputs_embeds = torch.cat([sos, decoder_inputs_embeds], dim=1)
+            output = self.T5(input_ids, decoder_inputs_embeds=decoder_inputs_embeds)
+            vae_latent = self.t5_to_vae(output.logits[:, -1:])
+            z_sequence.append(vae_latent)
+
+            if stopping_criteria == 'latent':
+                curr_z_sequence = torch.cat(z_sequence, dim=1)
+                pad_token = repeat(self.padding_token, '1 d -> b 1 d', b=input_ids.size(0)).to(decoder_inputs_embeds.device)
+                similarity = torch.nn.functional.cosine_similarity(curr_z_sequence, pad_token, dim=-1)
+                similarity = similarity[:, -stopping_after:] > self.padding_token_threshold
+                if torch.all(similarity.sum(-1) >= (stopping_after - stopping_errors)):
+                    # z_sequence = [curr_z_sequence[:, :-stopping_after]]
+                    z_sequence = [curr_z_sequence]
+                    break
+            elif stopping_criteria == 'pixel':
+                raise NotImplementedError
+
+        z_sequence = torch.cat(z_sequence, dim=1)
+        img = torch.clamp(self.vae.decode(self.z_rearrange(z_sequence)).sample, -1, 1)
+        return img
+
+
+    def generate(self,
+                     style_text=None,
+                     gen_text=None,
+                     style_img=None,
+                     input_ids=None,
+                     z_sequence=None,
+                     max_new_tokens=256,
+                     stopping_criteria='latent',
+                     stopping_after=10,
+                     stopping_patience=1,
+                     trim_image=True):
+        assert (gen_text is not None and style_text is not None) or input_ids is not None, 'Either gen_text and style_text or input_ids must be provided'
+        assert style_img is not None or z_sequence is not None, 'Either style_img or z_sequence must be provided'
+
+        if input_ids is None:
+            input_ids = self.tokenizer(gen_text + ' ' + style_text, return_tensors='pt', padding=True).input_ids
+            input_ids = input_ids.to(self.device)
+
+        if z_sequence is None:
+            _, z_sequence, _ = self._img_encode(style_img)
+        z_sequence = [z_sequence]
+
+        sos = repeat(self.sos.weight, '1 d -> b 1 d', b=input_ids.size(0))
+        pad_token = repeat(self.padding_token, '1 d -> b 1 d', b=input_ids.size(0))
+
+        for _ in range(max_new_tokens):
+            if len(z_sequence) == 0:
+                decoder_inputs_embeds = sos
+            else:
+                decoder_inputs_embeds = self.vae_to_t5(torch.cat(z_sequence, dim=1))
+                decoder_inputs_embeds = torch.cat([sos, decoder_inputs_embeds], dim=1)
+            output = self.T5(input_ids, decoder_inputs_embeds=decoder_inputs_embeds)
+            vae_latent = self.t5_to_vae(output.logits[:, -1:])
+            z_sequence.append(vae_latent)
+
+            if stopping_criteria == 'latent':
+                curr_z_sequence = torch.cat(z_sequence, dim=1)
+                similarity = torch.nn.functional.cosine_similarity(curr_z_sequence, pad_token, dim=-1)
+                similarity = similarity[:, -stopping_after:] > self.padding_token_threshold
+                if torch.all(similarity.sum(-1) >= (stopping_after - stopping_patience)):
+                    z_sequence = [curr_z_sequence[:, :-similarity.sum(-1)]] if trim_image else [curr_z_sequence]
+                    break
+            elif stopping_criteria == 'pixel':
+                raise NotImplementedError
+
+        z_sequence = torch.cat(z_sequence, dim=1)
+        img = torch.clamp(self.vae.decode(self.z_rearrange(z_sequence)).sample, -1, 1)
+        return img, z_sequence
+    
 
     def _img_encode(self, img, noise=0):
         posterior = self.vae.encode(img.float())
@@ -78,15 +168,15 @@ class Emuru(PreTrainedModel):
         if noise > 0:
             noise_sequence = z_sequence + torch.randn_like(z_sequence) * noise
 
-        decoder_inputs_embeds = self.query_emb(noise_sequence)
+        decoder_inputs_embeds = self.vae_to_t5(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)
-        ...
+        raise NotImplementedError("compute_padding_token not implemented")
+
 
     def compute_padding_token_threshold(self):
-        # Your compute_padding_token_threshold implementation (port over from your original code)
-        ...
+        raise NotImplementedError("compute_padding_token_threshold not implemented")