Update app.py

app.py

@@ -753,8 +753,56 @@ ko_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 ja_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ja-en")
 zh_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-en")
 
-from transformers import MarianMTModel, MarianTokenizer
 
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
+
+def translate_text(text):
+    try:
+        # M2M100은 다국어 번역을 한 모델로 처리할 수 있습니다
+        model_name = "facebook/m2m100_418M"
+        tokenizer = M2M100Tokenizer.from_pretrained(model_name)
+        model = M2M100ForConditionalGeneration.from_pretrained(model_name).to(device)
+        
+        # 언어 감지
+        def detect_language(text):
+            if any('\u3131' <= c <= '\u318E' or '\uAC00' <= c <= '\uD7A3' for c in text):
+                return 'ko'
+            elif any('\u3040' <= c <= '\u309F' or '\u30A0' <= c <= '\u30FF' for c in text):
+                return 'ja'
+            elif any('\u4e00' <= c <= '\u9fff' for c in text):
+                return 'zh'
+            return None
+        
+        src_lang = detect_language(text)
+        if src_lang is None:
+            return text
+            
+        tokenizer.src_lang = src_lang
+        encoded = tokenizer(text, return_tensors="pt").to(device)
+        generated_tokens = model.generate(
+            **encoded,
+            forced_bos_token_id=tokenizer.get_lang_id("en"),
+            max_length=128
+        )
+        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    except Exception as e:
+        print(f"Translation error: {e}")
+        return text
+
+@spaces.GPU
+@torch.no_grad()
+def generate_image(
+    prompt, width, height, guidance, inference_steps, seed,
+    do_img2img, init_image, image2image_strength, resize_img,
+    progress=gr.Progress(track_tqdm=True),
+):
+    translated_prompt = translate_text(prompt)
+    if translated_prompt != prompt:
+        print(f"Translated prompt: {translated_prompt}")
+    prompt = translated_prompt
+
+    if seed == 0:
+        seed = int(random.random() * 1000000)
 def translate_text(text, src_lang, model_name):
     try:
         tokenizer = MarianTokenizer.from_pretrained(model_name)
@@ -778,46 +826,6 @@ def translate_if_needed(prompt):
         return translate_text(prompt, 'zh', 'Helsinki-NLP/opus-mt-zh-en')
     return prompt
 
-@spaces.GPU
-@torch.no_grad()
-def generate_image(
-    prompt, width, height, guidance, inference_steps, seed,
-    do_img2img, init_image, image2image_strength, resize_img,
-    progress=gr.Progress(track_tqdm=True),
-):
-    translated_prompt = translate_if_needed(prompt)
-    if translated_prompt != prompt:
-        print(f"Translated prompt: {translated_prompt}")
-    prompt = translated_prompt
-
-    
-    # 한글, 일본어, 중국어 문자 감지
-    def contains_korean(text):
-        return any('\u3131' <= c <= '\u318E' or '\uAC00' <= c <= '\uD7A3' for c in text)
-        
-    def contains_japanese(text):
-        return any('\u3040' <= c <= '\u309F' or '\u30A0' <= c <= '\u30FF' or '\u4E00' <= c <= '\u9FFF' for c in text)
-    
-    def contains_chinese(text):
-        return any('\u4e00' <= c <= '\u9fff' for c in text)
-    
-    # 한글, 일본어, 중국어가 있으면 번역
-    if contains_korean(prompt):
-        translated_prompt = ko_translator(prompt, max_length=512)[0]['translation_text']
-        print(f"Translated Korean prompt: {translated_prompt}")
-        prompt = translated_prompt
-    elif contains_japanese(prompt):
-        translated_prompt = ja_translator(prompt, max_length=512)[0]['translation_text'] 
-        print(f"Translated Japanese prompt: {translated_prompt}")
-        prompt = translated_prompt
-    elif contains_chinese(prompt):
-        translated_prompt = zh_translator(prompt, max_length=512)[0]['translation_text']
-        print(f"Translated Chinese prompt: {translated_prompt}")
-        prompt = translated_prompt
-
-
-
-    
 
     if seed == 0:
         seed = int(random.random() * 1000000)