Update app.py

app.py

@@ -753,7 +753,30 @@ 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
 
+def translate_text(text, src_lang, model_name):
+    try:
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        model = MarianMTModel.from_pretrained(model_name)
+        model = model.to(device)
+        
+        inputs = tokenizer(text, return_tensors="pt", padding=True).to(device)
+        translated = model.generate(**inputs)
+        translated_text = tokenizer.batch_decode(translated, skip_special_tokens=True)[0]
+        return translated_text
+    except:
+        return text  # 번역 실패시 원본 텍스트 반환
+
+# 기존의 translator 정의 부분을 삭제하고 아래 코드로 대체
+def translate_if_needed(prompt):
+    if any('\u3131' <= c <= '\u318E' or '\uAC00' <= c <= '\uD7A3' for c in prompt):  # Korean
+        return translate_text(prompt, 'ko', 'Helsinki-NLP/opus-mt-ko-en')
+    elif any('\u3040' <= c <= '\u309F' or '\u30A0' <= c <= '\u30FF' for c in prompt):  # Japanese
+        return translate_text(prompt, 'ja', 'Helsinki-NLP/opus-mt-ja-en')
+    elif any('\u4e00' <= c <= '\u9fff' for c in prompt):  # Chinese
+        return translate_text(prompt, 'zh', 'Helsinki-NLP/opus-mt-zh-en')
+    return prompt
 
 @spaces.GPU
 @torch.no_grad()
@@ -762,7 +785,11 @@ def generate_image(
     do_img2img, init_image, image2image_strength, resize_img,
     progress=gr.Progress(track_tqdm=True),
 ):
-    translated_prompt = prompt
+    translated_prompt = translate_if_needed(prompt)
+    if translated_prompt != prompt:
+        print(f"Translated prompt: {translated_prompt}")
+    prompt = translated_prompt
+
     
     # 한글, 일본어, 중국어 문자 감지
     def contains_korean(text):
@@ -887,16 +914,14 @@ def create_demo():
             outputs=[init_image, image2image_strength, resize_img]
         )
 
-        examples = [
-            ["A magical fairy garden with glowing mushrooms and floating lanterns"],  # English
-            ["아름다운 벚꽃이 흩날리는 한옥 정원에서 한복을 입은 소녀"],  # Korean
-            ["夕暮れの富士山と桜の木の下で休んでいる可愛い柴犬"],  # Japanese
-            ["古老的中国庭园里，一只熊猫正在竹林中悠闲地吃着竹子"]   # Chinese
-        ]
-        
         gr.Examples(
-            examples=examples,
-            inputs=prompt,
+            examples=[
+                ["A magical fairy garden with glowing mushrooms and floating lanterns", 768, 768, 3.5, 30, 0, False, None, 0.8, True],  # English
+                ["아름다운 벚꽃이 흩날리는 한옥 정원에서 한복을 입은 소녀", 768, 768, 3.5, 30, 0, False, None, 0.8, True],  # Korean
+                ["夕暮れの富士山と桜の木の下で休んでいる可愛い柴犬", 768, 768, 3.5, 30, 0, False, None, 0.8, True],  # Japanese
+                ["古老的中国庭园里，一只熊猫正在竹林中悠闲地吃着竹子", 768, 768, 3.5, 30, 0, False, None, 0.8, True]   # Chinese
+            ],
+            inputs=[prompt, width, height, guidance, inference_steps, seed, do_img2img, init_image, image2image_strength, resize_img],
             outputs=[output_image, output_seed],
             fn=generate_image,
             cache_examples=True