Update app.py

app.py

@@ -4,8 +4,18 @@ import random
 import spaces
 import torch
 import time
+import os
 from diffusers import DiffusionPipeline
 from custom_pipeline import FLUXPipelineWithIntermediateOutputs
+from transformers import pipeline
+
+# Hugging Face 토큰 가져오기
+hf_token = os.getenv("HF_TOKEN")
+if not hf_token:
+    raise ValueError("HF_TOKEN environment variable is not set. Please set it to your Hugging Face token.")
+
+# 번역 모델 로드 (토큰 인증 추가)
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", use_auth_token=hf_token)
 
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
@@ -17,14 +27,35 @@ DEFAULT_INFERENCE_STEPS = 1
 # Device and model setup
 dtype = torch.float16
 pipe = FLUXPipelineWithIntermediateOutputs.from_pretrained(
-    "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
+    "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype, use_auth_token=hf_token
 ).to("cuda")
 torch.cuda.empty_cache()
 
+# 한글 메뉴 이름 dictionary
+korean_labels = {
+    "Generated Image": "생성된 이미지",
+    "Prompt": "프롬프트",
+    "Enhance Image": "이미지 향상",
+    "Advanced Options": "고급 옵션",
+    "Seed": "시드",
+    "Randomize Seed": "시드 무작위화",
+    "Width": "너비",
+    "Height": "높이",
+    "Inference Steps": "추론 단계",
+    "Inspiration Gallery": "영감 갤러리"
+}
+
+def translate_if_korean(text):
+    if any('\u3131' <= char <= '\u3163' or '\uac00' <= char <= '\ud7a3' for char in text):
+        return translator(text, use_auth_token=hf_token)[0]['translation_text']
+    return text
+
 # Inference function
 @spaces.GPU(duration=25)
 def generate_image(prompt, seed=42, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=False, num_inference_steps=DEFAULT_INFERENCE_STEPS):
     
+    prompt = translate_if_korean(prompt)
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
@@ -40,18 +71,17 @@ def generate_image(prompt, seed=42, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT,
             height=height,
             generator=generator
         ): 
-        latency = f"Latency: {(time.time()-start_time):.2f} seconds"    
+        latency = f"처리 시간: {(time.time()-start_time):.2f} 초"    
         yield img, seed, latency
 
-
 # Example prompts
 examples = [
-    "a tiny astronaut hatching from an egg on the moon",
-    "a cat holding a sign that says hello world",
-    "an anime illustration of a wiener schnitzel",
-    "a futuristic cityscape with flying cars and neon lights",
-    "Photo of a young woman with long, wavy brown hair tied in a bun and glasses. She has a fair complexion and is wearing subtle makeup, emphasizing her eyes and lips. She is dressed in a black top. The background appears to be an urban setting with a building facade, and the sunlight casts a warm glow on her face.",
-    "Imagine steve jobs as Star Wars movie character"
+    "달에서 알에서 부화하는 작은 우주 비행사",
+    "안녕하세요 세상이라고 쓰인 표지판을 들고 있는 고양이",
+    "비너 슈니첼의 애니메이션 일러스트레이션",
+    "하늘을 나는 자동차와 네온 불빛이 있는 미래적인 도시 풍경",
+    "긴 갈색 웨이브 머리를 올려 묶고 안경을 쓴 젊은 여성의 사진. 그녀는 흰 피부에 눈과 입술을 강조한 은은한 화장을 했습니다. 그녀는 검은색 상의를 입었습니다. 배경은 도시 건물 외관으로 보이며, 햇빛이 그녀의 얼굴에 따뜻한 빛을 비추고 있습니다.",
+    "스티브 잡스를 스타워즈 영화 캐릭터로 상상해보세요"
 ]
 
 css = """
@@ -60,37 +90,36 @@ footer {
 }
 """
 
-
 # --- Gradio UI ---
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
     with gr.Column(elem_id="app-container"):
 
         with gr.Row():
             with gr.Column(scale=3):
-                result = gr.Image(label="Generated Image", show_label=False, interactive=False)
+                result = gr.Image(label=korean_labels["Generated Image"], show_label=False, interactive=False)
             with gr.Column(scale=1):
                 prompt = gr.Text(
-                    label="Prompt",
-                    placeholder="Describe the image you want to generate...",
+                    label=korean_labels["Prompt"],
+                    placeholder="생성하고 싶은 이미지를 설명하세요...",
                     lines=3,
                     show_label=False,
                     container=False,
                 )
-                enhanceBtn = gr.Button("🚀 Enhance Image")
+                enhanceBtn = gr.Button(f"🚀 {korean_labels['Enhance Image']}")
 
-                with gr.Column("Advanced Options"):
+                with gr.Column(korean_labels["Advanced Options"]):
                     with gr.Row():
                         latency = gr.Text(show_label=False)
                     with gr.Row():
-                        seed = gr.Number(label="Seed", value=42, precision=0)
-                        randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
+                        seed = gr.Number(label=korean_labels["Seed"], value=42, precision=0)
+                        randomize_seed = gr.Checkbox(label=korean_labels["Randomize Seed"], value=False)
                     with gr.Row():
-                        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_WIDTH)
-                        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_HEIGHT)
-                        num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=4, step=1, value=DEFAULT_INFERENCE_STEPS)
+                        width = gr.Slider(label=korean_labels["Width"], minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_WIDTH)
+                        height = gr.Slider(label=korean_labels["Height"], minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_HEIGHT)
+                        num_inference_steps = gr.Slider(label=korean_labels["Inference Steps"], minimum=1, maximum=4, step=1, value=DEFAULT_INFERENCE_STEPS)
 
         with gr.Row():
-            gr.Markdown("### 🌟 Inspiration Gallery")
+            gr.Markdown(f"### 🌟 {korean_labels['Inspiration Gallery']}")
         with gr.Row():
             gr.Examples(
                 examples=examples,
@@ -122,4 +151,4 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
     )
 
 # Launch the app
-demo.launch() 
+demo.launch()