Update app.py

app.py

@@ -22,6 +22,24 @@ TMP_DIR = "/tmp/Trellis-demo"
 
 os.makedirs(TMP_DIR, exist_ok=True)
 
+# 상단의 import 부분에 추가
+from transformers import pipeline
+
+# 번역기 초기화
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+
+# 한글 감지 및 번역 함수
+def translate_korean_prompt(prompt: str) -> str:
+    """
+    한글이 포함된 프롬프트를 영어로 번역
+    """
+    def contains_korean(text):
+        return any(ord('가') <= ord(c) <= ord('힣') for c in text)
+    
+    if contains_korean(prompt):
+        translated = translator(prompt)[0]['translation_text']
+        return translated
+    return prompt
 
 def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
     """
@@ -78,7 +96,6 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     
     return gs, mesh, state['trial_id']
 
-
 @spaces.GPU
 def text_to_image(prompt: str, seed: int, randomize_seed: bool) -> Image.Image:
     """
@@ -87,8 +104,14 @@ def text_to_image(prompt: str, seed: int, randomize_seed: bool) -> Image.Image:
     if randomize_seed:
         seed = np.random.randint(0, MAX_SEED)
     
+    # 한글 프롬프트를 영어로 번역
+    english_prompt = translate_korean_prompt(prompt)
+    
+    # 프롬프트 형식 강제
+    formatted_prompt = f"wbgmsst, 3D, {english_prompt}, white background"
+    
     generator = torch.Generator(device="cuda").manual_seed(seed)
-    image = text2img_pipeline(prompt, generator=generator).images[0]
+    image = text2img_pipeline(formatted_prompt, generator=generator).images[0]
     return image