app.py

import gradio as gr
import google.generativeai as genai
from PIL import Image
import os
import json

# Gemini API 키 설정 (환경 변수에서 가져오거나 직접 입력)
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")

# Gemini API 초기화
genai.configure(api_key=GEMINI_API_KEY)

# 배경 JSON 파일 경로 설정
BACKGROUNDS_DIR = os.path.join(os.path.dirname(__file__), "background")

# JSON 파일이 존재하지 않을 경우 디렉토리 생성
if not os.path.exists(BACKGROUNDS_DIR):
    os.makedirs(BACKGROUNDS_DIR)

# JSON 파일 로드 함수
def load_background_json(filename):
    file_path = os.path.join(BACKGROUNDS_DIR, filename)
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        print(f"경고: {filename} 파일을 찾을 수 없습니다. 기본값을 사용합니다.")
        return {}

# 배경 데이터 로드
SIMPLE_BACKGROUNDS = load_background_json("simple_backgrounds.json")
STUDIO_BACKGROUNDS = load_background_json("studio_backgrounds.json")
NATURE_BACKGROUNDS = load_background_json("nature_backgrounds.json")
INDOOR_BACKGROUNDS = load_background_json("indoor_backgrounds.json")
ABSTRACT_BACKGROUNDS = load_background_json("abstract_backgrounds.json")

def generate_system_instruction():
    return """당신은 상품 이미지의 배경을 변경하기 위한 고품질 프롬프트를 생성하는 전문가입니다.
사용자가 제공하는 상품명, 배경 유형, 추가 요청사항을 바탕으로 미드저니(Midjourney)에 사용할 수 있는 
상세하고 전문적인 프롬프트를 영어로 생성해주세요.
다음 가이드라인을 반드시 따라야 합니다:

1. 상품을 "#1"로 지정하여 참조합니다. (예: "skincare tube (#1)")

2. *** 매우 중요: 상품의 원래 특성(디자인, 색상, 형태, 로고, 패키지 등)은 어떤 상황에서도 절대 변경하지 않습니다. ***

3. *** 상품의 본질적 특성은 유지하되, 자연스러운 환경 통합을 위한 조명과 그림자는 허용합니다: ***
   - 상품 자체의 색상, 디자인, 형태, 텍스처는 절대 수정하지 않습니다.
   - 환경과 자연스럽게 어울리는 그림자, 주변 조명 효과는 허용됩니다.
   - 상품에 물방울, 응축, 금, 은과 같은 추가 요소나 물리적 효과는 적용하지 않습니다.
   - 환경에 어울리는 자연스러운 빛 반사, 주변 조명, 그림자는 사실적 통합감을 위해 적용할 수 있습니다.

4. 이미지 비율은 정확히 1:1(정사각형) 형식으로 지정합니다. 프롬프트에 "square format", "1:1 ratio" 또는 "aspect ratio 1:1"을 명시적으로 포함합니다.

5. 상품은 반드시 정사각형 구도의 정중앙에 배치되어야 합니다.

6. 상품을 이미지의 주요 초점으로 부각시키고, 상품의 비율이 전체 이미지에서 크게 차지하도록 합니다.

7. 상품 이미지 컷아웃(#1)의 기본 형태와 색상은 유지하면서, 선택한 환경에 자연스럽게 통합되도록 합니다.

8. 고급스러운 상업적 이미지를 위한 다음 환경 요소들을 포함하세요:
   - 상품과 어울리는 주변 환경/배경 요소를 추가합니다. 예를 들어, 화장품 주변에 꽃이나 허브, 음료 제품 옆에 과일, 전자제품 근처에 현대적 소품 등.
   - 환경의 조명 효과(림 라이트, 백라이트, 소프트박스 등)를 설명합니다.
   - 상품이 환경에 자연스럽게 존재하는 것처럼 보이도록 적절한 그림자와 빛 표현을 포함합니다.
   - 상품의 용도나 장점을 간접적으로 암시하는 배경 요소를 포함합니다.
   - 프로페셔널한 상업 사진 효과(선택적 피사계 심도, 소프트 포커스, 스튜디오 조명 등)를 명시합니다.

9. 프롬프트에 다음 요소들을 명시적으로 포함하세요:
   - "highly detailed commercial photography"
   - "award-winning product photography"
   - "professional advertising imagery"
   - "studio quality"
   - "magazine advertisement quality"

10. 배경 환경 요소를 상품 카테고리에 맞게 선택합니다:
    - 스킨케어 제품: 깨끗한 욕실 선반, 우아한 화장대, 스파 같은 환경 등
    - 음료 제품: 세련된 테이블, 파티 환경, 야외 피크닉 장면 등
    - 전자 제품: 세련된 작업 공간, 현대적인 거실, 미니멀한 책상 등
    - 패션/의류: 세련된 쇼룸, 도시 거리, 엘레강스한 라이프스타일 환경 등
    - 식품 제품: 깔끔한 주방, 식탁, 요리 환경 등

11. 사용자가 제공한 구체적인 배경과 추가 요청사항을 정확히 반영합니다.

12. 프롬프트는 미드저니 AI에 최적화되어야 합니다.

13. 프롬프트 끝에 "--ar 1:1 --s 750 --q 2" 파라미터를 추가하여 미드저니에서 고품질 정사각형 비율을 강제합니다.

출력 형식은 영어로 된 단일 단락의 상세한 프롬프트여야 하며, 끝에 미드저니 파라미터가 포함되어야 합니다.
"""

