Test previous version

modeling_emuru.bkp.py

@@ -0,0 +1,316 @@
+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 repeat
+from torchvision.transforms import functional as F
+from typing import Optional, Tuple, List, Any
+from PIL import Image
+
+class Emuru(PreTrainedModel):
+    config_class = EmuruConfig
+    """
+    Emuru is a conditional generative model that integrates a T5-based decoder with a VAE
+    for image generation conditioned on text and style images.
+
+    Attributes:
+        config_class (Type): Configuration class for the model.
+        tokenizer (AutoTokenizer): Tokenizer loaded from the provided tokenizer configuration.
+        T5 (T5ForConditionalGeneration): T5 model adapted for conditional generation.
+        sos (nn.Embedding): Start-of-sequence embedding.
+        vae_to_t5 (nn.Linear): Linear projection from VAE latent space to T5 hidden space.
+        t5_to_vae (nn.Linear): Linear projection from T5 hidden space back to VAE latent space.
+        padding_token (nn.Parameter): Non-trainable parameter for padding tokens.
+        padding_token_threshold (nn.Parameter): Non-trainable parameter for padding token threshold.
+        vae (AutoencoderKL): Pre-trained Variational Autoencoder.
+        query_rearrange (Rearrange): Layer to rearrange VAE latent representations for queries.
+        z_rearrange (Rearrange): Layer to rearrange T5 outputs back to VAE latent dimensions.
+        mse_criterion (nn.MSELoss): Mean squared error loss function.
+    """
+
+    def __init__(self, config: EmuruConfig) -> None:
+        """
+        Initialize the Emuru model.
+
+        Args:
+            config (EmuruConfig): Configuration object containing model hyperparameters and paths.
+        """
+        super().__init__(config)
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_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.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)
+        self.padding_token_threshold = nn.Parameter(torch.empty(1), requires_grad=False)
+
+        self.vae = AutoencoderKL.from_pretrained(config.vae_config)
+        self.set_training(self.vae, False)
+
+        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.mse_criterion = nn.MSELoss()
+        self.init_weights()
+
+    def set_training(self, model: nn.Module, training: bool) -> None:
+        """
+        Set the training mode for a given model and freeze/unfreeze parameters accordingly.
+
+        Args:
+            model (nn.Module): The model to set the training mode for.
+            training (bool): If True, set the model to training mode; otherwise, evaluation mode.
+        """
+        model.train() if training else model.eval()
+        for param in model.parameters():
+            param.requires_grad = training
+
+    def forward(
+        self,
+        img: Optional[torch.Tensor] = None,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        noise: float = 0,
+        **kwargs: Any
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Forward pass of the model.
+
+        Args:
+            img (Optional[torch.Tensor]): Input image tensor.
+            input_ids (Optional[torch.Tensor]): Tokenized input IDs.
+            attention_mask (Optional[torch.Tensor]): Attention mask for the inputs.
+            noise (float): Amount of noise to add in image encoding.
+            **kwargs: Additional arguments.
+
+        Returns:
+            Tuple containing:
+                - mse_loss (torch.Tensor): Mean squared error loss.
+                - pred_latent (torch.Tensor): Predicted latent representations.
+                - z (torch.Tensor): Sampled latent vector from VAE.
+        """
+        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,
+        style_text: str,
+        gen_text: str,
+        style_img: torch.Tensor,
+        **kwargs: Any
+    ) -> Image.Image:
+        """
+        Generate an image by combining style and generation texts with a style image.
+
+        Args:
+            style_text (str): Style-related text prompt.
+            gen_text (str): Generation-related text prompt.
+            style_img (torch.Tensor): Style image tensor. Expected shape is either 3D or 4D.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            Image.Image: Generated image as a PIL image.
+        """
+        if style_img.ndim == 3:
+            style_img = style_img.unsqueeze(0)
+        elif style_img.ndim == 4:
+            pass
+        else:
+            raise ValueError('style_img must be 3D or 4D')
+        
+        texts = [style_text + ' ' + gen_text]
+        imgs, _, img_ends = self._generate(texts=texts, imgs=style_img, **kwargs)
+        imgs = (imgs + 1) / 2
+        return F.to_pil_image(imgs[0, ..., style_img.size(-1):img_ends.item()].detach().cpu())
+    
+    def generate_batch(
+        self,
+        style_texts: List[str],
+        gen_texts: List[str],
+        style_imgs: torch.Tensor,
+        lengths: List[int],
+        **kwargs: Any
+    ) -> List[Image.Image]:
+        """
+        Generate a batch of images from lists of style texts, generation texts, and style images.
+
+        Args:
+            style_texts (List[str]): List of style-related text prompts.
+            gen_texts (List[str]): List of generation-related text prompts.
+            style_imgs (torch.Tensor): Batch of style images (4D tensor).
+            lengths (List[int]): List of lengths corresponding to each image.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            List[Image.Image]: List of generated images as PIL images.
+        """
+        assert style_imgs.ndim == 4, 'style_imgs must be 4D'
+        assert len(style_texts) == len(style_imgs), 'style_texts and style_imgs must have the same length'
+        assert len(gen_texts) == len(style_imgs), 'gen_texts and style_imgs must have the same length'
+        texts = [style_text + ' ' + gen_text for style_text, gen_text in zip(style_texts, gen_texts)]
+        
+        imgs, _, img_ends = self._generate(texts=texts, imgs=style_imgs, lengths=lengths, **kwargs)
+        imgs = (imgs + 1) / 2
+
+        out_imgs = []
+        for i, end in enumerate(img_ends):
+            start = lengths[i]
+            out_imgs.append(F.to_pil_image(imgs[i, ..., start:end].detach().cpu()))
+        return out_imgs
+
+    def _generate(
+        self,
+        texts: Optional[List[str]] = None,
+        imgs: Optional[torch.Tensor] = None,
+        lengths: Optional[List[int]] = None,
+        input_ids: Optional[torch.Tensor] = None,
+        z_sequence: Optional[torch.Tensor] = None,
+        max_new_tokens: int = 256,
+        stopping_criteria: str = 'latent',
+        stopping_after: int = 10,
+        stopping_patience: int = 1
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Internal generation routine that combines textual and visual inputs to iteratively generate
+        latent representations and decode them into images.
+
+        Args:
+            texts (Optional[List[str]]): List of text prompts.
+            imgs (Optional[torch.Tensor]): Input image tensor.
+            lengths (Optional[List[int]]): Desired lengths for each image in latent space.
+            input_ids (Optional[torch.Tensor]): Tokenized input IDs.
+            z_sequence (Optional[torch.Tensor]): Precomputed latent sequence.
+            max_new_tokens (int): Maximum tokens to generate.
+            stopping_criteria (str): Criteria for stopping ('latent' or 'none').
+            stopping_after (int): Number of tokens to check for stopping condition.
+            stopping_patience (int): Patience parameter for stopping condition.
+
+        Returns:
+            Tuple containing:
+                - imgs (torch.Tensor): Generated images.
+                - canvas_sequence (torch.Tensor): Generated latent canvas sequence.
+                - img_ends (torch.Tensor): End indices for each generated image.
+        """
+        assert texts is not None or input_ids is not None, 'Either texts or input_ids must be provided'
+        assert imgs is not None or z_sequence is not None, 'Either imgs or z_sequence must be provided'
+
+        if input_ids is None:
+            input_ids = self.tokenizer(texts, return_tensors='pt', padding=True).input_ids
+            input_ids = input_ids.to(self.device)
+
+        if z_sequence is None:
+            _, z_sequence, _ = self._img_encode(imgs)
+        
+        if lengths is None:
+            lengths = [imgs.size(-1)] * imgs.size(0)
+        lengths = torch.tensor(lengths).to(self.device)
+        lengths = (lengths / 8).ceil().int()
+
+        z_sequence_mask = torch.zeros((z_sequence.size(0), lengths.max() + max_new_tokens))
+        z_sequence_mask = z_sequence_mask.bool().to(self.device)
+        for i, l in enumerate(lengths):
+            z_sequence_mask[i, :l] = True
+
+        canvas_sequence = z_sequence[:, :lengths.min()]
+        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))
+        seq_stops = torch.ones(z_sequence.size(0), dtype=torch.int) * -1
+
+        for token_idx in range(lengths.min(), lengths.max() + max_new_tokens):
+            if len(z_sequence) == 0:
+                decoder_inputs_embeds = sos
+            else:
+                decoder_inputs_embeds = self.vae_to_t5(canvas_sequence)
+                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:])
+
+            mask_slice = z_sequence_mask[:, token_idx].unsqueeze(-1)
+            if token_idx < z_sequence.size(1):
+                seq_slice = torch.where(mask_slice, z_sequence[:, token_idx], vae_latent[:, 0])
+            else:
+                seq_slice = vae_latent[:, 0]
+            canvas_sequence = torch.cat([canvas_sequence, seq_slice.unsqueeze(1)], dim=1)
+
+            if stopping_criteria == 'latent':
+                similarity = torch.nn.functional.cosine_similarity(canvas_sequence, pad_token, dim=-1)
+                windows = (similarity > self.padding_token_threshold).unfold(1, stopping_after, 1)
+                window_sums = windows.to(torch.int).sum(dim=2)
+
+                for i in range(similarity.size(0)):
+                    idx = (window_sums[i] > (stopping_after - stopping_patience)).nonzero(as_tuple=True)[0]
+                    if idx.numel() > 0:
+                        seq_stops[i] = idx[0].item()
+
+                if torch.all(seq_stops >= 0):
+                    break
+            elif stopping_criteria == 'none':
+                pass
+
+        imgs = torch.clamp(self.vae.decode(self.z_rearrange(canvas_sequence)).sample, -1, 1)
+        return imgs, canvas_sequence, seq_stops * 8
+    
+    def _img_encode(
+        self,
+        img: torch.Tensor,
+        noise: float = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Encode the input image into a latent representation using the VAE.
+
+        Args:
+            img (torch.Tensor): Input image tensor.
+            noise (float): Standard deviation of noise to add to the latent sequence.
+
+        Returns:
+            Tuple containing:
+                - decoder_inputs_embeds (torch.Tensor): Embeddings to be used as T5 decoder inputs.
+                - z_sequence (torch.Tensor): Rearranged latent sequence from the VAE.
+                - z (torch.Tensor): Sampled latent vector from the VAE.
+        """
+        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.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) -> None:
+        """
+        Compute and update the padding token.
+
+        Raises:
+            NotImplementedError: This method must be implemented.
+        """
+        raise NotImplementedError("compute_padding_token not implemented")
+
+    def compute_padding_token_threshold(self) -> None:
+        """
+        Compute and update the padding token threshold.
+
+        Raises:
+            NotImplementedError: This method must be implemented.
+        """
+        raise NotImplementedError("compute_padding_token_threshold not implemented")

modeling_emuru.py

@@ -1,46 +1,22 @@
+# modeling_emuru.py
 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 diffusers import AutoencoderKL
 from einops.layers.torch import Rearrange
-from einops import repeat
-from torchvision.transforms import functional as F
-from typing import Optional, Tuple, List, Any
-from PIL import Image
+from einops import rearrange, repeat
 
 class Emuru(PreTrainedModel):
-    config_class = EmuruConfig
-    """
-    Emuru is a conditional generative model that integrates a T5-based decoder with a VAE
-    for image generation conditioned on text and style images.
-
-    Attributes:
-        config_class (Type): Configuration class for the model.
-        tokenizer (AutoTokenizer): Tokenizer loaded from the provided tokenizer configuration.
-        T5 (T5ForConditionalGeneration): T5 model adapted for conditional generation.
-        sos (nn.Embedding): Start-of-sequence embedding.
-        vae_to_t5 (nn.Linear): Linear projection from VAE latent space to T5 hidden space.
-        t5_to_vae (nn.Linear): Linear projection from T5 hidden space back to VAE latent space.
-        padding_token (nn.Parameter): Non-trainable parameter for padding tokens.
-        padding_token_threshold (nn.Parameter): Non-trainable parameter for padding token threshold.
-        vae (AutoencoderKL): Pre-trained Variational Autoencoder.
-        query_rearrange (Rearrange): Layer to rearrange VAE latent representations for queries.
-        z_rearrange (Rearrange): Layer to rearrange T5 outputs back to VAE latent dimensions.
-        mse_criterion (nn.MSELoss): Mean squared error loss function.
-    """
-
-    def __init__(self, config: EmuruConfig) -> None:
-        """
-        Initialize the Emuru model.
-
-        Args:
-            config (EmuruConfig): Configuration object containing model hyperparameters and paths.
-        """
+    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)
@@ -54,51 +30,34 @@ class Emuru(PreTrainedModel):
         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)
 
+        # 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: nn.Module, training: bool) -> None:
-        """
-        Set the training mode for a given model and freeze/unfreeze parameters accordingly.
 
-        Args:
-            model (nn.Module): The model to set the training mode for.
-            training (bool): If True, set the model to training mode; otherwise, evaluation mode.
-        """
+    def set_training(self, model, training):
         model.train() if training else model.eval()
         for param in model.parameters():
             param.requires_grad = training
 
-    def forward(
-        self,
-        img: Optional[torch.Tensor] = None,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        noise: float = 0,
-        **kwargs: Any
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Forward pass of the model.
-
-        Args:
-            img (Optional[torch.Tensor]): Input image tensor.
-            input_ids (Optional[torch.Tensor]): Tokenized input IDs.
-            attention_mask (Optional[torch.Tensor]): Attention mask for the inputs.
-            noise (float): Amount of noise to add in image encoding.
-            **kwargs: Additional arguments.
-
-        Returns:
-            Tuple containing:
-                - mse_loss (torch.Tensor): Mean squared error loss.
-                - pred_latent (torch.Tensor): Predicted latent representations.
-                - z (torch.Tensor): Sampled latent vector from VAE.
-        """
+    # --- 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, 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)
@@ -108,182 +67,99 @@ class Emuru(PreTrainedModel):
         mse_loss = self.mse_criterion(vae_latent, z_sequence)
         return mse_loss, pred_latent, z
 
-    def generate(
-        self,
-        style_text: str,
-        gen_text: str,
-        style_img: torch.Tensor,
-        **kwargs: Any
-    ) -> Image.Image:
-        """
-        Generate an image by combining style and generation texts with a style image.
-
-        Args:
-            style_text (str): Style-related text prompt.
-            gen_text (str): Generation-related text prompt.
-            style_img (torch.Tensor): Style image tensor. Expected shape is either 3D or 4D.
-            **kwargs: Additional keyword arguments.
-
-        Returns:
-            Image.Image: Generated image as a PIL image.
-        """
-        if style_img.ndim == 3:
-            style_img = style_img.unsqueeze(0)
-        elif style_img.ndim == 4:
-            pass
-        else:
-            raise ValueError('style_img must be 3D or 4D')
-        
-        texts = [style_text + ' ' + gen_text]
-        imgs, _, img_ends = self._generate(texts=texts, imgs=style_img, **kwargs)
-        imgs = (imgs + 1) / 2
-        return F.to_pil_image(imgs[0, ..., style_img.size(-1):img_ends.item()].detach().cpu())
-    
-    def generate_batch(
-        self,
-        style_texts: List[str],
-        gen_texts: List[str],
-        style_imgs: torch.Tensor,
-        lengths: List[int],
-        **kwargs: Any
-    ) -> List[Image.Image]:
-        """
-        Generate a batch of images from lists of style texts, generation texts, and style images.
-
-        Args:
-            style_texts (List[str]): List of style-related text prompts.
-            gen_texts (List[str]): List of generation-related text prompts.
-            style_imgs (torch.Tensor): Batch of style images (4D tensor).
-            lengths (List[int]): List of lengths corresponding to each image.
-            **kwargs: Additional keyword arguments.
-
-        Returns:
-            List[Image.Image]: List of generated images as PIL images.
-        """
-        assert style_imgs.ndim == 4, 'style_imgs must be 4D'
-        assert len(style_texts) == len(style_imgs), 'style_texts and style_imgs must have the same length'
-        assert len(gen_texts) == len(style_imgs), 'gen_texts and style_imgs must have the same length'
-        texts = [style_text + ' ' + gen_text for style_text, gen_text in zip(style_texts, gen_texts)]
+
+    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)
         
