Update app.py

app.py

@@ -49,9 +49,18 @@ def setup_environment():
         raise ValueError("HF_TOKEN not found in environment variables")
     login(token=HF_TOKEN)
     return HF_TOKEN
+
 @spaces.GPU()
 def generate_image(prompt, mode, cfg_scale=7.0, steps=30, seed=None, width=512, height=768, lora_scale=0.85):
     try:
+        # 한글 처리
+        if contains_korean(prompt):
+            translator = get_translator()
+            translated = translator(prompt)[0]['translation_text']
+            actual_prompt = translated
+        else:
+            actual_prompt = prompt
+
         # 파이프라인 초기화
         pipe = DiffusionPipeline.from_pretrained(
             BASE_MODEL,
@@ -73,13 +82,13 @@ def generate_image(prompt, mode, cfg_scale=7.0, steps=30, seed=None, width=512,
         # 이미지 생성
         with torch.inference_mode():
             result = pipe(
-                prompt=f"{prompt} {trigger_word}",
+                prompt=f"{actual_prompt} {trigger_word}",
                 num_inference_steps=steps,
                 guidance_scale=cfg_scale,
                 width=width,
                 height=height,
                 generator=generator,
-                cross_attention_kwargs={"scale": lora_scale},
+                joint_attention_kwargs={"scale": lora_scale},
             ).images[0]
 
         # 메모리 정리
@@ -92,6 +101,13 @@ def generate_image(prompt, mode, cfg_scale=7.0, steps=30, seed=None, width=512,
         clear_memory()
         raise gr.Error(f"Generation failed: {str(e)}")
 
+def contains_korean(text):
+    return any(ord('가') <= ord(char) <= ord('힣') for char in text)
+
+@spaces.GPU()
+def get_translator():
+    return pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device="cuda")
+
 
 # 전역 변수 초기화
 fashion_pipe = None