Update app.py

app.py

@@ -777,6 +777,8 @@ TRANSLATORS = {
 
 translators_cache = {}
 
+from transformers import MarianMTModel, MarianTokenizer
+
 def get_translator(lang):
     """단일 번역기를 초기화하고 반환하는 함수"""
     if lang == "English":
@@ -785,12 +787,16 @@ def get_translator(lang):
     if lang not in translators_cache:
         try:
             model_name = TRANSLATORS[lang]
-            translator = pipeline(
-                "translation",
-                model=model_name,
-                device="cpu"  # CPU로 고정
-            )
-            translators_cache[lang] = translator
+            tokenizer = MarianTokenizer.from_pretrained(model_name)
+            model = MarianMTModel.from_pretrained(model_name)
+            
+            # CPU에서 실행
+            model = model.to("cpu")
+            
+            translators_cache[lang] = {
+                "model": model,
+                "tokenizer": tokenizer
+            }
             print(f"Successfully loaded translator for {lang}")
         except Exception as e:
             print(f"Error loading translator for {lang}: {e}")
@@ -803,64 +809,30 @@ def translate_prompt(prompt, source_lang):
     if source_lang == "English":
         return prompt
     
-    translator = get_translator(source_lang)
-    if translator is None:
+    translator_info = get_translator(source_lang)
+    if translator_info is None:
         print(f"No translator available for {source_lang}, using original prompt")
         return prompt
     
-    try:
-        translation = translator(prompt, max_length=512)
-        translated_text = translation[0]['translation_text']
-        print(f"Original ({source_lang}): {prompt}")
-        print(f"Translated: {translated_text}")
-        return translated_text
-    except Exception as e:
-        print(f"Translation error for {source_lang}: {e}")
-        return prompt
-
-        
-def get_translator(lang):
-    if lang == "English":
-        return None
-        
-    if lang not in translators_cache:
-        try:
-            model_name = TRANSLATORS[lang]
-            tokenizer = MarianTokenizer.from_pretrained(model_name)
-            model = MarianMTModel.from_pretrained(model_name).to("cpu")
-            
-            translators_cache[lang] = pipeline(
-                "translation",
-                model=model,
-                tokenizer=tokenizer,
-                device="cpu"
-            )
-        except Exception as e:
-            print(f"Error loading translator for {lang}: {e}")
-            return None
-            
-    return translators_cache[lang]
-
-def translate_text(text, translator_info):
-    if translator_info is None:
-        return text
-        
     try:
         tokenizer = translator_info["tokenizer"]
         model = translator_info["model"]
         
-        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        # 텍스트를 토큰화하고 모델 입력으로 변환
+        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        
+        # 번역 수행
         translated = model.generate(**inputs)
-        result = tokenizer.decode(translated[0], skip_special_tokens=True)
         
-        print(f"Original text: {text}")
-        print(f"Translated text: {result}")
+        # 번역된 텍스트 디코딩
+        translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
         
-        return result
+        print(f"Original ({source_lang}): {prompt}")
+        print(f"Translated: {translated_text}")
+        return translated_text
     except Exception as e:
-        print(f"Translation error: {e}")
-        return text
-
+        print(f"Translation error for {source_lang}: {e}")
+        return prompt
 
 @spaces.GPU
 @torch.no_grad()
@@ -869,7 +841,6 @@ def generate_image(
     do_img2img, init_image, image2image_strength, resize_img,
     progress=gr.Progress(track_tqdm=True),
 ):
-    # 번역 처리
     try:
         if source_lang != "English":
             translated_prompt = translate_prompt(prompt, source_lang)
@@ -880,6 +851,7 @@ def generate_image(
         print(f"Translation failed: {e}")
         translated_prompt = prompt
 
+
         
     if seed == 0:
         seed = int(random.random() * 1000000)