Update app.py

app.py

@@ -19,20 +19,28 @@ from refiners.solutions import BoxSegmenter
 from transformers import GroundingDinoForObjectDetection, GroundingDinoProcessor
 from diffusers import FluxPipeline
 # 상단에 import 추가
-from transformers import pipeline
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+# 상단에 import 추가
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqGeneration
+
+# 번역 모델 초기화
+model_name = "Helsinki-NLP/opus-mt-ko-en"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqGeneration.from_pretrained(model_name)
+translator = pipeline("translation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
 
-# 번역 함수 추가
 def translate_to_english(text: str) -> str:
     """한글 텍스트를 영어로 번역"""
-    if any(ord('가') <= ord(char) <= ord('힣') for char in text):
-        try:
-            translated = translator(text)[0]['translation_text']
+    try:
+        if any(ord('가') <= ord(char) <= ord('힣') for char in text):
+            translated = translator(text, max_length=128)[0]['translation_text']
+            print(f"Translated '{text}' to '{translated}'")
             return translated
-        except Exception as e:
-            print(f"Translation error: {e}")
-            return text
-    return text
+        return text
+    except Exception as e:
+        print(f"Translation error: {str(e)}")
+        return text
+
+
 
 BoundingBox = tuple[int, int, int, int]