-        imgs, _, img_ends = self._generate(texts=texts, imgs=style_imgs, lengths=lengths, **kwargs)
-        imgs = (imgs + 1) / 2
-
-        out_imgs = []
-        for i, end in enumerate(img_ends):
-            start = lengths[i]
-            out_imgs.append(F.to_pil_image(imgs[i, ..., start:end].detach().cpu()))
-        return out_imgs
-
-    def _generate(
-        self,
-        texts: Optional[List[str]] = None,
-        imgs: Optional[torch.Tensor] = None,
-        lengths: Optional[List[int]] = None,
-        input_ids: Optional[torch.Tensor] = None,
-        z_sequence: Optional[torch.Tensor] = None,
-        max_new_tokens: int = 256,
-        stopping_criteria: str = 'latent',
-        stopping_after: int = 10,
-        stopping_patience: int = 1
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Internal generation routine that combines textual and visual inputs to iteratively generate
-        latent representations and decode them into images.
-
-        Args:
-            texts (Optional[List[str]]): List of text prompts.
-            imgs (Optional[torch.Tensor]): Input image tensor.
-            lengths (Optional[List[int]]): Desired lengths for each image in latent space.
-            input_ids (Optional[torch.Tensor]): Tokenized input IDs.
-            z_sequence (Optional[torch.Tensor]): Precomputed latent sequence.
-            max_new_tokens (int): Maximum tokens to generate.
-            stopping_criteria (str): Criteria for stopping ('latent' or 'none').
-            stopping_after (int): Number of tokens to check for stopping condition.
-            stopping_patience (int): Patience parameter for stopping condition.
-
-        Returns:
-            Tuple containing:
-                - imgs (torch.Tensor): Generated images.
-                - canvas_sequence (torch.Tensor): Generated latent canvas sequence.
-                - img_ends (torch.Tensor): End indices for each generated image.
-        """
-        assert texts is not None or input_ids is not None, 'Either texts or input_ids must be provided'
-        assert imgs is not None or z_sequence is not None, 'Either imgs or z_sequence must be provided'
+        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(texts, return_tensors='pt', padding=True).input_ids
+            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(imgs)
-        
-        if lengths is None:
-            lengths = [imgs.size(-1)] * imgs.size(0)
-        lengths = torch.tensor(lengths).to(self.device)
-        lengths = (lengths / 8).ceil().int()
-
-        z_sequence_mask = torch.zeros((z_sequence.size(0), lengths.max() + max_new_tokens))
-        z_sequence_mask = z_sequence_mask.bool().to(self.device)
-        for i, l in enumerate(lengths):
-            z_sequence_mask[i, :l] = True
+            _, z_sequence, _ = self._img_encode(style_img)
+        z_sequence = [z_sequence]
 
-        canvas_sequence = z_sequence[:, :lengths.min()]
         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))
-        seq_stops = torch.ones(z_sequence.size(0), dtype=torch.int) * -1
 
-        for token_idx in range(lengths.min(), lengths.max() + max_new_tokens):
+        for _ in range(max_new_tokens):
             if len(z_sequence) == 0:
                 decoder_inputs_embeds = sos
             else:
-                decoder_inputs_embeds = self.vae_to_t5(canvas_sequence)
+                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:])
-
-            mask_slice = z_sequence_mask[:, token_idx].unsqueeze(-1)
-            if token_idx < z_sequence.size(1):
-                seq_slice = torch.where(mask_slice, z_sequence[:, token_idx], vae_latent[:, 0])
-            else:
-                seq_slice = vae_latent[:, 0]
-            canvas_sequence = torch.cat([canvas_sequence, seq_slice.unsqueeze(1)], dim=1)
+            z_sequence.append(vae_latent)
 
             if stopping_criteria == 'latent':
-                similarity = torch.nn.functional.cosine_similarity(canvas_sequence, pad_token, dim=-1)
-                windows = (similarity > self.padding_token_threshold).unfold(1, stopping_after, 1)
-                window_sums = windows.to(torch.int).sum(dim=2)
-
-                for i in range(similarity.size(0)):
-                    idx = (window_sums[i] > (stopping_after - stopping_patience)).nonzero(as_tuple=True)[0]
-                    if idx.numel() > 0:
-                        seq_stops[i] = idx[0].item()
-
-                if torch.all(seq_stops >= 0):
+                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 == 'none':
-                pass
+            elif stopping_criteria == 'pixel':
+                raise NotImplementedError
 
-        imgs = torch.clamp(self.vae.decode(self.z_rearrange(canvas_sequence)).sample, -1, 1)
-        return imgs, canvas_sequence, seq_stops * 8
+        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: torch.Tensor,
-        noise: float = 0
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Encode the input image into a latent representation using the VAE.
-
-        Args:
-            img (torch.Tensor): Input image tensor.
-            noise (float): Standard deviation of noise to add to the latent sequence.
-
-        Returns:
-            Tuple containing:
-                - decoder_inputs_embeds (torch.Tensor): Embeddings to be used as T5 decoder inputs.
-                - z_sequence (torch.Tensor): Rearranged latent sequence from the VAE.
-                - z (torch.Tensor): Sampled latent vector from the VAE.
-        """
+
+    def _img_encode(self, img, noise=0):
         posterior = self.vae.encode(img.float())
         z = posterior.latent_dist.sample()
         z_sequence = self.query_rearrange(z)
@@ -297,20 +173,10 @@ class Emuru(PreTrainedModel):
         decoder_inputs_embeds = torch.cat([sos, decoder_inputs_embeds], dim=1)
         return decoder_inputs_embeds, z_sequence, z
 
-    def compute_padding_token(self) -> None:
-        """
-        Compute and update the padding token.
 
-        Raises:
-            NotImplementedError: This method must be implemented.
-        """
+    def compute_padding_token(self):
         raise NotImplementedError("compute_padding_token not implemented")
 
-    def compute_padding_token_threshold(self) -> None:
-        """
-        Compute and update the padding token threshold.
 
-        Raises:
-            NotImplementedError: This method must be implemented.
-        """
-        raise NotImplementedError("compute_padding_token_threshold not implemented")
+    def compute_padding_token_threshold(self):
+        raise NotImplementedError("compute_padding_token_threshold not implemented")