app.py

# ===== CRITICAL: Import spaces FIRST before any CUDA operations =====
try:
    import spaces
    HF_SPACES = True
except ImportError:
    # If running locally, create a dummy decorator
    def spaces_gpu_decorator(duration=60):
        def decorator(func):
            return func
        return decorator
    spaces = type('spaces', (), {'GPU': spaces_gpu_decorator})()
    HF_SPACES = False
    print("Warning: Running without Hugging Face Spaces GPU allocation")

# ===== Now import other libraries =====
import random
import os
import uuid
import re
import time
from datetime import datetime

import gradio as gr
import numpy as np
import requests
import torch
from diffusers import DiffusionPipeline
from PIL import Image

# ===== OpenAI 설정 =====
from openai import OpenAI

# Add error handling for API key
try:
    client = OpenAI(api_key=os.getenv("LLM_API"))
except Exception as e:
    print(f"Warning: OpenAI client initialization failed: {e}")
    client = None

# ===== 프롬프트 증강용 스타일 프리셋 =====
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"
if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR, exist_ok=True)

# ===== 디바이스 & 모델 로드 =====
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

repo_id = "black-forest-labs/FLUX.1-dev"
adapter_id = "seawolf2357/kim-korea"

# Add error handling for model loading
try:
    pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
    pipeline.load_lora_weights(adapter_id)
    pipeline = pipeline.to(device)
    print("Model loaded successfully")
except Exception as e:
    print(f"Error loading model: {e}")
    pipeline = None

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

# ===== 한글 여부 판별 =====
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-4o-mini 사용). 영어 입력이면 그대로 반환."""
    if not is_korean(text):
        return text
    
    if client is None:
        print("Warning: OpenAI client not available, returning original text")
        return text

    for attempt in range(retries):
        try:
            res = client.chat.completions.create(
                model="gpt-4o-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 Exception as e:
            print(f"[translate] attempt {attempt + 1} failed: {e}")
            time.sleep(2)
    return text  # 번역 실패 시 원문 그대로

def enhance_prompt(text: str, retries: int = 3) -> str:
    """OpenAI를 통해 프롬프트를 증강하여 고품질 이미지 생성을 위한 상세한 설명으로 변환."""
    if client is None:
        print("Warning: OpenAI client not available, returning original text")
        return text

    for attempt in range(retries):
        try:
            res = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {
                        "role": "system",
                        "content": """You are an expert prompt engineer for image generation models. Enhance the given prompt to create high-quality, detailed images.

Guidelines:
- Add specific visual details (lighting, composition, colors, textures)
- Include technical photography terms (depth of field, focal length, etc.)
- Add atmosphere and mood descriptors
- Specify image quality terms (4K, ultra-detailed, professional, etc.)
- Keep the core subject and meaning intact
- Make it comprehensive but not overly long
- Focus on visual elements that will improve image generation quality

