Update app.py

app.py

@@ -7,7 +7,7 @@ import gradio as gr
 import numpy as np
 import os
 from huggingface_hub import snapshot_download
-from transformers import CLIPVisionModelWithProjection,CLIPImageProcessor
+from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor, pipeline
 from SAK.pipelines.pipeline_stable_diffusion_xl_chatglm_256_ipadapter_FaceID import StableDiffusionXLPipeline
 from SAK.models.modeling_chatglm import ChatGLMModel
 from SAK.models.tokenization_chatglm import ChatGLMTokenizer
@@ -18,8 +18,6 @@ from PIL import Image
 from insightface.app import FaceAnalysis
 from insightface.data import get_image as ins_get_image
 
-
-
 device = "cuda"
 ckpt_dir = snapshot_download(repo_id="SunderAli17/SAK")
 ckpt_dir_faceid = snapshot_download(repo_id="SunderAli17/SAK-IP-Adapter-FaceTransform-Plus")
@@ -44,6 +42,16 @@ pipe = StableDiffusionXLPipeline(
     force_zeros_for_empty_prompt = False,
 )
 
+# 번역 모델 초기화
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+
+# 프롬프트 처리 함수 추가
+def process_prompt(prompt):
+    if any('\u3131' <= char <= '\u3163' or '\uac00' <= char <= '\ud7a3' for char in prompt):
+        translated = translator(prompt)[0]['translation_text']
+        return prompt, translated
+    return prompt, prompt
+
 class FaceInfoGenerator():
     def __init__(self, root_dir = "./.insightface/"):
         self.app = FaceAnalysis(name = 'antelopev2', root = root_dir, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
@@ -86,6 +94,7 @@ def infer(prompt,
           guidance_scale = 5.0, 
           num_inference_steps = 50
         ):
+    original_prompt, english_prompt = process_prompt(prompt)
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
@@ -104,11 +113,11 @@ def infer(prompt,
     face_embeds = face_embeds.to(device, dtype = torch.float16)
 
     image = pipe(
-        prompt = prompt,
+        prompt = english_prompt,
         negative_prompt = negative_prompt, 
         height = 1024,
         width = 1024,
-        num_inference_steps= num_inference_steps, 
+        num_inference_steps = num_inference_steps, 
         guidance_scale = guidance_scale,
         num_images_per_prompt = 1,
         generator = generator,
@@ -116,25 +125,18 @@ def infer(prompt,
         face_insightface_embeds = face_embeds
     ).images[0]
 
-    return image, seed
-
+    return image, seed, original_prompt, english_prompt
 
 examples = [
     ["wearing a full suit sitting in a restaurant with candle lights", "image/test0.png"],
     ["Wild cowboy hat with western town and horses in the background", "image/test1.png"]
 ]
 
-
-
-
 def load_description(fp):
     with open(fp, 'r', encoding='utf-8') as f:
         content = f.read()
     return content
 
-
-
-
 css = """
 footer {
     visibility: hidden;
@@ -142,66 +144,67 @@ footer {
 """
 
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as SAK:
-
     with gr.Row():
         with gr.Column(elem_id="col-left"):
             with gr.Row():
                 prompt = gr.Textbox(
-                    label="Prompt",
-                    placeholder="Enter your prompt",
+                    label="프롬프트",
+                    placeholder="프롬프트를 입력하세요 (한글 또는 영어)",
                     lines=2
                 )
             with gr.Row():
-                image = gr.Image(label="Image", type="pil")
-            with gr.Accordion("Advanced Settings", open=False):
+                image = gr.Image(label="이미지", type="pil")
+            with gr.Accordion("고급 설정", open=False):
                 negative_prompt = gr.Textbox(
-                    label="Negative prompt",
-                    placeholder="Enter a negative prompt",
+                    label="네거티브 프롬프트",
+                    placeholder="네거티브 프롬프트를 입력하세요",
                     visible=True,
+                    value="nsfw, 얼굴 그림자, 저해상도, jpeg 아티팩트, 흐릿함, 열악함, 네온 조명"
                 )
                 seed = gr.Slider(
-                    label="Seed",
+                    label="시드",
                     minimum=0,
                     maximum=MAX_SEED,
                     step=1,
                     value=0,
                 )
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                randomize_seed = gr.Checkbox(label="시드 무작위화", value=True)
                 with gr.Row():
                     guidance_scale = gr.Slider(
-                        label="Guidance scale",
+                        label="가이던스 스케일",
                         minimum=0.0,
                         maximum=10.0,
                         step=0.1,
                         value=5.0,
                     )
                     num_inference_steps = gr.Slider(
-                        label="Number of inference steps",
+                        label="추론 단계 수",
                         minimum=10,
                         maximum=50,
                         step=1,
                         value=25,
                     )
             with gr.Row():
-                button = gr.Button("Run", elem_id="button")
+                button = gr.Button("실행", elem_id="button")
             
         with gr.Column(elem_id="col-right"):
-            result = gr.Image(label="Result", show_label=False)
-            seed_used = gr.Number(label="Seed Used")
+            result = gr.Image(label="결과", show_label=False)
+            seed_used = gr.Number(label="사용된 시드")
+            original_prompt_display = gr.Textbox(label="원본 프롬프트")
+            english_prompt_display = gr.Textbox(label="영어 프롬프트")
     
     with gr.Row():
         gr.Examples(
                 fn = infer,
                 examples = examples,
                 inputs = [prompt, image],
-                outputs = [result, seed_used],
+                outputs = [result, seed_used, original_prompt_display, english_prompt_display],
             )
 
     button.click(
         fn = infer,
         inputs = [prompt, image, negative_prompt, seed, randomize_seed, guidance_scale, num_inference_steps],
-        outputs = [result, seed_used]
+        outputs = [result, seed_used, original_prompt_display, english_prompt_display]
     )
 
-
 SAK.queue().launch(debug=True, share=True)