DALL-E mini explained | min(DALL-E) | Craiyon | ML Coding Series

Key Takeaways

what's up guys in this video i'm continuing on with the machine learning coding series and in this one i'll be covering dali mini so that's the open source implementation of the of openai's dahle model uh and when i actually say i'll be covering the lean mini what i mean i'll be cover covering this minimal port of the lee mini that's inference only and that's uh basically written in pie torch and the reason is it's much easier to go through this one and to set it up on my uh windows machine so we'll start with this one but in the next video i'll actually be covering the full the original uh jax code base of daley mini which is much more involved but this video will give you enough context to know what's going on okay so as usually i'll first give you the additional context uh in the form of explaining the papers behind this uh and one thing that's very important to understand is that the lea mini is not actually an implementation of the lee model it's more of a mixture and let me open up my onenote so it's more of a mixture of this paper called vqgans or teaming transformers for high resolution image synthesis so that's the paper that introduced the vq gans and i've covered this one and all of these papers actually pretty much in in my previous videos so i'll be linking those as we go along so this is the vq game paper and they also leverage this model which is an encoder decoder transformer model called bart and they use a modification of bart they use this glue from norm shazier we'll see all of those details a bit later and here is the original delete paper and then i'm gonna kind of contrast uh what the lean mini is with the actual dali paper they also have uh like a bits and biases also have a report which will like literally skim through i'm gonna show you uh some of the they actually have very nice training and inference uh diagrams here how the lead mini works and they also have uh comparisons between uh dali and daley mini and you can see there is quite a lot of bullet points here so i'm gonna go through all of those uh a bit later okay so let's first start with with actual papers i'm gonna skim i'm gonna literally just give you the gists because otherwise there's like four papers that would take too much time so let me go straight to it so this is the vq game paper and as i said i previously already covered it so if you want a deep dive go ahead and check it out i'm gonna link it somewhere here okay guys so again on a high level uh keeping a high level because i've covered this in in great depth uh in a separate video we have two components of eq gan basically uh one is the learning the visual code book here and that's this component here and then the second part is after you learn the visual code book then you basically train a transformer model to to auto aggressively model the latent space of visual vectors here so this is the transformer part so basically as i said we have the first part here and we have the transformer part here so how this is trained is very similar to vq vae again that's a paper i've covered before and a couple of differences are that they are using this thing called reconstruction loss uh and the second thing is they're using uh discriminator here so basically adversarial approach and because of that the images from vqgan are much crispier uh compared to what you get using vq vaes which is something that actually the lead version one uh was using uh quickly running through the pipeline here what you do is you feed an image here you then encode it into a lower result into a lower dimensional latent space here so you end up with some vectors as you can tell here and then what uh what what you do is you kind of snap them onto the closest uh visual uh vector that you learn in this in this code book so you can see that for example for this vector here the closest vector uh in as you can see here uh l2 by using the l2 as the definition of closeness uh you snap it on to vector number one so that means this one was closest to this one to this vector uh number one from the visual codebook okay so once you do that for all of the vectors from your low dimensional space here you snap them on to discrete vectors here then you pass them through a convolutional decoder to get the image out and then again i'm going to show you the losses in a second but you train them using a couple of losses such as the reconstruction loss the perceptual loss and the and the adversarial loss here and then the transformer part is i guess very self-explanatory uh once you have this whole pathway here that allows you to map your data set of images into a data set of of basically tokens because you can flatten out this in a in a raster order and you end up with you can directly map your image into a sequence of tokens then you basically can just train your transformer on this new data set of of basically sequences and that's how you learn uh how to well the data distribution uh of of interest okay so that's that's it on a high level uh here are a couple of formulas i'm going to show you so basically as i said what you do is you snap your your vectors so uh basically and i have some glitch with my onenote so you can see you'll you'll be seeing this thing kind of staying here unfortunately but yeah for whatever reason one note is glitching today uh so here is the vector that you estimate then you snap it on to the closest codebook vector here that's the the quantization uh step here so you first encode it uh then you quantize it and then you have the uh basically decoder part and that's how you get the reconstructed image x hat out and with that notation under our belt here is the actual loss uh the vq loss is just as you can see here a sum of couple of losses one of them is reconstruction loss which is l2 in this case here and then there there are a couple of these losses which i think they call commitment loss uh this one so the idea as you can see here is you put a stop gradient uh in this particular case here on on top of your code book vector and then you try and and and make sure to to make the to make your vectors closer to the codebook and here is the opposite here you you kind of put a stop gradient operator on top of your estimated um vectors and then you try and make your code book closer to the output of your encoder as you can see here uh so that's the that's that and then the vqgam actually adds a couple of uh modifications so first one is uh they replaced the l2 loss as you can see here for the um perceptual loss and they introduce an adversarial training procedure as well with the patch-based discriminator so you don't have a single scaler you have as you can as you could saw see on the on the image here you actually have uh basically uh you're doing a patch-based discrimination of whether that patch is real or or false stemming from real or false image and that helps additionally that kind of increases the bandwidth of the learning signal uh compared to just using a classical scaler okay and that's pretty much it that's what we again changes so we have the vq component uh and then we have now the gan component and again the difference here is that they are using the perceptual loss if you're not familiar with what that is uh check out some of my neural style transfer videos but basically the idea is to take uh like a pre-trained convolutional neural network such as like a maybe vgg pre-trained on like imagenet or something uh and then you instead of doing the l2 in the pixel space of the of your reconstructed image you actually feed it inside of that uh like a network and then you do l2 in in the latent space of that network where the features are much more abstract and you're you're basically trying to find um well you're basically capturing the semantics much more than than than high level noise in the picture space again refer to my video if you want to know no more details but this will be enough more than enough for the for the for the code walkthrough okay guys so let me show you one more thing in this paper uh and that's conditional uh synthesis so they mention it here so if the conditioning information c has spatial extent uh we first learned in order of eq again to obtain again an index based representation as you can see here um so you have what they say here you they'll have a height times width of these vectors where the set of vectors has is is this big there is like a big z sub c number of those vectors so that's the size of the codebook for that um condition conditional uh information so they're they're in let me give you a concrete example here this could be maybe uh like um semantic uh segmentation uh maps uh that you can then basically learn a vqgan on top of that type of a data set okay blah blah blah and then you have this newly obtained code book and then youtube regressive structure of the transformer we can then simply prepend r to s which is the uh well that's your your your your image from your target data set and r is is stems from the conditioning information and then basically you do the um negative log likelihood to these entries so in case that was not that clear let me quickly draw it here so again imagine uh you have basically some semantic segmentation um masks here so maybe something like this i'm gonna kind of try and uh draw it somehow so this is some type of a semantic segmentation image and you have a big data set of such images and now what you do is basically train a vq again such that you'll end up with a particular codebook uh for that vector for for the datasets this is going to be your your codebook for this particular uh image here so you then encode it you end up with uh basically um well you end up with h c times w c vectors you then flatten them out and you end up with a sequence and this sequence that which they denote as r is something you can prepend while training uh on your on your data set okay so now obviously for each of these images you have an actual uh image you care about which is the uh image from your from the data set you're trying to for which you're trying to learn the data distribution off so that there will be some image here which is actually original image and this is its semantic segmentation and then you do the same thing here so you'll you first learn the vqgan component so something like this and i'm really terrible at drawing today and and then you encode it here you flatten out this sequence i'm gonna denote it with a different color so we now have a sequence here and now how you're gonna train is you're going to prepend this vector r so we'll you'll grab the vector r from the semantic segmentation map and here you have the target image which they denote as s and now you just collect these tuples for your whole dataset so you basically collect this and then you learn how to to model uh that type of sequence and so how do you later when you want to generate uh a new image you literally you literally grab an arbitrary semantic segmentation map and you you pass it through the encoder here you fetch the uh this vector r and then you put it inside of your transformer so here you'll have a transformer you put the r here and then you auto aggressively keep on generating the sequence and then you pass that through the decoder here so once you generate that sequence here you pass it through the decoder and out comes the image so this will maybe be in practice like 256 tokens which you can then rearrange into 16 by 16 tokens something like this and then that that's like that's the that's the low uh dimensional representation and then you kind of as i said you pass it through the decoder and you get your image out so while i took some time to explain this is because um dali and and the lea mini uh use a very similar approach instead of and instead of using the uh semantic segmentation masks they're actually using text but the logic is super similar so just instead of doing this whole thing you pass a piece of text that's associated with your image here and then you learn how to model them together and that's it so literally by doing that you learn the um you learn the data distribution uh p x y and that's where the bard paper comes into picture so that's that's where this paper comes into play i'm gonna explain to you whoops um i'm gonna explain to you in a second like it's a it's a very simple it's basically like an encoder decoder transformer there is not too much to say there okay so again how bart will come into play here is instead of doing this pipeline with semantic segmentation masks you'll instead have uh an image and you'll have associated caption here so some piece of text here okay so have some piece of text here and you'll basically feed that text through uh encoder part of of uh of bart so let let me kind of denote the such as this like we have encoder part of bard then we hold we have the decoder part of bart decoder is obviously uh attending uh the encoder part that's how the transformers work and so you you grab this text here you basically feed it directly into the encoder part and out comes something from the decoder out comes this thing here and now you use this output of the encoder as r and that's it and now you just again you you'll now have these tuples of uh these conditioning vectors and your uh images from your data set and you just learn how to model the p x y so the the distribution of your uh like images and uh and captions and by the way this is literally this this these two boxes is everything you need to have like that's a mental model of part i have when i'm when i'm discussing it so the only thing we need to see to understand how bart works compared to all of the other transformers is i guess testing here so basically in your regular uh for example in birth model so invert uh what the pre-training objective looked like is you masked out so it's first is bi-directional so it's only encoder part of the of the uh encoder decoder original transformer structure whereas bart is both the encoder as well as the decoder so you're trying to just you just kind of mask out certain tokens and i think like you do that with like 10 20 um probability for each of the tokens and then you try to predict what the token at that particular position was so it's a simple uh like a audio encoding um objective and what's different with bart is you can kind of model you can kind of take a span of of tokens and then mask them out with a single mask token and then make sure to predict the original sequence here so because of this thing because of you can mask it out here however many tokens you want and then predict them here you have additional flexibility which you don't have with a pure encoder uh like transformers such as such as burt and i think they mentioned it somewhere here let me find that piece of text so the objective is called text infilling and i mentioned here a number of text bands are sampled with span lengths drawn from a poisson distribution with lambda parameter equals three and each span is replaced with a single mask token zero length spans correspond to the insertion of mask tokens and okay it's inspired blah blah blah so texting filling teaches the model to predict how many tokens are missing from a span and that's all like it's literally an encoder decoder transformer with this particular um denoising objective and they show a bit later here how it compares against the baselines using different types of like pre-training objectives such as just token masking albert or token deletion or texting fillings a bunch of different sentence shuffling document rotation et cetera et cetera and they compare that against the language model so that's just predicting the next token or ma masked language models such as birth etc and they show well competitive results okay that's the bird part the bard part and then there is a small modification called glue which i'll show you a bit later so glue variance improve transformer from this author gnome shazier who is actually the co-author of the original transformer paper so it's a very prolific author but before that um let me just quickly uh skim over over uh delhi version one which is which came from the paper zero shot text in regeneration again a super similar pipeline to vqgan again they mentioned here they have two stages so in the first stage they say here we train a discrete variational audit encoder to compress each of the 256 to 56 rgb images into 32 by 32 grid of image tokens each elements of which can assume 892 possible values so that that means that the code book size is this big and this reduces the context size of the transformer by a factor of 192 without a large degradation in visual quality because if you divide basically 256 times 256 times 3 because you have rgb image into 32 x 32 you end up with 192 okay and uh the the same authors of this paper uh well i think this is from yeah this is from open eye so basically they they previously trained this model called uh image gpt and uh in that model they were training they were auto aggressively learning how to model uh the the the images in the pixel space which is much more complex because uh transformers have quadratic complexity and obviously the image space can can grow like up to like uh well thousand by thousand images and so it becomes very quickly impossible to model in the pixel space so that's why this reduction this um mapping into a lower dimensional space here okay so again that's stage one uh and it's literally the same as vqgam without the perceptual loss and without the adversarial component and then they say here in stage two we concatenate up to 256 bpe by pairing coding uh text tokens with the 32 x 32 image tokens that come from the stage one and train a not regressive transformer to model the joint distribution over text and image tokens so basically what they said here is that so again they they train the vq vie which you can think as as same as vqgan so yeah they have the encoder part they have the decoder part they kind of have a low dimensional representation here and so what they did what they do instead uh so so they can kind of uh unroll this in a raster order to end up with a uh sequence of tokens s so to get r so to get the conditioning vector and i don't know why they're using r like i guess c c makes more sense but i'll just keep on using the same uh terminology so to get the r component uh that you'll use to to prepend here and then train or progressively train a transformer on top of it uh they just use as as you saw here they just have uh basically they take a piece of text and they bpe encoded so you start with the text and then you apply the tokenizer which gives you a sequence of tokens it just gives you a sequence of tokens and then i assume they just embed those using like a shallow like a simple uh well you have an embedding table you map that into tokens and then you you place that directly here to start um modeling whereas um as you saw dali mini is using bart encoder so that means they'll have additional layers of processing before they they output this r vector so the uh the conditioning vector okay guys that's that's pretty much it um let me quickly show you what they say here there is a lot of complexity hidden in here as i said i covered this one the lee one uh in a separate video so you can check it out they use stuff such as gumball soft max relaxation uh to train this model etc etc and also it's worth keeping in mind i think they mentioned it somewhere here let me just show you this so getting the model to train in 16-bit precision past one billion parameters without diverging was the most challenging part of this project so training this um big model well backing back a couple of years ago this was considered a big model uh was was extremely hard and um that's something i think boris the the main developer of the um the lee mini project is also struggling with but he's using jax and he he has like free tpus from google so that kind of helps because tpus have they did mention somewhere here the tpus have some less problems than than than gpus in certain aspects let me quickly show you one more thing again uh let me just explain this conditioning part so given a text image pair we bpe encode with lowercase captions using at most 256 tokens with vocab size 16384 and encode the image using 32 x 32 tokens with vocab size 8192 the image tokens are obtained using argmax sampling from the dvee uh vae encoded logits without adding any gum uh gumble noise uh finally the text and image tokens are concatenated and modeled auto aggressively as a single stream of data so that's something we saw multiple times already okay guys so that's it it's a very simple model as you saw here uh there is a lot of overlap between the ideas between vq game paper and ali version one they took the ideas from bart they took some other ideas uh boris is experimenting a lot as soon as a new paper comes out such as the deep net paper with novel initializations he tries an experiment with it okay so let me now show you the uh report that uh boris wrote with the other other co-authors again they have a very nice drawing of the training pipeline of dali mini model as you can see here again they have the uh bart encoder so you as you can see you you first have uh these um tuples of images and associated captions then you pass the caption to the bart encoder you get the conditioning vector here you take the image you pass it through the vqgan encoder you get the image tokens you feed them here so you can see here vqgan tokens and then you pass all of that through a bar decoder and you learn how to auto aggressively model these sequences uh simple stuff as for the inference we can see it here so you grab a piece of text you pass it to the bart encoder you feed it into the decoder and then you auto aggressively generate the image tokens uh you can generate multiple series of these image tokens because uh the process of generation is inherently stochastic then you pass all of those through the vqgan decoder you get some associated images so the token sequence from the low dimensional embedding space is projected into the image in the pixel space and then you can use something like clip to filter out the best images so you again feed the same caption to clip and finds the image that that is best described by this particular caption you can see in this particular example white snow covered mountain under blue sky during day time clip found that this image here is the best uh this is is uh well best described by this by this particular caption okay that's it let me quickly uh walk you through and i encourage you to check out the whole report uh in this in this awesome uh waiting biases uh report but let me show you uh some some differences that the uh authors uh of the lee mini model themselves mentioned so we are grateful for the research and pre-trained models published by open eye which were essentially building our model and not all the uh the details on the lead our public knowledge but here uh what we consider to be the main differences so dali uses a 12 billion parameter version of gpt3 in comparison our model is 27 times smaller with about 0.4 billion parameters now this report was written a while uh ago and so this was the the daley mini uh like they now have uh daley mega as well and also i think there is a mistake here because i'm fairly sure the original dolly uh paper use gpt2 not gpt3 which are fairly similar but still then they say we heavily leverage pre-trained models such as vqgan bart encoder and clip we mentioned all of those so far and while opening i had to train all their models from scratch our model architecture takes into account pre-trained models available and their efficiency deleting codes images using a larger number of tokens so they use 32 x 32 whereas the daily mini use uses a 16 by 16 from a smaller vocab and ali uses a vq vae while we use a vq gan okay then they say delete encodes text using fewer tokens at most 256 uh whereas the where dele mini uses 1024 and dali reads text and images as a single stream of data while we split them between this uh well bart encoder and decoder uh this also lets us use independent vocab for text and images daily reads the text through an auto regressive model while we use a bi-directional encoder okay so i made a small mistake when i was explaining how the league v1 works i'm gonna uh show you what i mean by that in in a second but let me finish here so the lee was trained on 250 million pairs of image and text while we used only 15 million pairs okay so the difference is scale obviously the model scale the dataset scale et cetera but also a lot of difference in the architecture and the training method like they don't have the gumball softmax they don't have the vq vae they're using vqgan they're using bart instead of using gpt et cetera et cetera so there's so dali mini is um well not really the lee model it's a it's a mixture of vqgans and bart and everything else so so yeah but still and that's why they are now called crayon not dali mini but in any case let me quickly go back to notebook and then we're jumping to code so i think i mentioned somewhere here that they bpe encode uh text and then they just learn how to auto aggressively uh model that and well i kind of did correctly explain it but not quite so what it do is so let me just uh draw this again here so we have the conditioning part and we have the target sequence so what they do how they learn this is they feed all of these into the transformer and then you learn how to predict all of the tokens so basically here so yeah again as a target you just take this same sequence and you shift it leftwards by by one uh by one step and that's what you're trying to to predict so that this is what you're trying to predict and then here ultimately you have the end of uh image token in this particular case because in this particular example the green part are the image tokens from the vq vae portion okay so that's it the main thing i want to stress here is that the the uh embeddings for the conditioning part are not shallow they're using uh deep uh transformer layers uh because as you as you can as you know like the uh transformer works by uh cross-attending um previous tokens here so it will be attending to deeper representations of this particular uh text here okay guys this there was a bit more than i wanted to to give you when it comes to context and and paper overviews but it is what it is let's now dig into the actual code okay so here we are in my vs code i'm gonna open up the i'm gonna hit the bug here and we're gonna slowly start uh stepping over the code again i'm going to skip all of the unnecessary details i just want to show you the similarity between what we just saw theoretically and the actual code implementation again this uh i'm showing you the code from amindaly repo and uh they only have the inference uh portion of the pipeline and they're um it's it's it's a port in into pytorch and it's whereas the original repo was written in jax we have some arguments here as you can see so i'm using the mega version and not the mini version uh fp16 set to false so that we have higher quality images but the footprint will be uh bigger so i have a 8 gigabyte tpu and otherwise you'll have problems executing this thing so let me show you uh we'll see how the gpu memory is slowly increasing as we're going through the code there's a couple of optimization in the code as well such as as soon as the encoder is created and and the image is encoded they uh dispose of the encoder uh so that's controlled by this parameter let me see whether they have it here i don't think it's in here they have a parameter that literally controls that that that behavior okay so again we have uh we're using mega we have a piece of text i just inputted a random caption here a photo of a funny dog let's see what we get uh there is some seed there is a grid size we're just using we are just generating a single image top cape for top case sampling is set to 256 uh image path uh so this is going to be directly dumped into the root repo root of of this repo so the image will be saved there the models are well already downloaded into this pre-trained i went ahead and downloaded them beforehand so that we don't have to wait and again as i said we're using fp16 so let's continue here so first off we're gonna generate the we're gonna create this uh mindling model and then we are gonna do this generation process so let's see what what happens here so again is reusable is set to false uh that's the parameter that's gonna control that optimization behavior we'll see that in a second so let me continue here so here it is nothing fancy there just storing these variables we have a number of text tokens that's going to specify the maximal size of the caption number of layers is going to be 24 for mega uh as you can see the difference between mega and mini is just in the number of layers number of attention heads also the size of the embedding space the size of the glue embedding space we'll get once we get there i'll briefly show you the paper as well for glue and then the text vocab count uh and the image vocab count kind of differ between these two models in any case we just form the paths here for the uh well by delete here they are literally referring to the bart part of the dahle model so that's kind of um potentially confusing um okay so basically i'm gonna skip over all of this we have a vocab json emerges text so that's gonna basically uh be used to form the tokenizer the text tokenizer okay and now we just form some some um pets which contain the pre-trained models weights and now we hit in initialize tokenizer let me just see where i have yep i have a break point here so basically it's already downloaded so we can ignore all of this uh and now we just open the vocab and here here it is so the let's see what's the size of the vocab so length of the vocab is 50 265. now we go and read this merge merges file uh that contains the basically all of the merges that happened through during the training of the bpe uh tokenizer check out my open ai clip video if you if you're confused by the tokenizer details i've covered that they are in a bit greater depth so let's continue here we just stored the vocab we loaded a couple of seconds ago we basically grabbed the the pairs from this uh like uh merges object and we split them into well into tuples to get the pair so let me show you an example of what the pair look like so pairs and i'm gonna print the first three pairs you can see i and n characters are merged together uh and this weird i don't know how to pronounce this one et cetera et cetera so there is there's gonna be a lot of those bpe merges uh stored in in the pairs this rank from pair is basically going to associate uh integer from zero to the length of pairs with each of the pairs and that's going to be used during the tokenization step i'm going to quickly uh later show you how the tokenization works on a very high level but i'm going to skip this part here because it's fairly intricate to to to describe and you can you can read on bpe yourself we can just treat the tokenizer as a black box uh that's fairly uh common structure between all of these models so i'm not gonna keep on uh explaining it every single time i did cover it in an open eyes clip uh video okay so continuing on because the is reusable flag is set we are not going to beforehand create encoder decoder and the tokenizer which is the vqgan decoder because otherwise the memory would literally explode so i don't have enough memory on my machine to do that so we're going to skip over that and now we have the generic image part okay so we have the model here and let's now generate the image again the encoder and everything else will now be built on the fly so we have the text again we have the photo of a funny dog blah blah same parameters before let's see what this does so there is this um generator uh function that they uh developed let me that's not that vital so we're just gonna generate a stream and then grab doing next we're gonna grab a single image from the stream but let's dig into the actual function okay nothing happened there because it's a generator so once the this line once we hit this line when i go f10 we're going to enter the actual function okay so uh similarly here everything interesting happens in the generate draw stream part so i'm going to ignore everything else here so now let's see what's going on okay so first off we um depending on the grid size we estimate the number of images because the grid size is one for us we just generate a single image so this squid they basically this project offers an like a an option to generate like a grid of multiple images here because of memory constraints i'm just dealing with single image so we can ignore the grid every single time uh okay so because we are in a verbose mode we are printing out some details here tokenizing text et cetera and now we hit the tokenization part so now we tokenize the our input caption before we feed it into the bart encoder so let's see what's going on here so they grab the ass like a separation token the cls token uh and the unknown token here uh there is some processing that they do on on top of the text so they do lower casing and then they encode it as ascii with this error set to ignore and then they do decoding so this is just going to do some simple primitive processing of text nothing fancy there let me let me print the text for you here let's see what happened basically nothing happened with this particular sentence because it was fairly simple and then you literally kind of split the sentence into words and splitting is happening so basically you use spaces to split the sentence into words and then each of the word of the words is going to be split into these sub words here because the this cat bite pair encoding which i'll skip height it's fairly complex to explain but it's just gonna split the word into bpe uh like subreddits and then each of the sub births is basically associated with the like with the specific index uh integer from from the vocab here and that's it on a high level and then we just kind of bracket that with the cls token in the beginning and that this c separation token uh at the at the end and that's how we form uh basically a sequence of tokens for our particular text okay then they do optionally here the cropping but because we are less than 64 tokens i think this is set to 64. let me see what this number is but i'm fairly sure it was 64. yep so if if it's like a longer than that length then they crop the the the tokens but for us it's not okay so just some printing because it's verbose mode blah blah blah uh now as you can see this is the interesting part so they they create this variable uh text tokens uh which is initialized as all ones and it has you can be confused by a couple of details here so first why why is ones is because one is a token in this vocab and that's kind of implicit here because this is as i said just a port from the original repo so sometimes sometimes things are hard to understand unless you well unless you understand what's going on so here here am i trying to explain to you that so that's one detail the second important detail is there is this number two and why they do that is because they have a basically analog thing to classify our free guidance that we saw in the fusion model so you can basically do the same thing for the class of auto regressive generative models and that's why they have two so as you see here they create the first sequence will just uh contain the cls and the uh uh separator token so it's going to be literally empty sequence everything else is going to be pad paddings so that's what we do on this line and then on the second line we add the actual tokens okay so let me show you uh what this thing is so text tokens if i print text tokens zero and then first maybe well seven tokens you can see it's an empty sequence because this is a cls this is a separator token and everything else is padding and the second one contains the actual uh tokens for our caption okay that's it and now we just convert that into tensor with a long data type we put it onto gpu and we continue on okay now here's the optimization i was mentioning because we're using this uh not reusable uh we're because we said this flag is reusable to false we now on the fly initialize the encoder so the encoder is going to be the vq gann encoder so let me step inside of here and let's see what's going on so again because i already downloaded the model we can skip over over all of these steps and we're going to focus on how this deliberate encoder is constructed so the interesting part is usually is going to be in the forward uh pass of of the of the model and uh after we construct it you can see here we just load the parameters of the pre-trained model that was trained in the actual daily mini uh repository we're gonna load those parameters and then we're gonna push this and this uh model onto the gpu okay so let's see how the bart encoder looks like here we are we have the the vocab size of 50k something we form the embedding tables for the uh well that's the text vocab uh because encoder again remember has the text vocab and the decoder will be generating the image token so it will have a separate uh image vocab okay so we generate that table the embedding vectors will be 2048 we then generate the positional encodings here this is going to be only 64 because if you recall before we had the cropping so that we make sure that the number of tokens that we feed into the bart encoder is always up to 64. never more than that and that's why we never need to go above 64 positional encodings okay now we start forming the um basically um encoder layers and we'll have 24 of these uh for the uh bart encoder for dolly mega i'm gonna hit go here and put a small break point uh in encoder layer which is a simple transformer uh block nothing fancy the interesting part is this glue so i'm gonna set uh one uh break point there uh and that's it so let's now continue execution here so we end up in the encoder layer as you can see it's just a simple transformer uh block we have the encoder we have the self attention module we have the some layer norms and we have glue which is just a modification of the uh your regular uh like uh mlp that's used the feed forward uh like a layer of the transformer black okay so let's see how glue looks like here is glue uh it's a very simple modification by noam chazier and i'm going to now open up the paper side by side to show you what this is and to make sense out of it but it's a very experimental finding so there is nothing uh well fundamental uh understanding you you can get from this okay guys here it is side by side uh here are the equations from the uh glue paper binomials here and you can see that this layer here implements exactly this line here so the idea here was they they kind of like he ablated a couple of um he made a couple of modifications why not using different activation functions instead of really why not using jelly or swish etc etc why not combine them combine them in a bit different ways using um similarly to the gated linear units so you can see so uh like the lead mini ended up using this variant here and let me just kind of convince you that that's indeed the case so let's see so we have we get the input x this is x and then we pass it through layer norm so ignoring that part we we pass it we then pass it through a linear layer so that's that corresponds to basically this uh bracket here and after that we we pass it through w so that's this part the sigma of of this so that's this uh term and then we have v which is formed from x as well from the input we also pass it through the linear layer so that's that's this part the x times big v and finally we do uh element wise uh product so hadamard product here and then we pass that product through uh a layer norm and so that's that's that's basically what happened here and finally we we do uh a feed forward to a linear layer so that's the w that's the that's the sorry that's the w-2 here and that's thus you can see that this simple floral pass of this glue model implements exactly this ffn glue from normations years paper and that's it now let's go back to the code i'm gonna uh remove this breakpoint so it's a simple sim simple modification instead of having just like a like a feed forward uh like a basically a linear layer for all by relu followed by a linear layer it's a bit more intricate and they showed experimentally that this gives a better performance but there is no ultimately no theory behind why this should work and that's it that's i guess deep learning let's continue i'm gonna now um we are now back into the encoder layer and again encoder layer is simply a transformer block so we have as you can see here we have some um layer norm and then a tension and the norm and then the residual connection and then then we pass through the feed four part and then we have again the residual connection and we return back the variable so that's it so i'm gonna basically ignore all of these break points and let's form the 24 layers of this bert of bart um encoder so i'm going to hit f5 we're going to form 24 layers and then we are now here we form additional layer norm additional air norm and then we form these token indices which are going to be used to index the positional embeddings here and that's it and we just do times two because remember we are doing the classifier free uh trick for all progressive models that's why we have a batch size of two here okay that's it that's the that's the initialization of the bart encoder now we are loading the weights here and then we kind of push those weights into the encoder structure and then we delete the parameters now you can already see that this is pushing my gpu somewhat well well actually not that much okay so i'm surprised this is actually not that big of a memory footprint i think some of the later components will be spiking it up but yeah let's continue so after the initialization of encoder let's see what else happens okay and now we grab the text tokens and we pass them through the bart encoder as we saw in the paper part of the video okay let's step into it so we're gonna hit the for a function of the bart encoder here so let's see what happens so we ask where the text tokens are not equal to one and one is again remember uh they're not using a constant here they're just using one to represent the pad tokens so it's kind of hard coded and dirty but it works okay so not equal to one that's the actual content that's where attention mask is gonna say uh true so let's now do something like this let me print this you'll see that it's basically true for the second component of the text tokens where we actually passed the the caption tokens we have a bunch of truths here but for the first one we only have true true for the first two because those are the cls and the separator token okay now what we do is we pass those tokens into the uh we embed them and then we just add the positional encodings using the the these uh well the the pose tokens indices uh that were previously constructed and that's it that's simple stuff there and now we just pass it to a layer norm and then we keep on uh processing uh those embeddings uh via the transformer bar transformer layers and that's it i'm just gonna hit f5 here we're gonna do a forward pass through the transformer and we end up with the final representation here called encoder state okay and because again this is the optimization i was mentioning we now delete the weights of the encoder let's see how the the memory footprint is gonna reduce okay so the memory footprint went up you can see it here and after we hit delete and empty cache you're gonna see how it kind of dips down hopefully i'm not sure what's going on um why is my gpu not reacting here not sure in any case let's continue so now we have the initialize decoder part i'm gonna hit step over and we enter the emit decoder so again it's already downloaded so we don't care about that part uh here we construct it and again we'll load the way it's a similar structure as as with the encoder uh now we are passing the image vocab etc so let's enter the constructor here so let's see what's the difference so uh here we form again the embedding table this time we have the image vocab size plus one because let me just think here um because they also have the they passed this uh this beginning of sentence token uh and that's the additional token that they need in addition to the um image to image code book so i think they mentioned this somewhere in the in the uh report let me show you this so basically it's called bos token okay so they say here at inference time blah blah so bos token is fed through the bar decoder um and yeah that's so that's basically what why there is plus one as i said not the cleanest way to do something but yeah this is just a port of the original repo and now we have we formed the uh table for the uh position um uh encodings and there is only 256 of those because uh as you recall from the paper part uh it's gonna the latent space is going to be 16 times 16 of those uh basically uh vectors that are then going to be fed into dvq and decoder okay let's continue here so we now generate a bunch of decoder layers and decoder layers have additionally the cross attention component because this is again encoder decoder transformer so let me see whether it makes sense to um to kind of show you that part let me just think okay i'm going to put a break point into the mid part here and we can maybe do a single pass uh here because they're doing something interesting with this attention state we'll see that a bit later and here as well i'm going to put the break point here and now let's continue okay so let's hit the first decoder layer it's again simply there is a self-attention part there is the decoder there is a cross attention part and there is the feed forward part so the glue variant in this particular case so that's pretty much it so i'm going to now remove the break point from the init and let's just generate a bunch of these layers oops let's now do that and we do that for 24 times again i think let me just print the layer count 24 okay so there's 24 decoder layers as well um okay we have some layer norms nothing fancy and finally we have we take the final embedding vectors and we map them into the space that has this dimensionality because remember we now want to sample those image tokens that's why we need to to to go from 2048 which was the internal model dimension we wanna uh go into into this dimension here okay let's continue on here and that's it okay that's pretty much it we load the weights again 24 layers of decoder blocks and finally then we we map into the image vocab uh size uh dimensionality space output space so that we can sample from it later okay again we do the parameters we put the decoder on top uh we push it to rg to my gpu so let's see whether something changes here nope uh it's a bit different behavior compared to what is we've seen before usually there is a spike bigger than this one happening and i'm tracking this slot here by the way okay in any case now comes the interesting part so here we are going to now start sampling from the decoder those image tokens and this one is already trained so everything example is going to be meaningful and then we're just going to pass it through the vqgan decoder and that's the that's the end of the program okay let's go here so um we're dealing with float 32 so this doesn't make much sense but with float30 uh float16 uh this kind of puts it puts all of the operations inside of this context in the float 16 regime which helps us save uh some memory uh so in the case when you have multiple images and this expanded indices make sense because you will not replicate the the textual captions for multiple images because you'll be generating multiple images for one uh particular caption but we don't care about that here so we just um basically grab the uh encoder state and and and the text tokens and those are of dimensionality so this is going to be 264 if you recall and this one is going to be i guess uh 2 2 000 something 2048 because that's the uh output of the encoder right so two oh actually yeah yeah 264 2048 because we just embedded this into this dimensionality here okay so we formed the mask again this is the padding token so mask is going to be uh true uh for for the non-pad tokens uh here is the we form this attention state which is gonna contain the keys and values from each of the layers and for each of the positions in the bart decoder that's just how they decided to implement this and finally image tokens is gonna contain that's the output array that's gonna contain two 256 uh tokens uh plus one because the first one will be the i think yeah the bos token the beginning of uh sentence token so let's continue we're going to initialize this vector initially as you can see here with this value here which is going to be the beginning of sentence tokens id so let me show you what that looks like i'm going to take the first five elements and all of them are going to be equal like this so that's how how they initially uh we'll initialize the image tokens and then later we're gonna keep on as we produce the image tokens we're gonna feed them back into this image tokens array okay so nothing fancy there uh continuing on we form these token indices and we have some settings like temperature set to 1 top k256 super conditioning factor which is used for the classifier free guidance part and that's it those are the settings and now we start iterating basically a for loop 256 times we're gonna generate a novel uh image uh vector uh and uh basically then we're gonna feed that into the vqgan decoder as i said multiple times already so we care about this part of the loop because this one will enter this one once we have 256 tokens um we'll enter the image grid from tokens function so let's focus on this part here so what happens here as you can see whatever we we take the image so i is initially zero which means we initially feed the bos so the beginning of of sentence token and we feed the output result into in into the first slot of this image tokens and that's how we start populating this the image tokens uh array again its dimensionality is one plus 256. it's precisely because of this behavior here we pass the tokens we pass the attention mask the encoder state the attention state which is going to contain the keys and values from the decoder and that's it and we we pass the token indices so we basically pass where we are in the generation uh in the generation sequence okay in the generated sequence okay now let's step inside of the uh forward function of the bar decoder so we grab the number of images uh we then basically make sure that the batch dimension is corresponds to to the super conditioning thingy so it's going to be two okay so it's going to be two there uh we grabbed the the list of previous tokens uh so previous tokens dimensionality is uh one so it's going to be initially just a bos token and then because we only have a single image this is just going to replicate uh the previous tokens we're going to have two bos tokens and that's it again the reason being the super conditioning that they are doing so the classifier free guidance thingy for progressive models some clamping i don't think this actually is needed and then so now we do the embedding but remember this time this is the these are the image embeddings okay so these are the as you can see here these are this is the embedding table this is a separate vocal vocabulary compared to the text that we use for the bard encoder part okay so we do the encoding then we add the positional um encodings here and now we have decoder state it's going to be i guess uh what two one two thousand something maybe yeah no to 2048 make sense okay because we have uh bos a single token we have two because of super conditioning and we just embedded them into this uh dimensionality here so all that makes sense let's continue on here now we add the additional dimensions now it's going to be 2 1 2048 and now we just iterate through the decoder uh layers of this decoder so we pass the decoder state the encoder state because we have the cross attention remember and we pass the attention state so tension state is going to be um basically a number of layers uh and then four because we we're gonna keep keys and values and times two because we have the the super conditioning and then finally we're gonna have 256 and 2048 so again 24 because we have 24 layers four because we are saving key and value for each of the layers but times two because we are using uh two sequences the one with uh empty sequence and the one with the actual caption because of super conditioning and then 256 is because we have 256 image tokens in the decoder and 2048 because that's the model dimension okay just kind of breaking down that the shape is always useful we pass the attention mask for the input caption and we pass the token now let's hit the forward function here so here we just generate the attention mask depending on where we are in the generation process so because token token index currently is zero we are we're just we're we're just have passed the bos so the beginning of census token and that's why this self-attention mass will currently just be active in the first in the zero slot and everything else is false so you can see here that the first will be true and then everything else is false and then we just kind of make sure that we have batch dimension of two because of super conditioning and then we just pass the equator state through the layer normalization nothing fancy there and now we pass it to the attention part so i'm going to skip the attention part is fairly um well actually let me see this is where we will store the keys and values so let me show you this part so let me just see whether i have uh um okay so we have a self-attention we have a break point there so i'm going to hit that breakpoint and let's see what we do so we store the keys and values and let's see what their dimensionality is so we have 2 1 2048 and the same for for a value vector we basically concatenate them and then we store it inside of this attention state and that's it and now grabbing from the attention state we're going to grab all of the previous keys up to this token which in this particular case is only these keys and values because we are just starting the generation but that's how the how the logic is implemented using this attention state we're basically fetching all of the previous keys and values and that's it and now we just pass it through the we just do the attention logic and we return back here that's it now we have the uh cross attention part we passed the decoder and encoder and that's it everything else is standard transformer logic we pass it now through the glue and that's a single step through the decoder layer i'm going to now remove the breakpoints here i'm going to remove the break point here as well and let's continue on here i'm going to hit f5 we went through all of the layers and now we pass it through again this final ln is just like layer norm and finally here is where the mapping happens so here we're gonna take the vector that's 2 1 2048 and we're going to map that into the corresponding into the image space so basically now the logits will be of 2 1 16 416 because that's the basically the image space dimensionality we grab the settings such as top k temperature super conditioning vector now this part is kind of um not completely clear because it's so hard coded so let me try and uh break it down so first we grab all of the um along the batch dimension we grab everything and that makes sense we take minus one because we only only only wanna we only care about the last the largest from the from the um from the last token because we're trying to sample from there that makes sense but this part doesn't make sense why is it um this number which is smaller than the number here and my hypothesis is that well it might be that the original vq gan checkpoint they used um had a bigger vocab compared to their target fine tuning after they've done the fine tuning of of of big guns not completely clear uh about about this part uh in any case so here is the super conditioning part so we grab the logits and as you can see here we grab now we grab the so these are the lodges that correspond to the empty sequence and these are the logics that correspond to the actual caption and we do the logic we saw in the diffusion models as well and that's how we combine the largest to fo to to form the final uh super condition distribution here okay so i think let me just check this i'm fairly sure so this is this uh reverse half wings twitter profile i think was the first one to have suggested this type of uh super conditioning for ultra aggressive models so so she said you can apply a similar trick to classifier free guidance to other regressive transformers to sample from a synthetic supercondition distribution uh blah blah blah training to try this and you can see how well basically you trade off the variance uh with the um quality of the images and finally this is what we do so you take the unconditional logics and you just basically add up on top of that this difference times the conditional scale so this is the same expression we are now using in in the code okay let me go back to the code here so that's basically if you if you cannot decompose this you get that same expression now we do some this is the top case sampling part so we do we sort and we as sort in the descending fashion such that the first lodge will be the biggest one and then they start descending okay and then is kept so we take the original logits and we find basically the 256th biggest logits here because here's a sorted one and because we have 256 here so we're going to grab the 255th biggest logit and use that as a threshold and all of those logics which are bigger than that threshold are kept that's how the top case sampling works then we do some uh this is just for for step stabilization purposes uh temperature exponentiation and then we just multiply with uh with this mask so again we are just keeping those logics that are that are uh basically big enough where big enough is defined by the top k parameter okay and finally we just sample from the distribution and we get the image token okay so we now have image token here that's the that's the token we sampled and uh in the next iteration of this um uh loop uh basically now we're gonna feed that token as the input and then we're gonna keep on ultra aggressively generating the image tokens so now i'm gonna uh try and skip uh all of this so i'm gonna go and disable all breakpoints and i'm gonna enable only this one here i'm gonna hit f5 and i'm gonna let it run until we generate all of the 256 uh image tokens basically okay here we are uh now let's enter the uh vqgan decoder as you can see here we only pass we ignore the first token in the image tokens because that was the bos so that was the beginning of sentence token so we ignore it so let's jump here i'm gonna enable all of the breakpoints uh let's see whether i have a breakpoint here yes i do so let's continue doing this okay here we are we delete the decoder that's going to hopefully release some memory and empty we empty the cache so here finally we have this tip in memory this is what i i'm used to and i'm not sure why that didn't happen previously when we were deleting the encoder uh yeah that was weird in any case now we can see it here okay so we now initialize the uh d tokenizer or the vqgan decoder so let's see what i have yeah i have a breakpoint there again similar structure as before uh we create the tokenizer object and then we just load the parameters so let's see how it looks like again we pass it the vocab count and the embedding count so it's going to have 256 tokens and this is the vocab count which corresponds to the this is the same size as what the bar decoder had right okay we formed the embedding table uh just some processing layers and decoder i'm going to ignore how the architecture looks like it's basically a bunch of like upsampling layers and attention blocks and resnet blocks nothing super um i guess informative so we just have resonant blocks some attention blocks which literally do attention uh treating uh image tokens as as well treating images as as sequences and thus we can just apply attention um just some up sampling as i said nothing worth digging deeper into to be honest uh you can go and check it out at your own pace if you want okay so we now load the parameters and we push the the tokenizer the vqgan on the gpu and now we do the forward pass through the uh this is the interesting part not that complex but still interesting so we have said that these are 256 tokens so that shape is 256. again i'm not sure why they're doing the clamping because we we we are certain that this will be inside of these boundaries already so there is some weird thing happening there uh there's a consequence of that cropping we saw in the when the with the logits so that that might explain it okay grid size is one because as i said we're generating just a single image and now there is a difference whether um well basically the the shape of this um how do we feed the the image tokens into the actual convolutional decoder part and we're going to enter this branch because the seamless is set to to false and we just embed using the uh embedding table we embed our tokens and i assume this is the same embedding table as the one that we had in the bart decoder that should be shared weight and now we just um do this view thingy and we end up with 16 16 256 and that's precisely the latent space volume that we saw in the papers that we're going to now feed into the vqgan convolutional uh decoder part okay some permutations some processing with the cnn basically this is a like a single uh channel cnn sorry single uh the spatial extent of the cnn is one times one and that then we just up sample this this image and we end up finally with a set of shape 256 to 56 and three channels and that's our image and now we do some clipping uh renormalization back into the zero 255 range and that's it some again manipulations on the shape and let's see what we whoops let's see what we end up with here so 256 to 56.3 so that's our rgb image and that's it guys basically we now uh push the image we convert it into uh unsigned integer uh 8-bit format we push it onto cpu from the cuda and we basically convert it into numpy and now we basically yield back up p a pil image build image and that's it and then we just save the image and that's it now there is also this print ascii from image function that's very neat actually so again ascii from image makes sense when you're running this from a console it's just gonna print the image [Music] in the ascii format i guess okay that's it let me step over this and we are pretty much done that's the end of the program let me now show you the generated image i should have it somewhere here let me open it up and here it is so here's the uh what was our caption our caption was um a photo of a funny dog and here is a dog i'm not sure whether it's funny but uh basically that's it let me try a different caption and i'll get back to you once we once once it's wrong maybe like something a bit more complex this might break the dali mega but it's worth a shot a photo of a funny dog maybe riding a bicycle bicycle let me run this one and let's see what we get okay guys here is the actual image uh it's kind of weird you can kind of see it's trying to merge the concept of a dog within with the concept of a bicycle but it's still not that expressive and big to achieve that so scale is vital for these types of models and that's it uh do let me know whether you found this video useful as as usually uh leave down the comments below if you have any feedback for me and until next time bye bye [Music]