Update app.py

app.py

@@ -20,8 +20,6 @@ from transformers import CLIPTextModel, CLIPTokenizer
 from transformers import T5EncoderModel, T5Tokenizer
 from transformers import pipeline, AutoTokenizer, MarianMTModel
 
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
-# ---------------- Encoders ----------------
 
 
 class HFEmbedder(nn.Module):
@@ -793,38 +791,62 @@ TRANSLATORS = {
 # 번역기 캐시 딕셔너리
 translators_cache = {}
 
+# 번역기 초기화 부분 수정
 def get_translator(lang):
+    if lang == "English":
+        return None
+        
     if lang not in translators_cache:
-        model_name = TRANSLATORS[lang]
         try:
+            model_name = TRANSLATORS[lang]
             tokenizer = AutoTokenizer.from_pretrained(model_name)
             model = MarianMTModel.from_pretrained(
                 model_name,
-                torch_dtype=torch.float16,
+                torch_dtype=torch.float32,  # float16 대신 float32 사용
                 low_cpu_mem_usage=True
-            )
-            translators_cache[lang] = pipeline(
-                "translation", 
-                model=model, 
-                tokenizer=tokenizer,
-                device=-1  # CPU 사용
-            )
+            ).to("cpu")  # 명시적으로 CPU 지정
+            
+            translators_cache[lang] = {
+                "model": model,
+                "tokenizer": tokenizer
+            }
         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)
+        translated = model.generate(**inputs)
+        result = tokenizer.decode(translated[0], skip_special_tokens=True)
+        
+        print(f"Original text: {text}")
+        print(f"Translated text: {result}")
+        
+        return result
+    except Exception as e:
+        print(f"Translation error: {e}")
+        return text
+
 def translate_prompt(prompt, source_lang):
     if source_lang == "English":
         return prompt
     
     translator = get_translator(source_lang)
     if translator is None:
-        print(f"Translation failed for {source_lang}, using original prompt")
+        print(f"No translator available for {source_lang}, using original prompt")
         return prompt
         
     try:
-        translated = translator(prompt, max_length=512)[0]['translation_text']
+        translated = translate_text(prompt, translator)
         print(f"Translated from {source_lang}: {translated}")
         return translated
     except Exception as e:
@@ -840,6 +862,7 @@ def generate_image(
 ):
     try:
         translated_prompt = translate_prompt(prompt, source_lang)
+        print(f"Using prompt: {translated_prompt}")
     except Exception as e:
         print(f"Translation failed: {e}")
         translated_prompt = prompt