Update main.py

main.py

@@ -17,7 +17,8 @@ from imgurpython import ImgurClient
 import requests
 import tempfile
 from datetime import datetime
-from transformers import pipeline
+from transformers import pipeline,MarianMTModel, MarianTokenizer
+
 
 # 設定 Google AI API 金鑰
 genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
@@ -125,7 +126,17 @@ def generate_image_hf(prompt):
     else:
         print("圖片生成失敗:", response.text)
         return None
-        
+
+def translate_zh_to_en(text):
+    model_name = "Helsinki-NLP/opus-mt-zh-en"
+    tokenizer = MarianTokenizer.from_pretrained(model_name)
+    model = MarianMTModel.from_pretrained(model_name)
+    
+    inputs = tokenizer(text, return_tensors="pt", padding=True)
+    translated = model.generate(**inputs)
+    result = tokenizer.decode(translated[0], skip_special_tokens=True)
+    
+    return result
 #==========================
 # 使用者上傳圖片
 #==========================
@@ -198,8 +209,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_zh)
+        translation = translate_zh_to_en(prompt)
         prompt_en = translation[0]['translation_text']
 
         # 生成圖片