Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -14,8 +14,8 @@ import gc
 
 torch.set_grad_enabled(False)
 torch.set_num_threads(os.cpu_count())
-torch.backends.cudnn.enabled = True
-torch.backends.cudnn.allow_tf32 = True
+torch.backends.cudnn.enabled = False
+torch.backends.cudnn.allow_tf16 = False
 
 MIN_DALLE_REPO = 'https://huggingface.co/kuprel/min-dalle/resolve/main/'
 IMAGE_TOKEN_COUNT = 256
@@ -25,7 +25,7 @@ class MinDalle:
     def __init__(
         self,
         models_root: str = 'pretrained',
-        dtype: torch.dtype = torch.float32,
+        dtype: torch.dtype = torch.float16,
         device: str = None,
         is_mega: bool = True, 
         is_reusable: bool = True,
@@ -188,7 +188,7 @@ class MinDalle:
         if len(tokens) > self.text_token_count: 
             tokens = tokens[:self.text_token_count]
         if is_verbose: print("{} text tokens".format(len(tokens)), tokens)
-        text_tokens = numpy.ones((2, 64), dtype=numpy.int32)
+        text_tokens = numpy.ones((2, 64), dtype=numpy.int16)
         text_tokens[0, :2] = [tokens[0], tokens[-1]]
         text_tokens[1, :len(tokens)] = tokens
         text_tokens = torch.tensor(
@@ -232,9 +232,11 @@ class MinDalle:
         token_indices = torch.arange(IMAGE_TOKEN_COUNT, device=self.device)
         settings = torch.tensor(
             [temperature, top_k, supercondition_factor], 
-            dtype=torch.float32,
+            dtype=torch.float16,
             device=self.device
         )
+        
+
         for i in range(IMAGE_TOKEN_COUNT):
         
             if(st.session_state.page != 0):
@@ -243,7 +245,7 @@ class MinDalle:
 
             #torch.cuda.empty_cache() 
             #torch.cpu.empty_cache()                
-            gc.collect()
+            #gc.collect()
         
             image_tokens[i + 1], attention_state = self.decoder.forward(
                 settings=settings,
@@ -254,7 +256,7 @@ class MinDalle:
                 token_index=token_indices[[i]]
             )
 
-            if ((i + 1) % 32 == 0 and progressive_outputs) or i + 1 == 256:
+            if ((i + 1) % 16 == 0 and progressive_outputs) or i + 1 == 256:
                 yield self.image_grid_from_tokens(
                     image_tokens=image_tokens[1:].T,
                     is_seamless=is_seamless,