CPU only for now

app.py

@@ -3,10 +3,10 @@
 
 Automatically generated by Colab.
 
-Original file is located at
-    https://colab.research.google.com/drive/1I47sLakpuwERGzn-XoNct67mwiDS1mQD
 """
 
+DEVICE = 'cpu'
+
 import matplotlib.pyplot as plt
 import matplotlib
 
@@ -108,7 +108,7 @@ def latent_code_from_text(text,):# args):
     coded1 =torch.Tensor.long(coded1)
     with torch.no_grad():
         x0 = coded1
-        x0 = x0.to('cuda')
+        x0 = x0.to(DEVICE)
         pooled_hidden_fea = model_vae.encoder(x0, attention_mask=(x0 > 0).float())[1]
         mean, logvar = model_vae.encoder.linear(pooled_hidden_fea).chunk(2, -1)
         latent_z = mean.squeeze(1)  
@@ -128,9 +128,9 @@ def text_from_latent_code(latent_z):
         length= length, # Chunyuan: Fix length; or use <EOS> to complete a sentence
         temperature=.5,
         top_k=100,
-        top_p=.95,
-        device='cuda',
-        decoder_tokenizer = tokenizer_decoder
+        top_p=.98,
+        device=DEVICE,
+        decoder_tokenizer=tokenizer_decoder
     )
     text_x1 = tokenizer_decoder.decode(out[0,:].tolist(), clean_up_tokenization_spaces=True)
     text_x1 = text_x1.split()[1:-1]
@@ -158,7 +158,7 @@ encoder_config_class, encoder_model_class, encoder_tokenizer_class = MODEL_CLASS
 model_encoder = encoder_model_class.from_pretrained(encoder_path, latent_size=latent_size)
 tokenizer_encoder = encoder_tokenizer_class.from_pretrained('bert-base-cased', do_lower_case=True)
 
-model_encoder.to('cuda')
+model_encoder.to(DEVICE)
 if block_size <= 0:
     block_size = tokenizer_encoder.max_len_single_sentence  # Our input block size will be the max possible for the model
 block_size = min(block_size, tokenizer_encoder.max_len_single_sentence)
@@ -167,7 +167,7 @@ block_size = min(block_size, tokenizer_encoder.max_len_single_sentence)
 decoder_config_class, decoder_model_class, decoder_tokenizer_class = MODEL_CLASSES['gpt2']
 model_decoder = decoder_model_class.from_pretrained(decoder_path, latent_size=latent_size)
 tokenizer_decoder = decoder_tokenizer_class.from_pretrained('gpt2', do_lower_case=False)
-model_decoder.to('cuda')
+model_decoder.to(DEVICE)
 if block_size <= 0:
     block_size = tokenizer_decoder.max_len_single_sentence  # Our input block size will be the max possible for the model
 block_size = min(block_size, tokenizer_decoder.max_len_single_sentence)
@@ -185,10 +185,10 @@ assert tokenizer_decoder.pad_token == '<PAD>'
 
 
 # Evaluation
-model_vae = VAE(model_encoder, model_decoder, tokenizer_encoder, tokenizer_decoder, SimpleNamespace(**{'latent_size': latent_size, 'device':'cuda'}))
+model_vae = VAE(model_encoder, model_decoder, tokenizer_encoder, tokenizer_decoder, SimpleNamespace(**{'latent_size': latent_size, 'device':DEVICE}))
 model_vae.load_state_dict(checkpoint['model_state_dict'])
 print("Pre-trained Optimus is successfully loaded")
-model_vae.to('cuda').to(torch.bfloat16)
+model_vae.to(DEVICE).to(torch.bfloat16)
 model_vae = torch.compile(model_vae)
 
 l = latent_code_from_text('A photo of a mountain.')[0]
@@ -222,17 +222,17 @@ def generate(prompt, in_embs=None,):
   if prompt != '':
     print(prompt)
     in_embs = in_embs / in_embs.abs().max() * .6 if in_embs != None else None
-    in_embs = 1 * in_embs.to('cuda') + 1 * latent_code_from_text(prompt)[0] if in_embs != None else latent_code_from_text(prompt)[0]
+    in_embs = 1 * in_embs.to(DEVICE) + 1 * latent_code_from_text(prompt)[0] if in_embs != None else latent_code_from_text(prompt)[0]
   else:
     print('From embeds.')
   in_embs = in_embs / in_embs.abs().max() * .6
-  in_embs = in_embs.to('cuda').to(torch.bfloat16)
+  in_embs = in_embs.to(DEVICE).to(torch.bfloat16)
   plt.close('all')
   plt.hist(np.array(in_embs.detach().to('cpu').to(torch.float)).flatten(), bins=5)
   plt.savefig('real_im_emb_plot.jpg')
     
   
-  text = text_from_latent_code(in_embs).replace('<unk> ', '')
+  text = ' '.join(text_from_latent_code(in_embs).replace( '<unk>', '').split())
   in_embs = latent_code_from_text(text)[0]
   print(text)
   return text, in_embs.to('cpu')

requirements.txt

@@ -1,4 +1,5 @@
 gradio
+boto3
 numpy
 scikit-learn
 pandas