.

app.py

@@ -125,7 +125,7 @@ def text_from_latent_code(latent_z):
         model=model_vae.decoder,
         context=context_tokens,
         past=past,
-        length= length, # Chunyuan: Fix length; or use <EOS> to complete a sentence
+        length=length, # Chunyuan: Fix length; or use <EOS> to complete a sentence
         temperature=.5,
         top_k=100,
         top_p=.98,
@@ -190,6 +190,8 @@ model_vae.load_state_dict(checkpoint['model_state_dict'])
 print("Pre-trained Optimus is successfully loaded")
 model_vae.to(DEVICE).to(torch.bfloat16)
 model_vae = torch.compile(model_vae)
+model_vae.encoder = torch.compile(model_vae.encoder)
+model_vae.decoder = torch.compile(model_vae.decoder)
 
 l = latent_code_from_text('A photo of a mountain.')[0]
 t = text_from_latent_code(l)

requirements.txt

@@ -12,3 +12,4 @@ sentencepiece
 peft
 tensorflow_hub
 tensorflow==2.14.0
+sacremoses