Update app.py

app.py

@@ -22,8 +22,8 @@ from torch import Tensor, nn
 from transformers import CLIPTextModel, CLIPTokenizer
 from transformers import T5EncoderModel, T5Tokenizer
 # from optimum.quanto import freeze, qfloat8, quantize
-
-
+from transformers import pipeline
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 # ---------------- Encoders ----------------
 
 
@@ -759,6 +759,12 @@ def generate_image(
     do_img2img, init_image, image2image_strength, resize_img,
     progress=gr.Progress(track_tqdm=True),
 ):
+    translated_prompt = prompt
+    if any('\u3131' <= c <= '\u318E' or '\uAC00' <= c <= '\uD7A3' for c in prompt):
+        translated_prompt = translator(prompt, max_length=512)[0]['translation_text']
+        print(f"Translated prompt: {translated_prompt}")
+        prompt = translated_prompt     
+
     if seed == 0:
         seed = int(random.random() * 1000000)
         
@@ -809,8 +815,9 @@ def generate_image(
     x = rearrange(x[0], "c h w -> h w c")
     img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
     
-    return img, seed
-
+ 
+    return img, seed, translated_prompt
+    
 css = """
 footer {
     visibility: hidden;