Update main.py

main.py

@@ -18,7 +18,7 @@ import requests
 import tempfile
 from datetime import datetime
 
-from transformers import AutoModel, AutoTokenizer, MarianMTModel
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from transformers import pipeline
 
 
@@ -129,6 +129,17 @@ def generate_image_hf(prompt):
         print("圖片生成失敗:", response.text)
         return None
 
+
+
+def translate_zh_to_en(text):
+    model_name = "Helsinki-NLP/opus-mt-zh-en"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    
+    inputs = tokenizer(text, return_tensors="pt")
+    outputs = model.generate(**inputs)
+    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return translated_text
     
 #==========================
 # 使用者上傳圖片
@@ -202,10 +213,7 @@ def handle_image_message(event):
         line_bot_api.reply_message(event.reply_token, TextSendMessage(text="圖片生成中～ 請稍候.....✨"))
         client = ImgurClient(client_id, client_secret, access_token, refresh_token)
 
-        translator = pipeline("translation", model='Helsinki-NLP/opus-mt-zh-en') 
-        translation = translator(prompt)
-        prompt_en = translation[0]['translation_text']
-
+        prompt_en = translate_zh_to_en(prompt)
         # 生成圖片
         image_binary = generate_image_hf(prompt_en)