Update app.py

app.py

@@ -4,11 +4,15 @@ import torch
 from PIL import Image
 from diffusers import DiffusionPipeline
 import random
+from transformers import pipeline
 
 torch.backends.cudnn.deterministic = True
 torch.backends.cudnn.benchmark = False
 torch.backends.cuda.matmul.allow_tf32 = True
 
+# 번역 모델 초기화
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+
 base_model = "black-forest-labs/FLUX.1-dev"
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
 
@@ -21,7 +25,18 @@ pipe.to("cuda")
 MAX_SEED = 2**32-1
 
 @spaces.GPU()
-def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+def translate_and_generate(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+    # 한글 감지 및 번역
+    def contains_korean(text):
+        return any(ord('가') <= ord(char) <= ord('힣') for char in text)
+    
+    if contains_korean(prompt):
+        # 한글을 영어로 번역
+        translated = translator(prompt)[0]['translation_text']
+        actual_prompt = translated
+    else:
+        actual_prompt = prompt
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device="cuda").manual_seed(seed)
@@ -33,7 +48,7 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
             progress(i / steps * 100, f"Processing step {i} of {steps}...")
 
     image = pipe(
-        prompt=f"{prompt} {trigger_word}",
+        prompt=f"{actual_prompt} {trigger_word}",
         num_inference_steps=steps,
         guidance_scale=cfg_scale,
         width=width,
@@ -75,8 +90,8 @@ with gr.Blocks(css=css) as app:
             with gr.Column(scale=3):
                 with gr.Group(elem_classes="parameter-box"):
                     prompt = gr.TextArea(
-                        label="✍️ Your Prompt",
-                        placeholder="Describe the image you want to generate...",
+                        label="✍️ Your Prompt (한글 또는 영어)",
+                        placeholder="이미지를 설명하세요... (한글 입력시 자동으로 영어로 번역됩니다)",
                         lines=5
                     )
                     
@@ -156,7 +171,7 @@ with gr.Blocks(css=css) as app:
     )
     
     generate_button.click(
-        run_lora,
+        translate_and_generate,
         inputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale],
         outputs=[result, seed]
     )