def generate_prompt_with_gemini(product_name, background_info, additional_info=""):
    if not GEMINI_API_KEY:
        return "Gemini API 키가 설정되지 않았습니다. 환경 변수 GEMINI_API_KEY를 설정하거나 코드에 직접 입력하세요."
    try:
        prompt_request = f"""
상품명: {product_name}
배경 유형: {background_info.get('english', 'studio')}
배경 카테고리: {background_info.get('category', '')}
배경 이름: {background_info.get('name', '')}
추가 요청사항: {additional_info}

중요 요구사항:
1. 상품이 크게 부각되고 이미지에서 중심적인 위치를 차지하도록 프롬프트를 생성해주세요.
2. 이미지는 정확히 1:1 비율(정사각형)이어야 합니다.
3. 상품은 정사각형 프레임의 정중앙에 위치해야 합니다.
4. 상품의 디자인, 색상, 형태, 로고 등 본질적 특성은 절대 수정하지 마세요.
5. 환경과의 자연스러운 통합을 위한 조명 효과와 그림자는 포함해주세요.
6. 상품을 더 돋보이게 하는 배경 환경을 설명해주세요.
7. 고급스러운 상업 광고 품질의 이미지가 되도록 환경 설명을 해주세요.
8. 프롬프트 끝에 미드저니 파라미터 "--ar 1:1 --s 750 --q 2"를 추가해주세요.

한국어 입력 내용을 영어로 적절히 번역하여 반영해주세요.
"""
        model = genai.GenerativeModel(
            'gemini-2.0-flash',
            system_instruction=generate_system_instruction()
        )
        response = model.generate_content(
            prompt_request,
            generation_config=genai.types.GenerationConfig(
                temperature=0.7,
                top_p=0.95,
                top_k=64,
                max_output_tokens=1024,
            )
        )
        response_text = response.text.strip()
        if "--ar 1:1" not in response_text:
            response_text = response_text.rstrip(".") + ". --ar 1:1 --s 750 --q 2"
        return response_text
    except Exception as e:
        return f"프롬프트 생성 중 오류가 발생했습니다: {str(e)}"

