Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import gc
+import re
 import gradio as gr
 import numpy as np
 import torch
@@ -8,10 +9,14 @@ import spaces
 import config
 import utils
 import logging
-from PIL import Image, PngImagePlugin
+import time
 from datetime import datetime
+from typing import List, Dict, Tuple, Optional
+from PIL import Image, PngImagePlugin
 from diffusers.models import AutoencoderKL
 from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
+from transformers import pipeline as translation_pipeline
+
 from config import (
     MODEL,
     MIN_IMAGE_SIZE,
@@ -26,8 +31,6 @@ from config import (
     aspect_ratios,
     style_list,
 )
-import time
-from typing import List, Dict, Tuple, Optional
 
 # Enhanced logging configuration
 logging.basicConfig(
@@ -50,10 +53,27 @@ torch.backends.cuda.matmul.allow_tf32 = True
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
+# 번역 파이프라인 초기화 (한글 → 영어)
+translator = translation_pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+
 class GenerationError(Exception):
     """Custom exception for generation errors"""
     pass
 
+def translate_if_korean(prompt: str) -> str:
+    """프롬프트에 한글이 포함되어 있으면 영어로 번역"""
+    if re.search(r'[ㄱ-ㅎㅏ-ㅣ가-힣]', prompt):
+        logger.info("Korean detected in prompt. Translating to English...")
+        try:
+            translation = translator(prompt)[0]['translation_text']
+            logger.info(f"Translation result: {translation}")
+            return translation
+        except Exception as e:
+            logger.error(f"Translation error: {e}")
+            # 번역 실패 시 원본 프롬프트 사용
+            return prompt
+    return prompt
+
 def validate_prompt(prompt: str) -> str:
     """Validate and clean up the input prompt."""
     if not isinstance(prompt, str):
@@ -69,6 +89,9 @@ def validate_prompt(prompt: str) -> str:
     # Only check if the prompt is completely empty or only whitespace
     if not prompt or prompt.isspace():
         raise GenerationError("Prompt cannot be empty")
+    
+    # 번역 적용: 한글이 감지되면 영어로 변환
+    prompt = translate_if_korean(prompt)
     return prompt.strip()
 
 def validate_dimensions(width: int, height: int) -> None:
@@ -124,7 +147,7 @@ def generate(
 
         # Process prompts
         if add_quality_tags:
-            prompt = "masterpiece, high score, great score, absurdres, {prompt}".format(prompt=prompt)
+            prompt = f"masterpiece, high score, great score, absurdres, {prompt}"
 
         prompt, negative_prompt = utils.preprocess_prompt(
             styles, style_selector, prompt, negative_prompt
@@ -257,121 +280,151 @@ else:
 # Process styles
 styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
 
-with gr.Blocks(css="style.css", theme="Nymbo/Nymbo_Theme_5") as demo:
+# 사용자 인터페이스 (UI) 개선: CSS 스타일 및 레이아웃 수정
+custom_css = """
+/* 배경 및 글자 색상 변경 */
+body {
+    background-color: #f7f9fc;
+    color: #333;
+    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+}
+
+/* 헤더 스타일 */
+.header {
+    text-align: center;
+    padding: 20px;
+}
+.header .title {
+    font-size: 3em;
+    font-weight: bold;
+    color: #2c3e50;
+}
+.header .subtitle {
+    font-size: 1.2em;
+    color: #7f8c8d;
+}
+a {
+    text-decoration: none;
+    color: #3498db;
+}
+
+/* Discord 버튼 스타일 */
+.discord-btn {
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    padding: 10px 20px;
+    background: #7289da;
+    color: white;
+    border-radius: 8px;
+    font-weight: bold;
+    margin-top: 20px;
+}
+.discord-btn:hover {
+    background: #5b6eae;
+}
+.discord-icon {
+    width: 24px;
+    height: 24px;
+    margin-right: 8px;
+}
+
+/* Gradio 갤러리 스타일 개선 */
+.gradio-gallery {
+    border: none;
+    box-shadow: none;
+}
+"""
+
+with gr.Blocks(css=custom_css, theme="default") as demo:
+    # 상단 헤더
     gr.HTML(
         """
         <div class="header">
             <div class="title">ANIM4GINE</div>
-            <div class="subtitle">Gradio demo for <a href="https://huggingface.co/CagliostroLab/Animagine-XL-4.0" target="_blank">Animagine XL 4.0</a></div>
+            <div class="subtitle">A sleek and intuitive demo for <a href="https://huggingface.co/CagliostroLab/Animagine-XL-4.0" target="_blank">Animagine XL 4.0</a></div>
         </div>
-        """,
+        """
     )
     
+    # 좌측에 입력 및 설정, 우측에 결과 출력 배치
     with gr.Row():
-        with gr.Column(scale=2):
-            with gr.Group():
-                prompt = gr.Text(
-                    label="Prompt",
-                    max_lines=5,
-                    placeholder="Describe what you want to generate",
-                    info="Enter your image generation prompt here. Be specific and descriptive for better results.",
-                )
-                negative_prompt = gr.Text(
-                    label="Negative Prompt",
-                    max_lines=5,
-                    placeholder="Describe what you want to avoid",
-                    value=DEFAULT_NEGATIVE_PROMPT,
-                    info="Specify elements you don't want in the image.",
-                )
-                add_quality_tags = gr.Checkbox(
-                    label="Quality Tags",
-                    value=True,
-                    info="Add quality-enhancing tags to your prompt automatically.",
-                )
-            with gr.Accordion(label="More Settings", open=False):
-                with gr.Group():
-                    aspect_ratio_selector = gr.Radio(
-                        label="Aspect Ratio",
-                        choices=aspect_ratios,
-                        value=DEFAULT_ASPECT_RATIO,
-                        container=True,
-                        info="Choose the dimensions of your image.",
-                    )
-                with gr.Group(visible=False) as custom_resolution:
-                    with gr.Row():
-                        custom_width = gr.Slider(
-                            label="Width",
-                            minimum=MIN_IMAGE_SIZE,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=8,
-                            value=1024,
-                            info=f"Image width (must be between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE})",
-                        )
-                        custom_height = gr.Slider(
-                            label="Height",
-                            minimum=MIN_IMAGE_SIZE,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=8,
-                            value=1024,
-                            info=f"Image height (must be between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE})",
-                        )
-                with gr.Group():
-                    use_upscaler = gr.Checkbox(
-                        label="Use Upscaler",
-                        value=False,
-                        info="Enable high-resolution upscaling.",
-                    )
-                    with gr.Row() as upscaler_row:
-                        upscaler_strength = gr.Slider(
-                            label="Strength",
-                            minimum=0,
-                            maximum=1,
-                            step=0.05,
-                            value=0.55,
-                            visible=False,
-                            info="Control how much the upscaler affects the final image.",
-                        )
-                        upscale_by = gr.Slider(
-                            label="Upscale by",
-                            minimum=1,
-                            maximum=1.5,
-                            step=0.1,
-                            value=1.5,
-                            visible=False,
-                            info="Multiplier for the final image resolution.",
-                        )
-                with gr.Accordion(label="Advanced Parameters", open=False):
-                    with gr.Group():
-                        style_selector = gr.Dropdown(
-                            label="Style Preset",
-                            interactive=True,
-                            choices=list(styles.keys()),
-                            value="(None)",
-                            info="Apply a predefined style to your generation.",
-                        )
+        with gr.Column(scale=3):
+            with gr.Tabs():
+                with gr.TabItem("Generate"):
                     with gr.Group():
-                        sampler = gr.Dropdown(
-                            label="Sampler",
-                            choices=sampler_list,
-                            interactive=True,
-                            value="Euler a",
-                            info="Different samplers can produce varying results.",
+                        prompt = gr.Textbox(
+                            label="Prompt",
+                            lines=4,
+                            placeholder="Describe what you want to generate...",
+                            info="Enter your image generation prompt here. 한글 입력 시 자동으로 영어로 번역됩니다.",
                         )
-                    with gr.Group():
-                        seed = gr.Slider(
-                            label="Seed",
-                            minimum=0,
-                            maximum=utils.MAX_SEED,
-                            step=1,
-                            value=0,
-                            info="Set a specific seed for reproducible results.",
+                        negative_prompt = gr.Textbox(
+                            label="Negative Prompt",
+                            lines=4,
+                            placeholder="Describe what you want to avoid",
+                            value=DEFAULT_NEGATIVE_PROMPT,
+                            info="Specify elements you don't want in the image.",
                         )
-                        randomize_seed = gr.Checkbox(
-                            label="Randomize seed",
+                        add_quality_tags = gr.Checkbox(
+                            label="Quality Tags",
                             value=True,
-                            info="Generate a new random seed for each image.",
+                            info="Automatically add quality-enhancing tags to your prompt.",
                         )
-                    with gr.Group():
+                    with gr.Accordion(label="More Settings", open=False):
+                        with gr.Column():
+                            aspect_ratio_selector = gr.Radio(
+                                label="Aspect Ratio",
+                                choices=aspect_ratios,
+                                value=DEFAULT_ASPECT_RATIO,
+                                container=True,
+                                info="Choose the dimensions of your image.",
+                            )
+                            with gr.Row(visible=False) as custom_resolution:
+                                custom_width = gr.Slider(
+                                    label="Width",
+                                    minimum=MIN_IMAGE_SIZE,
+                                    maximum=MAX_IMAGE_SIZE,
+                                    step=8,
+                                    value=1024,
+                                    info=f"Image width (between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE})",
+                                )
+                                custom_height = gr.Slider(
+                                    label="Height",
+                                    minimum=MIN_IMAGE_SIZE,
+                                    maximum=MAX_IMAGE_SIZE,
+                                    step=8,
+                                    value=1024,
+                                    info=f"Image height (between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE})",
+                                )
+                    with gr.Accordion(label="Advanced Parameters", open=False):
+                        with gr.Row():
+                            style_selector = gr.Dropdown(
+                                label="Style Preset",
+                                choices=list(styles.keys()),
+                                value="(None)",
+                                info="Apply a predefined style to your generation.",
+                            )
+                            sampler = gr.Dropdown(
+                                label="Sampler",
+                                choices=sampler_list,
+                                value="Euler a",
+                                info="Different samplers can produce varying results.",
+                            )
+                        with gr.Row():
+                            seed = gr.Slider(
+                                label="Seed",
+                                minimum=0,
+                                maximum=utils.MAX_SEED,
+                                step=1,
+                                value=0,
+                                info="Set a specific seed for reproducible results.",
+                            )
+                            randomize_seed = gr.Checkbox(
+                                label="Randomize seed",
+                                value=True,
+                                info="Generate a new random seed for each image.",
+                            )
                         with gr.Row():
                             guidance_scale = gr.Slider(
                                 label="Guidance scale",
@@ -382,48 +435,84 @@ with gr.Blocks(css="style.css", theme="Nymbo/Nymbo_Theme_5") as demo:
                                 info="Higher values make the image more closely match your prompt.",
                             )
                             num_inference_steps = gr.Slider(
-                                label="Number of inference steps",
+                                label="Inference steps",
                                 minimum=1,
                                 maximum=50,
                                 step=1,
                                 value=25,
-                                info="More steps generally mean higher quality but slower generation.",
+                                info="More steps generally yield higher quality images.",
                             )
+                        with gr.Row():
+                            use_upscaler = gr.Checkbox(
+                                label="Use Upscaler",
+                                value=False,
+                                info="Enable high-resolution upscaling.",
+                            )
+                            upscaler_strength = gr.Slider(
+                                label="Upscaler Strength",
+                                minimum=0,
+                                maximum=1,
+                                step=0.05,
+                                value=0.55,
+                                visible=False,
+                                info="Control how strongly the upscaler affects the image.",
+                            )
+                            upscale_by = gr.Slider(
+                                label="Upscale By",
+                                minimum=1,
+                                maximum=1.5,
+                                step=0.1,
+                                value=1.5,
+                                visible=False,
+                                info="Multiplier for the final image resolution.",
+                            )
+                with gr.TabItem("Examples"):
+                    gr.Markdown(
+                        """
+                        ### Example Prompts
+                        - **Scenic Landscape:** A breathtaking view of a mountain landscape during sunrise.
+                        - **Cyberpunk City:** A futuristic cyberpunk city with neon lights and towering skyscrapers.
+                        - **Fantasy Character:** A majestic wizard with a long beard and glowing magical staff.
+                        """
+                    )
+                    gr.Examples(
+                        examples=examples,
+                        inputs=prompt,
+                        outputs=[],
+                        cache_examples=CACHE_EXAMPLES,
+                    )
+            run_button = gr.Button("Generate", variant="primary", elem_id="generate-button")
         
-        with gr.Column(scale=3):
-            with gr.Blocks():
-                run_button = gr.Button("Generate", variant="primary", elem_id="generate-button")
+        with gr.Column(scale=4):
             result = gr.Gallery(
                 label="Generated Images",
-                columns=1,
-                height='768px',
-                preview=True,
+                columns=2,
+                height="600px",
                 show_label=True,
+                elem_classes="gradio-gallery",
             )
             with gr.Accordion(label="Generation Parameters", open=False):
                 gr_metadata = gr.JSON(
                     label="Image Metadata",
                     show_label=True,
                 )
-            gr.Examples(
-                examples=examples,
-                inputs=prompt,
-                outputs=[result, gr_metadata],
-                fn=lambda *args, **kwargs: generate(*args, use_upscaler=True, **kwargs),
-                cache_examples=CACHE_EXAMPLES,
-            )
     
-    # Discord button in a new full row
+    # Discord 버튼을 하단 중앙에 배치
     with gr.Row():
         gr.HTML(
             """
-            <a href="https://discord.com/invite/cqh9tZgbGc" target="_blank" class="discord-btn">
-                <svg class="discord-icon" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 127.14 96.36"><path fill="currentColor" d="M107.7,8.07A105.15,105.15,0,0,0,81.47,0a72.06,72.06,0,0,0-3.36,6.83A97.68,97.68,0,0,0,49,6.83,72.37,72.37,0,0,0,45.64,0,105.89,105.89,0,0,0,19.39,8.09C2.79,32.65-1.71,56.6.54,80.21h0A105.73,105.73,0,0,0,32.71,96.36,77.7,77.7,0,0,0,39.6,85.25a68.42,68.42,0,0,1-10.85-5.18c.91-.66,1.8-1.34,2.66-2a75.57,75.57,0,0,0,64.32,0c.87.71,1.76,1.39,2.66,2a68.68,68.68,0,0,1-10.87,5.19,77,77,0,0,0,6.89,11.1A105.25,105.25,0,0,0,126.6,80.22h0C129.24,52.84,122.09,29.11,107.7,8.07ZM42.45,65.69C36.18,65.69,31,60,31,53s5-12.74,11.43-12.74S54,46,53.89,53,48.84,65.69,42.45,65.69Zm42.24,0C78.41,65.69,73.25,60,73.25,53s5-12.74,11.44-12.74S96.23,46,96.12,53,91.08,65.69,84.69,65.69Z"/></svg>
-                <span class="discord-text">Join our Discord Server</span>
-            </a>
+            <div style="width:100%; display:flex; justify-content:center;">
+                <a href="https://discord.com/invite/cqh9tZgbGc" target="_blank" class="discord-btn">
+                    <svg class="discord-icon" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 127.14 96.36">
+                        <path fill="currentColor" d="M107.7,8.07A105.15,105.15,0,0,0,81.47,0a72.06,72.06,0,0,0-3.36,6.83A97.68,97.68,0,0,0,49,6.83,72.37,72.37,0,0,0,45.64,0,105.89,105.89,0,0,0,19.39,8.09C2.79,32.65-1.71,56.6.54,80.21h0A105.73,105.73,0,0,0,32.71,96.36,77.7,77.7,0,0,0,39.6,85.25a68.42,68.42,0,0,1-10.85-5.18c.91-.66,1.8-1.34,2.66-2a75.57,75.57,0,0,0,64.32,0c.87.71,1.76,1.39,2.66,2a68.68,68.68,0,0,1-10.87,5.19,77,77,0,0,0,6.89,11.1A105.25,105.25,0,0,0,126.6,80.22h0C129.24,52.84,122.09,29.11,107.7,8.07ZM42.45,65.69C36.18,65.69,31,60,31,53s5-12.74,11.43-12.74S54,46,53.89,53,48.84,65.69,42.45,65.69Zm42.24,0C78.41,65.69,73.25,60,73.25,53s5-12.74,11.44-12.74S96.23,46,96.12,53,91.08,65.69,84.69,65.69Z"/>
+                    </svg>
+                    <span class="discord-text">Join our Discord Server</span>
+                </a>
+            </div>
             """
         )
-
+    
+    # UI 동적 업데이트: 업스케일러 옵션 및 커스텀 해상도 슬라이더 표시 제어
     use_upscaler.change(
         fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
         inputs=use_upscaler,
@@ -439,7 +528,7 @@ with gr.Blocks(css="style.css", theme="Nymbo/Nymbo_Theme_5") as demo:
         api_name=False,
     )
 
-    # Combine all triggers including keyboard shortcuts
+    # 트리거 설정: 프롬프트 입력, 버튼 클릭 시 시드 랜덤화 후 이미지 생성 실행
     gr.on(
         triggers=[
             prompt.submit,