Update app.py

app.py

@@ -12,15 +12,17 @@ from rudalle import get_rudalle_model, get_tokenizer, get_vae
 tqdm.__init__ = partialmethod(tqdm.__init__, disable=True)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-translation_model = FSMTForConditionalGeneration.from_pretrained("facebook/wmt19-en-ru", torch_dtype=torch.float16)
+translation_model = FSMTForConditionalGeneration.from_pretrained("facebook/wmt19-en-ru", torch_dtype=torch.float16).to(device)
 tokenizer = FSMTTokenizer.from_pretrained("facebook/wmt19-en-ru")
-translation_pipe = pipeline("translation", model=translation_model, tokenizer=tokenizer, device=0)
 dalle = get_rudalle_model("Malevich", pretrained=True, fp16=True, device=device)
 tokenizer = get_tokenizer()
 vae = get_vae().to(device)
 
 def translation_wrapper(text: str):
-    return translation_pipe(text)[0]["translation_text"]
+    input_ids = tokenizer.encode(text, return_tensors="pt")
+    outputs = translation_model.generate(input_ids)
+    decoded = tokenizer.decode(outputs[0].float(), skip_special_tokens=True)
+    return decoded
 
 def dalle_wrapper(prompt: str):
     top_k, top_p = random.choice([