Update app.py

app.py

@@ -8,20 +8,30 @@ from huggingface_hub import InferenceClient
 import subprocess
 import torch
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForCausalLM
-import random
+from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
 
+# 설치 과정은 생략 가능하며 필요한 경우에만 실행
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
+# Hugging Face 토큰 설정
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
 
+# 번역 모델 추가
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
-# Initialize Florence model
+# Florence 모델 초기화
 device = "cuda" if torch.cuda.is_available() else "cpu"
 florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to(device).eval()
 florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
 
-# Florence caption function
+# 한글 프롬프트 번역 함수
+def translate_prompt(prompt):
+    if any("\uAC00" <= char <= "\uD7A3" for char in prompt):  # 한글이 포함된 경우
+        translated = translator(prompt, max_length=512)[0]['translation_text']
+        return translated
+    return prompt
+
+# Florence 캡션 함수
 def florence_caption(image):
     if not isinstance(image, Image.Image):
         image = Image.fromarray(image)
@@ -42,13 +52,14 @@ def florence_caption(image):
         image_size=(image.width, image.height)
     )
     return parsed_answer["<MORE_DETAILED_CAPTION>"]
-    
-# Load JSON files
+
+# JSON 파일 로드 함수
 def load_json_file(file_name):
     file_path = os.path.join("data", file_name)
     with open(file_path, "r") as file:
         return json.load(file)
 
+# JSON 데이터 로드
 ARTFORM = load_json_file("artform.json")
 PHOTO_TYPE = load_json_file("photo_type.json")
 BODY_TYPES = load_json_file("body_types.json")
@@ -68,6 +79,7 @@ COMPOSITION = load_json_file("composition.json")
 POSE = load_json_file("pose.json")
 BACKGROUND = load_json_file("background.json")
 
+# 프롬프트 생성 클래스
 class PromptGenerator:
     def __init__(self, seed=None):
         self.rng = random.Random(seed)
@@ -147,6 +159,7 @@ class PromptGenerator:
         components = []
         custom = kwargs.get("custom", "")
         if custom:
+            custom = translate_prompt(custom)  # 번역 적용
             components.append(custom)
         is_photographer = kwargs.get("artform", "").lower() == "photography" or (
             kwargs.get("artform", "").lower() == "random"
@@ -390,13 +403,19 @@ You are allowed to make up film and branding names, and do them like 80's, 90's
             print(f"An error occurred: {e}")
             return f"Error occurred while processing the request: {str(e)}"
 
+css = """
+footer {
+    visibility: hidden;
+}
+"""
+
 def create_interface():
     prompt_generator = PromptGenerator()
     huggingface_node = HuggingFaceInferenceNode()
 
-    with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
+    with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
         
-        gr.HTML("<h1>My AI Art Prompt Generator</h1>")
+        gr.HTML("<h1>Flux Prompt Generator</h1>")
 
         with gr.Row():
             with gr.Column(scale=2):
@@ -533,3 +552,4 @@ def create_interface():
 if __name__ == "__main__":
     demo = create_interface()
     demo.launch()
+