Update app.py

app.py

@@ -1,19 +1,37 @@
 import random
 import os
 import uuid
+import re
+import time
 from datetime import datetime
+
 import gradio as gr
 import numpy as np
-import spaces
+import requests
 import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
 
-# Create permanent storage directory
+# ===== OpenAI 설정 =====
+from openai import OpenAI
+
+client = OpenAI(api_key=os.getenv("LLM_API"))  # 환경 변수에 API 키가 있어야 합니다.
+
+# ===== 프롬프트 증강용 스타일 프리셋 =====
+STYLE_PRESETS = {
+    "None": "",
+    "Realistic Photo": "photorealistic, 8k, ultra-detailed, cinematic lighting, realistic skin texture",
+    "Oil Painting": "oil painting, rich brush strokes, canvas texture, baroque lighting",
+    "Comic Book": "comic book style, bold ink outlines, cel shading, vibrant colors",
+    "Watercolor": "watercolor illustration, soft gradients, splatter effect, pastel palette",
+}
+
+# ===== 저장 폴더 =====
 SAVE_DIR = "saved_images"  # Gradio will handle the persistence
 if not os.path.exists(SAVE_DIR):
     os.makedirs(SAVE_DIR, exist_ok=True)
 
+# ===== 디바이스 & 모델 로드 =====
 device = "cuda" if torch.cuda.is_available() else "cpu"
 repo_id = "black-forest-labs/FLUX.1-dev"
 adapter_id = "seawolf2357/kim-korea"  # 특정 정치인을 학습한 LoRA 모델
@@ -25,171 +43,142 @@ pipeline = pipeline.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def save_generated_image(image, prompt):
-    # Generate unique filename with timestamp
+# ===== 한글 여부 판별 =====
+HANGUL_RE = re.compile(r"[\u3131-\u318E\uAC00-\uD7A3]+")
+
+def is_korean(text: str) -> bool:
+    return bool(HANGUL_RE.search(text))
+
+# ===== 번역 & 증강 함수 =====
+
+def openai_translate(text: str, retries: int = 3) -> str:
+    """한글을 영어로 번역 (OpenAI GPT-4.1-mini 사용). 영어 입력이면 그대로 반환."""
+    if not is_korean(text):
+        return text
+
+    for attempt in range(retries):
+        try:
+            res = client.chat.completions.create(
+                model="gpt-4.1-mini",
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "Translate the following Korean prompt into concise, descriptive English suitable for an image generation model. Keep the meaning, do not add new concepts."
+                    },
+                    {"role": "user", "content": text}
+                ],
+                temperature=0.3,
+                max_tokens=256,
+            )
+            return res.choices[0].message.content.strip()
+        except (requests.exceptions.RequestException, Exception) as e:
+            print(f"[translate] attempt {attempt + 1} failed: {e}")
+            time.sleep(2)
+    return text  # 번역 실패 시 원문 그대로
+
+def prepare_prompt(user_prompt: str, style_key: str) -> str:
+    """한글이면 번역하고, 선택한 스타일 프리셋을 붙여서 최종 프롬프트를 만든다."""
+    prompt_en = openai_translate(user_prompt)
+    style_suffix = STYLE_PRESETS.get(style_key, "")
+    if style_suffix:
+        final_prompt = f"{prompt_en}, {style_suffix}"
+    else:
+        final_prompt = prompt_en
+    return final_prompt
+
+# ===== 이미지 저장 =====
+
+def save_generated_image(image: Image.Image, prompt: str) -> str:
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     unique_id = str(uuid.uuid4())[:8]
     filename = f"{timestamp}_{unique_id}.png"
     filepath = os.path.join(SAVE_DIR, filename)
-    
-    # Save the image
     image.save(filepath)
-    
-    # Save metadata
+
+    # 메타데이터 저장
     metadata_file = os.path.join(SAVE_DIR, "metadata.txt")
     with open(metadata_file, "a", encoding="utf-8") as f:
         f.write(f"{filename}|{prompt}|{timestamp}\n")
-    
     return filepath
 
