Update app.py

app.py

@@ -11,7 +11,8 @@ import gradio as gr
 import torch
 from diffusers import FluxPipeline
 from PIL import Image
-
+from transformers import pipeline
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 # Hugging Face 토큰 설정
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN is None:
@@ -195,14 +196,29 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             )
             
             gallery.value = load_gallery()
-    
+
+
     @spaces.GPU
     def process_and_save_image(height, width, steps, scales, prompt, seed):
         global pipe
+        
+        # 한글 감지 및 번역
+        def contains_korean(text):
+            return any(ord('가') <= ord(c) <= ord('힣') for c in text)
+        
+        # 프롬프트 전처리
+        if contains_korean(prompt):
+            # 한글을 영어로 번역
+            translated = translator(prompt)[0]['translation_text']
+            prompt = translated
+        
+        # 프롬프트 형식 강제
+        formatted_prompt = f"wbgmsst, 3D, {prompt} ,white background"
+        
         with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
             try:
                 generated_image = pipe(
-                    prompt=[prompt],
+                    prompt=[formatted_prompt],  # 수정된 프롬프트 사용
                     generator=torch.Generator().manual_seed(int(seed)),
                     num_inference_steps=int(steps),
                     guidance_scale=float(scales),
@@ -220,6 +236,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 print(f"Error in image generation: {str(e)}")
                 return None, load_gallery()
     
+
+    
     def update_seed():
         return get_random_seed()