Update app.py

app.py

@@ -4,11 +4,15 @@ import random
 import torch
 from diffusers import DiffusionPipeline
 import spaces
+from transformers import pipeline
 
 # 기본 설정
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# 한국어-영어 번역 모델 로드 (CPU에서 실행)
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device="cpu")
+
 # 모델 로드
 pipe = DiffusionPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell",
@@ -78,11 +82,32 @@ GRADIO_EXAMPLES = [
     for example in EXAMPLES
 ]
 
+# 한국어 감지 함수
+def contains_korean(text):
+    for char in text:
+        if ord('가') <= ord(char) <= ord('힣'):
+            return True
+    return False
+
+# 필요시 번역 후 추론 함수
 @spaces.GPU()
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
+    # 한국어 감지 및 번역
+    original_prompt = prompt
+    translated = False
+    
+    if contains_korean(prompt):
+        translated = True
+        translation = translator(prompt)
+        prompt = translation[0]['translation_text']
+    
+    # 랜덤 시드 설정
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+    
     generator = torch.Generator().manual_seed(seed)
+    
+    # 모델 실행
     image = pipe(
         prompt=prompt,
         width=width,
@@ -91,7 +116,12 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
         generator=generator,
         guidance_scale=0.0
     ).images[0]
-    return image, seed
+    
+    # 번역 정보 반환
+    if translated:
+        return image, seed, original_prompt, prompt
+    else:
+        return image, seed, None, None
 
 # CSS 스타일 (기존 구조 유지)
 css = """
@@ -164,6 +194,14 @@ css = """
 .examples-area {
     flex: 1 !important;
 }
+.translation-info {
+    background-color: #f8f9fa;
+    border-left: 4px solid #17a2b8;
+    padding: 10px 15px;
+    margin-top: 10px;
+    border-radius: 4px;
+    font-size: 14px;
+}
 """
 
 with gr.Blocks(css=css) as demo:
@@ -178,13 +216,25 @@ with gr.Blocks(css=css) as demo:
         with gr.Column(elem_id="input-column", scale=2):
             with gr.Group(elem_classes="input-box"):
                 prompt = gr.Text(
-                    label="Design Prompt",
-                    placeholder="Enter your product design concept details...",
+                    label="Design Prompt (한국어 또는 영어로 입력하세요)",
+                    placeholder="Enter your product design concept details in Korean or English...",
                     lines=10,
                     elem_classes="prompt-input"
                 )
                 run_button = gr.Button("Generate Design", variant="primary")
                 result = gr.Image(label="Generated Design")
+                
+                # 번역 정보 표시 영역
+                original_prompt = gr.Textbox(visible=False)
+                translated_prompt = gr.Textbox(visible=False)
+                translation_info = gr.Markdown(visible=False, elem_classes="translation-info")
+                
+                # 번역 정보 업데이트 함수
+                def update_translation_info(original, translated):
+                    if original and translated:
+                        return gr.update(visible=True, value=f"🔄 Korean prompt was translated to English:\n\n**Original:** {original}\n\n**Translated:** {translated}")
+                    else:
+                        return gr.update(visible=False)
             
             with gr.Accordion("Advanced Settings", open=False):
                 seed = gr.Slider(
@@ -247,7 +297,15 @@ with gr.Blocks(css=css) as demo:
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps],
-        outputs=[result, seed]
+        outputs=[result, seed, original_prompt, translated_prompt]
+    )
+    
+    # 번역 정보 업데이트 이벤트
+    gr.on(
+        triggers=[original_prompt.change, translated_prompt.change],
+        fn=update_translation_info,
+        inputs=[original_prompt, translated_prompt],
+        outputs=[translation_info]
     )
 
 if __name__ == "__main__":
@@ -258,4 +316,4 @@ if __name__ == "__main__":
         share=False,
         show_error=True,
         debug=True
-    )
+    )