-@spaces.GPU(duration=60)
-def inference(
-    prompt,
-    seed=42,
-    randomize_seed=True,
-    width=1024,
-    height=768,
-    guidance_scale=3.5,
-    num_inference_steps=30,
-    lora_scale=1.0,
-    progress=None,
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+# ===== Diffusion 호출 =====
+
+def run_pipeline(prompt: str, seed: int, width: int, height: int, guidance_scale: float, num_steps: int, lora_scale: float):
     generator = torch.Generator(device=device).manual_seed(int(seed))
-    
-    image = pipeline(
+    result = pipeline(
         prompt=prompt,
         guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
+        num_inference_steps=num_steps,
         width=width,
         height=height,
         generator=generator,
         joint_attention_kwargs={"scale": lora_scale},
     ).images[0]
-    
-    # Save the generated image
-    filepath = save_generated_image(image, prompt)
-    
-    # Return just the image and seed
-    return image, seed
+    return result
 
-# 예시 문구: 특정 정치인 Mr. KIM의 다양한 상황을 묘사
+# ===== Gradio inference 래퍼 =====
 
+@spaces.GPU(duration=60)
+def generate_image(
+    user_prompt: str,
+    style_key: str,
+    seed: int = 42,
+    randomize_seed: bool = True,
+    width: int = 1024,
+    height: int = 768,
+    guidance_scale: float = 3.5,
+    num_inference_steps: int = 30,
+    lora_scale: float = 1.0,
+    progress=None,
+):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    # 1) 번역 + 증강
+    final_prompt = prepare_prompt(user_prompt, style_key)
+
+    # 2) 파이프라인 호출
+    image = run_pipeline(final_prompt, seed, width, height, guidance_scale, num_inference_steps, lora_scale)
+
+    # 3) 저장
+    save_generated_image(image, final_prompt)
+
+    return image, seed
+
+# ===== 예시 프롬프트 (한국어/영어 혼용 허용) =====
 examples = [
-    "Mr. KIM holding up a 'Fighting!' banner with both hands, showing patriotic pride and determination for national excellence. ",   
+    "김 후보가 태극기를 들고 힘찬 미소를 짓는 모습을 8K로",   # 한글 예시 (자동 번역)
     "Mr. KIM raising both arms in celebration with a triumphant expression, showing victory and hope for the future.",
-    "Mr. KIM jogging in a park wearing athletic gear, demonstrating healthy lifestyle and energetic leadership qualities.",  
-    "Mr. KIM warmly shaking hands with female citizens in a crowded street, showing genuine care and connection with women voters. ",
-    "Mr. KIM at a campaign rally, pointing toward the horizon with an inspiring gesture while female and kids audience members applaud.  ",
-    "Mr. KIM participating in a community event, surrounded by enthusiastic female supporters cheering ",
-    "Mr. KIM visiting a local market, engaging in friendly conversation with female vendors and shopkeepers. ",
-    "Mr. KIM walking through a university campus, discussing education policies with female students and professors. ",    
-    "Mr. KIM delivering a powerful speech in front of a large crowd with confident gestures and determined expression. ",
-    "Mr. KIM in a dynamic interview setting, passionately outlining his visions for the future.",
-    "Mr. KIM preparing for an important debate, surrounded by paperwork, looking focused and resolute. ",
+    "김 후보가 공원에서 조깅 중 건강한 리더십을 보여주는 장면",  # 한글 예시
 ]
 
-# UI를 붉은 계열 그라디에이션으로 디자인
+# ===== 커스텀 CSS (붉은 톤 유지) =====
 custom_css = """
 :root {
-    --color-primary: #8F1A3A; /* 붉은 톤의 메인 컬러 */
-    --color-secondary: #FF4B4B; /* 포인트 컬러(밝은 빨강) */
+    --color-primary: #8F1A3A;
+    --color-secondary: #FF4B4B;
     --background-fill-primary: linear-gradient(to right, #FFF5F5, #FED7D7, #FEB2B2);
 }
-
-footer {
-    visibility: hidden;
-}
-
-.gradio-container {
-    background: var(--background-fill-primary);
-}
-
-.title {
-    color: var(--color-primary) !important;
-    font-size: 3rem !important;
-    font-weight: 700 !important;
-    text-align: center;
-    margin: 1rem 0;
-    text-shadow: 2px 2px 4px rgba(0,0,0,0.05);
-    font-family: 'Playfair Display', serif;
-}
-
-.subtitle {
-    color: #4A5568 !important;
-    font-size: 1.2rem !important;
-    text-align: center;
-    margin-bottom: 1.5rem;
-    font-style: italic;
-}
-
-.collection-link {
-    text-align: center;
-    margin-bottom: 2rem;
-    font-size: 1.1rem;
-}
-
-.collection-link a {
-    color: var(--color-primary);
-    text-decoration: underline;
-    transition: color 0.3s ease;
-}
-
-.collection-link a:hover {
-    color: var(--color-secondary);
-}
-
-.model-description {
-    background-color: rgba(255, 255, 255, 0.8);
-    border-radius: 12px;
-    padding: 24px;
-    margin: 20px 0;
-    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.05);
-    border-left: 5px solid var(--color-primary);
-}
-
-button.primary {
-    background-color: var(--color-primary) !important;
-    transition: all 0.3s ease;
-    color: #fff !important;
-}
-
-button:hover {
-    transform: translateY(-2px);
-    box-shadow: 0 5px 15px rgba(0,0,0,0.1);
-}
-
-.input-container {
-    border-radius: 10px;
-    box-shadow: 0 2px 8px rgba(0,0,0,0.05);
-    background-color: rgba(255, 255, 255, 0.6);
-    padding: 20px;
-    margin-bottom: 1rem;
-}
-
-.advanced-settings {
-    margin-top: 1rem;
-    padding: 1rem;
-    border-radius: 10px;
-    background-color: rgba(255, 255, 255, 0.6);
-}
-
-.example-region {
-    background-color: rgba(255, 255, 255, 0.5);
-    border-radius: 10px;
-    padding: 1rem;
-    margin-top: 1rem;
-}
+footer {visibility: hidden;}
+.gradio-container {background: var(--background-fill-primary);} 
+.title {color: var(--color-primary)!important; font-size:3rem!important; font-weight:700!important; text-align:center; margin:1rem 0; font-family:'Playfair Display',serif;}
+.subtitle {color:#4A5568!important; font-size:1.2rem!important; text-align:center; margin-bottom:1.5rem; font-style:italic;}
+.collection-link {text-align:center; margin-bottom:2rem; font-size:1.1rem;}
+.collection-link a {color:var(--color-primary); text-decoration:underline; transition:color .3s ease;}
+.collection-link a:hover {color:var(--color-secondary);} 
+.model-description{background:rgba(255,255,255,.8); border-radius:12px; padding:24px; margin:20px 0; box-shadow:0 4px 12px rgba(0,0,0,.05); border-left:5px solid var(--color-primary);} 
+button.primary{background:var(--color-primary)!important; color:#fff!important; transition:all .3s ease;} 
+button:hover{transform:translateY(-2px); box-shadow:0 5px 15px rgba(0,0,0,.1);} 
+.input-container{border-radius:10px; box-shadow:0 2px 8px rgba(0,0,0,.05); background:rgba(255,255,255,.6); padding:20px; margin-bottom:1rem;} 
+.advanced-settings{margin-top:1rem; padding:1rem; border-radius:10px; background:rgba(255,255,255,.6);} 
+.example-region{background:rgba(255,255,255,.5); border-radius:10px; padding:1rem; margin-top:1rem;} 
 """
 
+# ===== Gradio UI =====
 with gr.Blocks(css=custom_css, analytics_enabled=False) as demo:
     gr.HTML('<div class="title">Mr. KIM in KOREA</div>')
-    
-    # 컬렉션 링크 또는 안내문을 필요 시 수정/삭제
     gr.HTML('<div class="collection-link"><a href="https://huggingface.co/collections/openfree/painting-art-ai-681453484ec15ef5978bbeb1" target="_blank">Visit the LoRA Model Collection</a></div>')
-    
-    # 모델 설명: 특정 정치인에 대한 LoRA 모델임을 언급
+
     with gr.Group(elem_classes="model-description"):
         gr.HTML("""
         <p>
@@ -199,85 +188,39 @@ with gr.Blocks(css=custom_css, analytics_enabled=False) as demo:
         </p>
         """)
 
-    # 메인 UI
-    with gr.Column(elem_id="col-container"):
+    # ===== 메인 입력 =====
+    with gr.Column():
         with gr.Row(elem_classes="input-container"):
-            prompt = gr.Text(
-                label="Prompt",
-                max_lines=1,
-                placeholder="Enter your prompt (add [trigger] at the end)",
-                value=examples[0]  # 기본 예시
-            )
-            run_button = gr.Button("Generate", variant="primary", scale=0)
+            user_prompt = gr.Text(label="Prompt", max_lines=1, value=examples[0])
+            style_select = gr.Radio(label="Style Preset", choices=list(STYLE_PRESETS.keys()), value="None", interactive=True)
+            run_button = gr.Button("Generate", variant="primary")
 
-        result = gr.Image(label="Generated Image")
-        seed_output = gr.Number(label="Seed", visible=True)
+        result_image = gr.Image(label="Generated Image")
+        seed_output = gr.Number(label="Seed")
 
+        # ===== 고급 설정 =====
         with gr.Accordion("Advanced Settings", open=False, elem_classes="advanced-settings"):
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=42,
-            )
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
             with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
-                )
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=768,
-                )
-
+                width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
+                height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=768)
             with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=3.5,