Example:
Input: "A man giving a speech"
Output: "A professional man giving an inspiring speech at a podium, dramatic lighting with warm spotlights, confident posture and gestures, high-resolution 4K photography, sharp focus, cinematic composition, bokeh background with audience silhouettes, professional event setting, detailed facial expressions, realistic skin texture"
"""
                    },
                    {"role": "user", "content": f"Enhance this prompt for high-quality image generation: {text}"}
                ],
                temperature=0.7,
                max_tokens=512,
            )
            return res.choices[0].message.content.strip()
        except Exception as e:
            print(f"[enhance] attempt {attempt + 1} failed: {e}")
            time.sleep(2)
    return text  # 증강 실패 시 원문 그대로

def prepare_prompt(user_prompt: str, style_key: str, enhance_prompt_enabled: bool = False) -> str:
    """한글이면 번역하고, 프롬프트 증강 옵션이 활성화되면 증강하고, 선택한 스타일 프리셋을 붙여서 최종 프롬프트를 만든다."""
    # 1. 번역 (한글인 경우)
    prompt_en = openai_translate(user_prompt)
    
    # 2. 프롬프트 증강 (활성화된 경우)
    if enhance_prompt_enabled:
        prompt_en = enhance_prompt(prompt_en)
        print(f"Enhanced prompt: {prompt_en}")
    
    # 3. 스타일 프리셋 적용
    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)
    image.save(filepath)

    # 메타데이터 저장
    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

# ===== Diffusion 호출 =====

def run_pipeline(prompt: str, seed: int, width: int, height: int, guidance_scale: float, num_steps: int, lora_scale: float):
    if pipeline is None:
        raise ValueError("Model pipeline not loaded")
    
    generator = torch.Generator(device=device).manual_seed(int(seed))
    result = pipeline(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_steps,
        width=width,
        height=height,
        generator=generator,
        joint_attention_kwargs={"scale": lora_scale},
    ).images[0]
    return result

# ===== Gradio inference 래퍼 =====

@spaces.GPU(duration=60)
def generate_image(
    user_prompt: str,
    style_key: str,
    enhance_prompt_enabled: bool = False,
    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,
):
    try:
        if randomize_seed:
            seed = random.randint(0, MAX_SEED)

        # 1) 번역 + 증강
        final_prompt = prepare_prompt(user_prompt, style_key, enhance_prompt_enabled)
        print(f"Final prompt: {final_prompt}")

        # 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
    
    except Exception as e:
        print(f"Error generating image: {e}")
        # Return a placeholder or error message
        error_image = Image.new('RGB', (width, height), color='red')
        return error_image, seed

# ===== 예시 프롬프트 (한국어/영어 혼용 허용) =====

examples = [
    "Mr. KIM이 두 손으로 'Fighting!' 현수막을 들고 있는 모습, 애국심과 국가 발전에 대한 의지를 보여주고 있다.",   
    "Mr. KIM이 양팔을 들어 올리며 승리의 표정으로 환호하는 모습, 승리와 미래에 대한 희망을 보여주고 있다.",
    "Mr. KIM이 운동복을 입고 공원에서 조깅하는 모습, 건강한 생활습관과 활기찬 리더십을 보여주고 있다.",  
    "Mr. KIM이 붐비는 거리에서 여성 시민들과 따뜻하게 악수하는 모습, 여성 유권자들에 대한 진정한 관심과 소통을 보여주고 있다.",
    "Mr. KIM이 선거 유세장에서 지평선을 향해 손가락으로 가리키며 영감을 주는 제스처를 취하고 있고, 여성들과 아이들이 박수를 치고 있다.",
    "Mr. KIM이 지역 행사에 참여하여 열정적으로 응원하는 여성 지지자들에게 둘러싸여 있는 모습.",
    "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.",
]

# ===== 커스텀 CSS (진한 붉은색 고급 디자인) =====
custom_css = """
:root {
    --color-primary: #E91E63;
    --color-secondary: #FCE4EC;
    --color-accent: #F8BBD9;
    --color-rose: #F06292;
    --color-gold: #FFB74D;
    --color-warm-gray: #F5F5F5;
    --color-dark-gray: #424242;
    --background-primary: linear-gradient(135deg, #FAFAFA 0%, #F5F5F5 50%, #EEEEEE 100%);
    --background-accent: linear-gradient(135deg, #FCE4EC 0%, #F8BBD9 100%);
    --text-primary: #212121;
    --text-secondary: #757575;
    --shadow-soft: 0 4px 20px rgba(0, 0, 0, 0.08);
    --shadow-medium: 0 8px 30px rgba(0, 0, 0, 0.12);
    --border-radius: 16px;
}

/* 전체 배경 */
footer {visibility: hidden;}
.gradio-container {
    background: var(--background-primary) !important;
    min-height: 100vh;
    font-family: 'Inter', 'Noto Sans KR', sans-serif;
}

/* 타이틀 스타일 */
.title {
    color: var(--text-primary) !important;
    font-size: 3rem !important;
    font-weight: 700 !important;
    text-align: center;
    margin: 2rem 0;
    background: linear-gradient(135deg, var(--color-primary) 0%, var(--color-rose) 50%, var(--color-gold) 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    letter-spacing: -0.02em;
}

.subtitle {
    color: var(--text-secondary) !important;
    font-size: 1.2rem !important;
    text-align: center;
    margin-bottom: 2rem;
    font-weight: 400;
}

.collection-link {
    text-align: center;
    margin-bottom: 2rem;
    font-size: 1rem;
}

.collection-link a {
    color: var(--color-primary);
    text-decoration: none;
    transition: all 0.3s ease;
    font-weight: 500;
    border-bottom: 1px solid transparent;
}

.collection-link a:hover {
    color: var(--color-rose);
    border-bottom-color: var(--color-rose);
}

/* 심플한 카드 스타일 */
.model-description {
    background: rgba(255, 255, 255, 0.9);
    border: 1px solid rgba(233, 30, 99, 0.1);
    border-radius: var(--border-radius);
    padding: 2rem;
    margin: 1.5rem 0;
    box-shadow: var(--shadow-soft);
    backdrop-filter: blur(10px);
    -webkit-backdrop-filter: blur(10px);
}

.model-description p {
    color: var(--text-primary) !important;
    font-size: 1rem;
    line-height: 1.6;
    margin: 0;
}

/* 버튼 스타일 */
button.primary {
    background: var(--background-accent) !important;
    color: var(--color-primary) !important;
    border: 1px solid var(--color-accent) !important;
    border-radius: 12px !important;
    box-shadow: var(--shadow-soft) !important;
    transition: all 0.2s ease !important;
    font-weight: 600 !important;
    font-size: 0.95rem !important;
}

button.primary:hover {
    background: linear-gradient(135deg, var(--color-accent) 0%, var(--color-secondary) 100%) !important;
    transform: translateY(-1px) !important;
    box-shadow: var(--shadow-medium) !important;
}

/* 입력 컨테이너 */
.input-container {
    background: rgba(255, 255, 255, 0.8);
    border: 1px solid rgba(233, 30, 99, 0.15);
    border-radius: var(--border-radius);
    padding: 1.5rem;
    margin-bottom: 1.5rem;
    box-shadow: var(--shadow-soft);
    backdrop-filter: blur(10px);
    -webkit-backdrop-filter: blur(10px);
}

/* 고급 설정 */
.advanced-settings {
    background: rgba(255, 255, 255, 0.6);
    border: 1px solid rgba(233, 30, 99, 0.1);
    border-radius: var(--border-radius);
    padding: 1.5rem;
    margin-top: 1rem;
    box-shadow: var(--shadow-soft);
    backdrop-filter: blur(8px);
    -webkit-backdrop-filter: blur(8px);
}

/* 예시 영역 */
.example-region {
    background: rgba(252, 228, 236, 0.3);
    border: 1px solid rgba(233, 30, 99, 0.15);
    border-radius: var(--border-radius);
    padding: 1.5rem;
    margin-top: 1rem;
    box-shadow: var(--shadow-soft);
}

/* 프롬프트 입력칸 스타일 */
.large-prompt textarea {
    min-height: 120px !important;
    font-size: 15px !important;
    line-height: 1.5 !important;
    background: rgba(255, 255, 255, 0.9) !important;
    border: 2px solid rgba(233, 30, 99, 0.2) !important;
    border-radius: 12px !important;
    color: var(--text-primary) !important;
    transition: all 0.3s ease !important;
    padding: 1rem !important;
}

.large-prompt textarea:focus {
    border-color: var(--color-primary) !important;
    box-shadow: 0 0 0 3px rgba(233, 30, 99, 0.1) !important;
    outline: none !important;
}

.large-prompt textarea::placeholder {
    color: var(--text-secondary) !important;
    font-style: italic;
}

/* 생성 버튼 */
.small-generate-btn {
    max-width: 140px !important;
    height: 48px !important;
    font-size: 15px !important;
    padding: 12px 24px !important;
    border-radius: 12px !important;
    font-weight: 600 !important;
}

/* 프롬프트 증강 섹션 */
.prompt-enhance-section {
    background: linear-gradient(135deg, rgba(255, 183, 77, 0.1) 0%, rgba(252, 228, 236, 0.2) 100%);
    border: 1px solid rgba(255, 183, 77, 0.3);
    border-radius: var(--border-radius);
    padding: 1.2rem;
    margin-top: 1rem;
    box-shadow: var(--shadow-soft);
}

/* 스타일 프리셋 섹션 */
.style-preset-section {
    background: linear-gradient(135deg, rgba(248, 187, 217, 0.15) 0%, rgba(252, 228, 236, 0.2) 100%);
    border: 1px solid rgba(233, 30, 99, 0.2);
    border-radius: var(--border-radius);
    padding: 1.2rem;
    margin-top: 1rem;
    box-shadow: var(--shadow-soft);
}

/* 라벨 텍스트 */
label {
    color: var(--text-primary) !important;
    font-weight: 600 !important;
    font-size: 0.95rem !important;
}

/* 정보 텍스트 */
.gr-info, .gr-textbox-info {
    color: var(--text-secondary) !important;
    font-size: 0.85rem !important;
    line-height: 1.4 !important;
}

/* 예시 마크다운 */
.example-region h3 {
    color: var(--text-primary) !important;
    font-weight: 600 !important;
    margin-bottom: 1rem !important;
}

/* 폼 요소들 */
input[type="radio"], input[type="checkbox"] {
    accent-color: var(--color-primary) !important;
}

input[type="range"] {
    accent-color: var(--color-primary) !important;
}

/* 결과 이미지 컨테이너 */
.image-container {
    border-radius: var(--border-radius) !important;
    overflow: hidden !important;
    box-shadow: var(--shadow-medium) !important;
    background: rgba(255, 255, 255, 0.9) !important;
    border: 1px solid rgba(233, 30, 99, 0.1) !important;
}

/* 슬라이더 컨테이너 스타일링 */
.gr-slider {
    margin: 0.5rem 0 !important;
}

/* 아코디언 스타일 */
.gr-accordion {
    border: 1px solid rgba(233, 30, 99, 0.15) !important;
    border-radius: var(--border-radius) !important;
    background: rgba(255, 255, 255, 0.7) !important;
}

.gr-accordion-header {
    background: var(--background-accent) !important;
    color: var(--color-primary) !important;
    font-weight: 600 !important;
    border-radius: var(--border-radius) var(--border-radius) 0 0 !important;
}

/* 부드러운 애니메이션 */
.model-description, .input-container, .prompt-enhance-section, .style-preset-section {
    animation: fadeInUp 0.4s ease-out;
}

@keyframes fadeInUp {
    from {
        opacity: 0;
        transform: translateY(20px);
    }
    to {
        opacity: 1;
        transform: translateY(0);
    }
}

/* 전체적인 텍스트 가독성 향상 */
* {
    -webkit-font-smoothing: antialiased;
    -moz-osx-font-smoothing: grayscale;
}

/* 드롭다운 및 셀렉트 스타일 */
select, .gr-dropdown {
    background: rgba(255, 255, 255, 0.9) !important;
    border: 1px solid rgba(233, 30, 99, 0.2) !important;
    border-radius: 8px !important;
    color: var(--text-primary) !important;
}

/* 체크박스와 라디오 버튼 개선 */
.gr-checkbox, .gr-radio {
    background: transparent !important;
}

/* 전체 컨테이너 여백 조정 */
.gr-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
    padding: 2rem 1rem !important;
}

/* 모바일 반응형 */
@media (max-width: 768px) {
    .title {
        font-size: 2.2rem !important;
    }
    
    .model-description, .input-container, .advanced-settings, .example-region {
        padding: 1rem !important;
        margin: 1rem 0 !important;
    }
    
    .large-prompt textarea {
        min-height: 100px !important;
        font-size: 14px !important;
    }
}
"""

# ===== Gradio UI =====
def create_interface():
    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 (LoRA 모델 컬렉션 방문)</a></div>')

        with gr.Group(elem_classes="model-description"):
            gr.HTML("""
            <p>
            <strong>Research Purpose Model (연구 목적 모델)</strong><br>
            This model is trained on specific person's facial features and appearance using LoRA technology for research purposes.<br>
            <small style="opacity: 0.8;">본 모델은 연구 목적으로 특정인의 얼굴과 외모를 LoRA 기술로 학습한 모델입니다.</small><br><br>
            
            <strong>Usage Guidelines (사용 가이드라인)</strong><br>
            Please ensure appropriate use and avoid unauthorized usage beyond research purposes.<br>
            <small style="opacity: 0.8;">목적 외의 용도로 무단 사용하지 않도록 유의해 주세요.</small><br><br>
            
            <strong>Optimal Results (최적 결과를 위한 팁)</strong><br>
            For best results, make sure to include 'kim' in your example prompts.<br>
            <small style="opacity: 0.8;">최적의 결과를 얻기 위해 예시 프롬프트 사용 시 반드시 'kim'을 포함하여 주세요.</small>
            </p>
            """)

        # ===== 메인 입력 =====
        with gr.Column():
            with gr.Row(elem_classes="input-container"):
                with gr.Column(scale=4):
                    user_prompt = gr.Text(
                        label="Prompt (프롬프트)", 
                        max_lines=5, 
                        value=examples[0],
                        elem_classes="large-prompt",
                        placeholder="Enter your image description here... (이미지 설명을 입력하세요...)"
                    )
                with gr.Column(scale=1):
                    run_button = gr.Button(
                        "Generate (생성)", 
                        variant="primary",
                        elem_classes="small-generate-btn"
                    )
            
            # 프롬프트 증강 옵션 (생성 버튼 아래)
            with gr.Group(elem_classes="prompt-enhance-section"):
                enhance_prompt_checkbox = gr.Checkbox(
                    label="🚀 Prompt Enhancement (프롬프트 증강)", 
                    value=False,
                    info="Automatically improve your prompt using OpenAI API for high-quality image generation (OpenAI API를 사용하여 고품질 이미지 생성을 위해 프롬프트를 자동으로 개선합니다)"
                )
            
            # 스타일 프리셋 섹션
            with gr.Group(elem_classes="style-preset-section"):
                style_select = gr.Radio(
                    label="🎨 Style Preset (스타일 프리셋)", 
                    choices=list(STYLE_PRESETS.keys()), 
                    value="None", 
                    interactive=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)
                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)
                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="Inference steps (추론 단계)", minimum=1, maximum=50, step=1, value=30)
                    lora_scale = gr.Slider(label="LoRA scale (LoRA 스케일)", 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=user_prompt, cache_examples=False)

        # ===== 이벤트 =====
        run_button.click(
            fn=generate_image,
            inputs=[
                user_prompt,
                style_select,
                enhance_prompt_checkbox,
                seed,
                randomize_seed,
                width,
                height,
                guidance_scale,
                num_inference_steps,
                lora_scale,
            ],
            outputs=[result_image, seed_output],
        )
    
    return demo

# ===== 애플리케이션 실행 =====
if __name__ == "__main__":
    demo = create_interface()
    demo.queue()
    demo.launch()