def create_app():
    with gr.Blocks(title="고급 상품 이미지 배경 프롬프트 생성기") as demo:
        gr.Markdown("# 고급 상품 이미지 배경 프롬프트 생성기")
        gr.Markdown("상품 이미지를 업로드하고 옵션을 선택하면 고품질 상업용 프롬프트가 생성됩니다.")
        
        # 상품명 섹션
        with gr.Row():
            product_name = gr.Textbox(label="상품명 (한국어 입력)", placeholder="예: 스킨케어 튜브, 텀블러 등", interactive=True)
        
        with gr.Row():
            # 좌측 컬럼 - 입력 컨트롤
            with gr.Column(scale=1):
                image_input = gr.Image(label="상품 이미지 업로드", type="pil")
                
                # 배경 타입 탭 (심플, 스튜디오, 자연, 실내, 추상)
                gr.Markdown("## 배경 선택 (한 가지만 선택)")
                
                with gr.Tabs() as background_tabs:
                    with gr.TabItem("심플 배경"):
                        simple_dropdown = gr.Dropdown(
                            choices=list(SIMPLE_BACKGROUNDS.keys()),
                            value=list(SIMPLE_BACKGROUNDS.keys())[0] if SIMPLE_BACKGROUNDS else None,
                            label="심플 배경 선택",
                            interactive=True
                        )
                    
                    with gr.TabItem("스튜디오 배경"):
                        studio_dropdown = gr.Dropdown(
                            choices=list(STUDIO_BACKGROUNDS.keys()),
                            value=list(STUDIO_BACKGROUNDS.keys())[0] if STUDIO_BACKGROUNDS else None,
                            label="스튜디오 배경 선택",
                            interactive=True
                        )
                    
                    with gr.TabItem("자연 환경"):
                        nature_dropdown = gr.Dropdown(
                            choices=list(NATURE_BACKGROUNDS.keys()),
                            value=list(NATURE_BACKGROUNDS.keys())[0] if NATURE_BACKGROUNDS else None,
                            label="자연 환경 선택",
                            interactive=True
                        )
                    
                    with gr.TabItem("실내 환경"):
                        indoor_dropdown = gr.Dropdown(
                            choices=list(INDOOR_BACKGROUNDS.keys()),
                            value=list(INDOOR_BACKGROUNDS.keys())[0] if INDOOR_BACKGROUNDS else None,
                            label="실내 환경 선택",
                            interactive=True
                        )
                    
                    with gr.TabItem("추상/특수 배경"):
                        abstract_dropdown = gr.Dropdown(
                            choices=list(ABSTRACT_BACKGROUNDS.keys()),
                            value=list(ABSTRACT_BACKGROUNDS.keys())[0] if ABSTRACT_BACKGROUNDS else None,
                            label="추상/특수 배경 선택",
                            interactive=True
                        )
                
                # 추가 요청사항
                additional_info = gr.Textbox(
                    label="추가 요청사항 (선택사항)",
                    placeholder="예: 고급스러운 느낌, 밝은 조명 등",
                    lines=3,
                    interactive=True
                )
                
                submit_btn = gr.Button("프롬프트 생성", variant="primary")
            
            # 우측 컬럼 - 출력 결과
            with gr.Column(scale=1):
                prompt_output = gr.Textbox(label="생성된 프롬프트", lines=10)
                image_preview = gr.Image(label="업로드된 이미지 (#1)", type="pil")
                preview_html = gr.HTML("프롬프트 생성 시 여기에 미리보기가 표시됩니다.")
        
        # 선택된 배경 정보 가져오기
        def get_selected_background_info(active_tab, simple, studio, nature, indoor, abstract):
            if active_tab == "심플 배경" and simple:
                return {
                    "category": "심플 배경",
                    "name": simple,
                    "english": SIMPLE_BACKGROUNDS.get(simple, "white background")
                }
            elif active_tab == "스튜디오 배경" and studio:
                return {
                    "category": "스튜디오 배경",
                    "name": studio,
                    "english": STUDIO_BACKGROUNDS.get(studio, "product photography studio")
                }
            elif active_tab == "자연 환경" and nature:
                return {
                    "category": "자연 환경",
                    "name": nature,
                    "english": NATURE_BACKGROUNDS.get(nature, "natural environment")
                }
            elif active_tab == "실내 환경" and indoor:
                return {
                    "category": "실내 환경",
                    "name": indoor,
                    "english": INDOOR_BACKGROUNDS.get(indoor, "indoor environment")
                }
            elif active_tab == "추상/특수 배경" and abstract:
                return {
                    "category": "추상/특수 배경",
                    "name": abstract,
                    "english": ABSTRACT_BACKGROUNDS.get(abstract, "abstract background")
                }
            else:
                return {
                    "category": "기본 배경",
                    "name": "화이트 배경",
                    "english": "white background"
                }
        
        # 프롬프트 생성 함수
        def generate_output(image, active_tab, simple, studio, nature, indoor, abstract, product_text, additional_text):
            if image is None:
                gr.Warning("이미지를 업로드해주세요.")
                return "이미지를 업로드해주세요.", None, "이미지를 업로드 후 프롬프트를 생성해주세요."
            
            product_text = product_text.strip() or "제품"
            
            # 배경 정보 가져오기
            background_info = get_selected_background_info(active_tab, simple, studio, nature, indoor, abstract)
            
            try:
                prompt = generate_prompt_with_gemini(product_text, background_info, additional_text)
                
                if not GEMINI_API_KEY:
                    gr.Warning("Gemini API 키가 설정되지 않았습니다. 키를 발급받아 설정해주세요.")
                    prompt = """
[Gemini API 키 누락]
API 키 설정 방법:
1. 환경 변수: export GEMINI_API_KEY="your-api-key"
2. 코드 내 직접 입력: GEMINI_API_KEY = "your-api-key"
키 발급: https://makersuite.google.com/
"""
                    return prompt, image, "API 키를 설정해야 합니다."
                
                # 프롬프트 요약 HTML 생성
                preview = f"""
                <div style="padding:10px; border:1px solid #ddd; border-radius:8px; margin-top:10px;">
                    <h3>프롬프트 요약</h3>
                    <p><strong>총 길이:</strong> {len(prompt)} 글자</p>
                    <p><strong>배경:</strong> {background_info['category']} &gt; {background_info['name']}</p>
                    <p><strong>주요 요소:</strong> {", ".join([kw for kw in ["commercial photography", "product", "square format", "centered", "detailed"] if kw.lower() in prompt.lower()])}</p>
                    <p><strong>미드저니 파라미터:</strong> {" ".join([param for param in ["--ar 1:1", "--s 750", "--q 2"] if param in prompt])}</p>
                </div>
                """
                
                return prompt, image, preview
            except Exception as e:
                error_msg = f"프롬프트 생성 중 오류 발생: {str(e)}"
                gr.Error(error_msg)
                return error_msg, image, "오류가 발생했습니다."
        
        # 프롬프트 생성 버튼 클릭 이벤트
        submit_btn.click(
            fn=generate_output,
            inputs=[
                image_input,
                background_tabs,
                simple_dropdown,
                studio_dropdown,
                nature_dropdown,
                indoor_dropdown,
                abstract_dropdown,
                product_name,
                additional_info
            ],
            outputs=[
                prompt_output,
                image_preview,
                preview_html
            ]
        )
    
    return demo

if __name__ == "__main__":
    app = create_app()
    app.launch()