Update app.py

app.py

@@ -5,7 +5,7 @@ import torch
 from PIL import Image
 import os
 
-from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
+from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor, pipeline
 from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256_ipadapter import StableDiffusionXLPipeline
 from kolors.models.modeling_chatglm import ChatGLMModel
 from kolors.models.tokenization_chatglm import ChatGLMTokenizer
@@ -53,6 +53,8 @@ if hasattr(pipe.unet, 'encoder_hid_proj'):
 
 pipe.load_ip_adapter(f'{root_dir}/weights/Kolors-IP-Adapter-Plus', subfolder="", weight_name=["ip_adapter_plus_general.bin"])
 
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
@@ -61,6 +63,9 @@ def infer(prompt, ip_adapter_image, ip_adapter_scale=0.5, negative_prompt="", se
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
+    # Translate prompt if it's in Korean
+    translated_prompt = translator(prompt, src_lang="ko", tgt_lang="en")[0]['translation_text']
+    
     generator = torch.Generator(device="cuda").manual_seed(seed)
     pipe.to("cuda")
     image_encoder.to("cuda")
@@ -68,7 +73,7 @@ def infer(prompt, ip_adapter_image, ip_adapter_scale=0.5, negative_prompt="", se
     pipe.set_ip_adapter_scale([ip_adapter_scale])
     
     image = pipe(
-        prompt=prompt,
+        prompt=translated_prompt,
         ip_adapter_image=[ip_adapter_image],
         negative_prompt=negative_prompt,
         height=height,
@@ -82,9 +87,9 @@ def infer(prompt, ip_adapter_image, ip_adapter_scale=0.5, negative_prompt="", se
     return image, seed
 
 examples = [
-    ["강아지", "minta.jpeg", 0.4],
-    ["카피바라라", "king-min.png", 0.5],
-    ["고양이", "blue_hair.png", 0.5],
+    ["A dog", "minta.jpeg", 0.4],
+    ["A capybara", "king-min.png", 0.5],
+    ["A cat", "blue_hair.png", 0.5],
     ["", "meow.jpeg", 1.0],
 ]