is_reusable

image_from_text.py

@@ -44,9 +44,9 @@ def generate_image(
     image_path: str,
     token_count: int
 ):
-    is_expendable = True
+    is_reusable = False
     if is_torch:
-        image_generator = MinDalleTorch(is_mega, is_expendable, token_count)
+        image_generator = MinDalleTorch(is_mega, is_reusable, token_count)
 
         if token_count < image_generator.config['image_length']:
             image_tokens = image_generator.generate_image_tokens(text, seed)
@@ -56,7 +56,7 @@ def generate_image(
             image = image_generator.generate_image(text, seed)
 
     else:
-        image_generator = MinDalleFlax(is_mega, is_expendable=True)
+        image_generator = MinDalleFlax(is_mega, is_reusable)
         image = image_generator.generate_image(text, seed)
 
     save_image(image, image_path)

min_dalle/min_dalle_flax.py

@@ -9,11 +9,11 @@ from .models.dalle_bart_decoder_flax import DalleBartDecoderFlax
 
 
 class MinDalleFlax(MinDalleBase):
-    def __init__(self, is_mega: bool, is_expendable: bool = False):
+    def __init__(self, is_mega: bool, is_reusable: bool = True):
         super().__init__(is_mega)
-        self.is_expendable = is_expendable
+        self.is_reusable = is_reusable
         print("initializing MinDalleFlax")
-        if not is_expendable:
+        if is_reusable:
             self.init_encoder()
             self.init_decoder()
             self.init_detokenizer()
@@ -48,12 +48,12 @@ class MinDalleFlax(MinDalleBase):
     def generate_image(self, text: str, seed: int) -> Image.Image:
         text_tokens = self.tokenize_text(text)
 
-        if self.is_expendable: self.init_encoder()
+        if not self.is_reusable: self.init_encoder()
         print("encoding text tokens")
         encoder_state = self.encoder(text_tokens)
-        if self.is_expendable: del self.encoder
+        if not self.is_reusable: del self.encoder
 
-        if self.is_expendable:
+        if not self.is_reusable:
             self.init_decoder()
             params = self.model_params.pop('decoder')
         else:
@@ -65,13 +65,13 @@ class MinDalleFlax(MinDalleBase):
             jax.random.PRNGKey(seed),
             params
         )
-        if self.is_expendable: del self.decoder
+        if not self.is_reusable: del self.decoder
 
         image_tokens = torch.tensor(numpy.array(image_tokens))
 
-        if self.is_expendable: self.init_detokenizer()
+        if not self.is_reusable: self.init_detokenizer()
         print("detokenizing image")
         image = self.detokenizer.forward(image_tokens).to(torch.uint8)
-        if self.is_expendable: del self.detokenizer
+        if not self.is_reusable: del self.detokenizer
         image = Image.fromarray(image.to('cpu').detach().numpy())
         return image

min_dalle/min_dalle_torch.py

@@ -16,14 +16,14 @@ class MinDalleTorch(MinDalleBase):
     def __init__(
         self, 
         is_mega: bool, 
-        is_expendable: bool = False,
+        is_reusable: bool = True,
         token_count: int = 256
     ):
         super().__init__(is_mega)
-        self.is_expendable = is_expendable
+        self.is_reusable = is_reusable
         self.token_count = token_count
         print("initializing MinDalleTorch")
-        if not is_expendable:
+        if is_reusable:
             self.init_encoder()
             self.init_decoder()
             self.init_detokenizer()
@@ -84,24 +84,24 @@ class MinDalleTorch(MinDalleBase):
         text_tokens = torch.tensor(text_tokens).to(torch.long)
         if torch.cuda.is_available(): text_tokens = text_tokens.cuda()
 
-        if self.is_expendable: self.init_encoder()
+        if not self.is_reusable: self.init_encoder()
         print("encoding text tokens")
         encoder_state = self.encoder.forward(text_tokens)
-        if self.is_expendable: del self.encoder
+        if not self.is_reusable: del self.encoder
 
-        if self.is_expendable: self.init_decoder()
+        if not self.is_reusable: self.init_decoder()
         print("sampling image tokens")
         torch.manual_seed(seed)
         image_tokens = self.decoder.forward(text_tokens, encoder_state)
-        if self.is_expendable: del self.decoder
+        if not self.is_reusable: del self.decoder
         return image_tokens
         
 
     def generate_image(self, text: str, seed: int) -> Image.Image:
         image_tokens = self.generate_image_tokens(text, seed)
-        if self.is_expendable: self.init_detokenizer()
+        if not self.is_reusable: self.init_detokenizer()
         print("detokenizing image")
         image = self.detokenizer.forward(image_tokens).to(torch.uint8)
-        if self.is_expendable: del self.detokenizer
+        if not self.is_reusable: del self.detokenizer
         image = Image.fromarray(image.to('cpu').detach().numpy())
         return image