File size: 13,982 Bytes
19f124c
 
 
5f68c1b
19f124c
 
5f68c1b
 
 
 
19f124c
 
5f68c1b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19f124c
5f68c1b
 
19f124c
5f68c1b
19f124c
 
 
 
 
 
0021de3
19f124c
 
 
 
 
5f68c1b
19f124c
 
 
 
 
 
 
 
5f68c1b
 
 
69ab272
5f68c1b
 
 
 
19f124c
 
 
 
5f68c1b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0021de3
 
 
 
 
 
 
 
69ab272
605e556
5f68c1b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
935404c
5f68c1b
 
 
 
 
605e556
5f68c1b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
69ab272
 
5f68c1b
69ab272
 
 
5f68c1b
 
 
 
 
 
 
 
 
 
 
69ab272
5f68c1b
69ab272
 
5f68c1b
69ab272
5f68c1b
69ab272
 
 
5f68c1b
69ab272
 
 
5f68c1b
 
 
 
 
 
 
69ab272
 
5f68c1b
6caadef
5f68c1b
 
 
 
 
 
 
 
69ab272
5f68c1b
 
69ab272
5f68c1b
 
69ab272
5f68c1b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0021de3
 
 
 
 
 
 
 
69ab272
0021de3
 
 
 
5f68c1b
 
 
69ab272
5f68c1b
 
 
69ab272
 
5f68c1b
 
 
0021de3
5f68c1b
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
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):
    """
    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.
    """
    config_class = EmuruConfig

    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_name_or_path)

        t5_config = T5Config.from_pretrained(config.t5_name_or_path)
        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_name_or_path)
        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

    @torch.inference_mode()
    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())
    
    @torch.inference_mode()
    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, min(stopping_after, similarity.size(-1)), 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")