tanner/min-dalle-test
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

is_reusable

Browse Source
This commit is contained in:
Brett Kuprel 2022-06-30 11:25:24 -04:00
parent 3e64e868ef
commit f951424e38
4 changed files with 64 additions and 167 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
image_from_text.py
View File

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

187
min_dalle.ipynb vendored
View File

File diff suppressed because one or more lines are too long

18
min_dalle/min_dalle_flax.py
View File

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

18
min_dalle/min_dalle_torch.py
View File

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