-                )
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=30,
-                )
-                lora_scale = gr.Slider(
-                    label="LoRA scale",
-                    minimum=0.0,
-                    maximum=1.0,
-                    step=0.1,
-                    value=1.0,
-                )
+                guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=3.5)
+                num_inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=50, step=1, value=30)
+                lora_scale = gr.Slider(label="LoRA scale", minimum=0.0, maximum=1.0, step=0.1, value=1.0)
 
+        # ===== 예시 영역 =====
         with gr.Group(elem_classes="example-region"):
             gr.Markdown("### Examples")
-            gr.Examples(
-                examples=examples,
-                inputs=prompt,
-                outputs=None,  # Don't auto-run examples
-                fn=None,       # No function to run for examples - just fill the prompt
-                cache_examples=False,
-            )
+            gr.Examples(examples=examples, inputs=user_prompt, cache_examples=False)
 
-    # 이벤트 핸들러
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=inference,
+    # ===== 이벤트 =====
+    run_button.click(
+        fn=generate_image,
         inputs=[
-            prompt,
+            user_prompt,
+            style_select,
             seed,
             randomize_seed,
             width,
@@ -286,8 +229,9 @@ with gr.Blocks(css=custom_css, analytics_enabled=False) as demo:
             num_inference_steps,
             lora_scale,
         ],
-        outputs=[result, seed_output],
+        outputs=[result_image, seed_output],
     )
 
+
 demo.queue()
 demo.launch()