Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -39,10 +39,9 @@ class MinDalle:
         self.dtype = dtype
         self.is_verbose = is_verbose
         self.text_token_count = 64
-        #Changed
-        self.layer_count = 24 if is_mega else 6
-        self.attention_head_count = 32 if is_mega else 8
-        self.embed_count = 2048 if is_mega else 512
+        self.layer_count = 24 if is_mega else 12
+        self.attention_head_count = 32 if is_mega else 16
+        self.embed_count = 2048 if is_mega else 1024
         self.glu_embed_count = 4096 if is_mega else 2730
         self.text_vocab_count = 50272 if is_mega else 50264
         self.image_vocab_count = 16415 if is_mega else 16384
@@ -238,29 +237,27 @@ class MinDalle:
         for i in range(IMAGE_TOKEN_COUNT):
             if(st.session_state.page != 0):
                 break
-         
             st.session_state.bar.progress(i/IMAGE_TOKEN_COUNT)
-      
-            #torch.cuda.empty_cache() 
-            #torch.cpu.empty_cache()                
-            #with torch.cuda.amp.autocast(dtype=self.dtype):
-            image_tokens[i + 1], attention_state = self.decoder.forward(
-                settings=settings,
-                attention_mask=attention_mask,
-                encoder_state=encoder_state,
-                attention_state=attention_state,
-                prev_tokens=image_tokens[i],
-                token_index=token_indices[[i]]
-            )
 
-            # with torch.cuda.amp.autocast(dtype=torch.float32):
-            if ((i + 1) % 32 == 0 and progressive_outputs) or i + 1 == 256:
-                yield self.image_grid_from_tokens(
-                    image_tokens=image_tokens[1:].T,
-                    is_seamless=is_seamless,
-                    is_verbose=is_verbose
+            torch.cuda.empty_cache() 
+            #torch.cpu.empty_cache()                
+            with torch.cuda.amp.autocast(dtype=self.dtype):
+                image_tokens[i + 1], attention_state = self.decoder.forward(
+                    settings=settings,
+                    attention_mask=attention_mask,
+                    encoder_state=encoder_state,
+                    attention_state=attention_state,
+                    prev_tokens=image_tokens[i],
+                    token_index=token_indices[[i]]
                 )
-           
+
+            with torch.cuda.amp.autocast(dtype=torch.float32):
+                if ((i + 1) % 32 == 0 and progressive_outputs) or i + 1 == 256:
+                    yield self.image_grid_from_tokens(
+                        image_tokens=image_tokens[1:].T,
+                        is_seamless=is_seamless,
+                        is_verbose=is_verbose
+                    )
 
     def generate_image_stream(self, *args, **kwargs) -> Iterator[Image.Image]:
         image_stream = self.generate_raw_image_stream(*args, **kwargs)
@@ -278,6 +275,7 @@ class MinDalle:
             image = image.reshape([grid_size ** 2, 2 ** 8, 2 ** 8, 3])
             yield image
 
+
     def generate_image(self, *args, **kwargs) -> Image.Image:
         image_stream = self.generate_image_stream(
             *args, **kwargs,