app.py

import os
import logging
import re
import json
import gradio as gr
from google import genai
from google.genai import types
import google.generativeai as genai_generative
from dotenv import load_dotenv

load_dotenv()

# ------------------- 로깅 설정 -------------------
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# ------------------- 배경 디렉토리 설정 -------------------
BACKGROUNDS_DIR = "./background"
if not os.path.exists(BACKGROUNDS_DIR):
    os.makedirs(BACKGROUNDS_DIR)
    logger.info(f"배경 디렉토리를 생성했습니다: {BACKGROUNDS_DIR}")

# ------------------- 전역 변수 설정 -------------------
SIMPLE_BACKGROUNDS = {}
STUDIO_BACKGROUNDS = {}
NATURE_BACKGROUNDS = {}
INDOOR_BACKGROUNDS = {}
TECHNOLOGY_BACKGROUNDS = {}
COLORFUL_PATTERN_BACKGROUNDS = {}
ABSTRACT_BACKGROUNDS = {}
JEWELRY_BACKGROUNDS = {}  # 쥬얼리 배경 전역 변수 추가


# ------------------- 배경 JSON 파일 로드 함수 -------------------
def load_background_json(filename):
    """배경 JSON 파일 로드 함수"""
    file_path = os.path.join(BACKGROUNDS_DIR, filename)
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            logger.info(f"{filename} 파일을 성공적으로 로드했습니다. {len(data)} 항목 포함.")
            return data
    except FileNotFoundError:
        logger.info(f"{filename} 파일이 없습니다.")
        return {}
    except Exception as e:
        logger.warning(f"{filename} 파일 로드 중 오류 발생: {str(e)}.")
        return {}

# ------------------- 배경 옵션 초기화 함수 -------------------
def initialize_backgrounds():
    """모든 배경 옵션 초기화 함수"""
    global SIMPLE_BACKGROUNDS, STUDIO_BACKGROUNDS, NATURE_BACKGROUNDS, INDOOR_BACKGROUNDS
    global TECHNOLOGY_BACKGROUNDS, COLORFUL_PATTERN_BACKGROUNDS, ABSTRACT_BACKGROUNDS
    global JEWELRY_BACKGROUNDS  # 쥬얼리 배경 추가

    
    # 디렉토리 내 모든 파일 로깅
    logger.info(f"Backgrounds 디렉토리 경로: {BACKGROUNDS_DIR}")
    logger.info(f"디렉토리 내 파일 목록: {os.listdir(BACKGROUNDS_DIR)}")
    
    # 각 배경 파일 로드
    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")
    TECHNOLOGY_BACKGROUNDS = load_background_json("tech-backgrounds-final.json")  # 파일 이름 수정
    COLORFUL_PATTERN_BACKGROUNDS = load_background_json("colorful-pattern-backgrounds.json")  # 파일 이름 수정
    ABSTRACT_BACKGROUNDS = load_background_json("abstract_backgrounds.json")
    JEWELRY_BACKGROUNDS = load_background_json("jewelry_backgrounds.json")

    
    # 기본값 설정 (파일이 없거나 비어있는 경우)
    if not SIMPLE_BACKGROUNDS:
        SIMPLE_BACKGROUNDS = {"클래식 화이트": "clean white background with soft even lighting"}
    if not STUDIO_BACKGROUNDS:
        STUDIO_BACKGROUNDS = {"미니멀 플랫레이": "minimalist flat lay with clean white background"}
    if not NATURE_BACKGROUNDS:
        NATURE_BACKGROUNDS = {"열대 해변": "tropical beach with crystal clear water"}
    if not INDOOR_BACKGROUNDS:
        INDOOR_BACKGROUNDS = {"미니멀 스칸디나비안 거실": "minimalist Scandinavian living room"}
    if not TECHNOLOGY_BACKGROUNDS:
        TECHNOLOGY_BACKGROUNDS = {"다이나믹 스플래시": "dynamic water splash interaction with product"}
    if not COLORFUL_PATTERN_BACKGROUNDS:
        COLORFUL_PATTERN_BACKGROUNDS = {"화려한 꽃 패턴": "vibrant floral pattern backdrop"}
    if not ABSTRACT_BACKGROUNDS:
        ABSTRACT_BACKGROUNDS = {"네온 라이트": "neon light abstract background with vibrant glowing elements"}
    if not JEWELRY_BACKGROUNDS:
        JEWELRY_BACKGROUNDS = {"클래식 화이트 실크": "pristine white silk fabric backdrop"}
    
    logger.info("모든 배경 옵션 초기화 완료")

# 배경 드롭다운 초기화를 위한 함수 추가
def initialize_dropdowns():
    """드롭다운 메뉴 초기화 함수"""
    # 각 배경 유형별 드롭다운 선택 목록 생성
    simple_choices = list(SIMPLE_BACKGROUNDS.keys())
    studio_choices = list(STUDIO_BACKGROUNDS.keys())
    nature_choices = list(NATURE_BACKGROUNDS.keys())
    indoor_choices = list(INDOOR_BACKGROUNDS.keys())
    tech_choices = list(TECHNOLOGY_BACKGROUNDS.keys())
    colorful_choices = list(COLORFUL_PATTERN_BACKGROUNDS.keys())
    abstract_choices = list(ABSTRACT_BACKGROUNDS.keys())
    jewelry_choices = list(JEWELRY_BACKGROUNDS.keys())


    
    return {
        "simple": simple_choices,
        "studio": studio_choices,
        "nature": nature_choices,
        "indoor": indoor_choices,
        "tech": tech_choices,
        "colorful": colorful_choices,
        "abstract": abstract_choices,
        "jewelry": jewelry_choices,  # 새로 추가
    }

# ------------------- 프롬프트 관련 함수 -------------------
def filter_prompt_only(prompt):
    """Gemini의 설명 및 불필요한 메시지를 제거하고 실제 프롬프트만 추출하는 함수"""
    
    # 코드 블록 내부의 프롬프트 찾기
    code_block_pattern = r"```\s*(.*?)```"
    code_match = re.search(code_block_pattern, prompt, re.DOTALL)
    if code_match:
        return code_match.group(1).strip()
    
    # Midjourney 파라미터를 포함하는 프롬프트 부분 찾기
    if "--ar 1:1" in prompt:
        lines = prompt.split('\n')
        prompt_lines = []
        in_prompt = False
        
        for line in lines:
            # 프롬프트 시작 부분 인식 (일반적인 제품 설명이나 'Magazine-worthy' 같은 키워드로 시작)
            if (not in_prompt and 
                ("product" in line.lower() or 
                 "magazine" in line.lower() or 
                 "commercial" in line.lower() or 
                 "photography" in line.lower())):
                in_prompt = True
                prompt_lines.append(line)
            # 이미 프롬프트 영역에 있는 경우 계속 추가
            elif in_prompt:
                # 설명이나 메타 텍스트가 시작되면 중단
                if "explanation" in line.lower() or "let me know" in line.lower():
                    break
                prompt_lines.append(line)
        
        # 프롬프트 라인 합치기
        if prompt_lines:
            return '\n'.join(prompt_lines).strip()
    
    # 위 방법으로 찾지 못한 경우 원본 반환
    return prompt.strip()

def get_selected_background_info(bg_type, simple, studio, nature, indoor, tech, colorful, abstract, jewelry):
    """선택된 배경 정보를 가져오는 함수"""
    if bg_type == "심플 배경":
        return {
            "category": "심플 배경",
            "name": simple,
            "english": SIMPLE_BACKGROUNDS.get(simple, "white background")
        }
    elif bg_type == "스튜디오 배경":
        return {
            "category": "스튜디오 배경",
            "name": studio,
            "english": STUDIO_BACKGROUNDS.get(studio, "product photography studio")
        }
    elif bg_type == "자연 환경":
        return {
            "category": "자연 환경",
            "name": nature,
            "english": NATURE_BACKGROUNDS.get(nature, "natural environment")
        }
    elif bg_type == "실내 환경":
        return {
            "category": "실내 환경",
            "name": indoor,
            "english": INDOOR_BACKGROUNDS.get(indoor, "indoor environment")
        }
    elif bg_type == "테크놀로지 배경":
        return {
            "category": "테크놀로지 배경",
            "name": tech,
            "english": TECHNOLOGY_BACKGROUNDS.get(tech, "technology environment")
        }
    elif bg_type == "컬러풀 패턴 배경":
        return {
            "category": "컬러풀 패턴 배경",
            "name": colorful,
            "english": COLORFUL_PATTERN_BACKGROUNDS.get(colorful, "colorful pattern background")
        }
    elif bg_type == "추상/특수 배경":
        return {
            "category": "추상/특수 배경",
            "name": abstract,
            "english": ABSTRACT_BACKGROUNDS.get(abstract, "abstract background")
        }

    elif bg_type == "쥬얼리 배경":
        return {
            "category": "쥬얼리 배경",
            "name": jewelry,
            "english": JEWELRY_BACKGROUNDS.get(jewelry, "jewelry backdrop")
        }

    else:
        return {
            "category": "기본 배경",
            "name": "화이트 배경",
            "english": "white background"
        }

# ------------------- 프롬프트 생성 함수 -------------------
def generate_enhanced_system_instruction():
    """향상된 시스템 인스트럭션 생성 함수"""
    return """당신은 상품 이미지의 배경을 변경하기 위한 최고 품질의 프롬프트를 생성하는 전문가입니다.
사용자가 제공하는 상품명, 배경 유형, 추가 요청사항을 바탕으로 미드저니(Midjourney)에 사용할 수 있는 상세하고 전문적인 프롬프트를 영어로 생성해주세요.

다음 가이드라인을 반드시 준수해야 합니다:

1. 상품을 "#1"로 지정하여 참조합니다. (예: "skincare tube (#1)")
2. *** 매우 중요: 상품의 원래 특성(디자인, 색상, 형태, 로고, 패키지 등)은 어떤 상황에서도 절대 변경하지 않습니다. ***
3. *** 상품의 본질적 특성을 유지하되, 상품에 포커스를 맞춰 모든 세부 사항이 선명하게 드러나도록 하며, 
   8K 해상도(8K resolution), 오버샤프닝 없는 초고화질(ultra high definition without oversharpening)로 렌더링되어야 합니다. ***
4. 이미지 비율은 정확히 1:1(정사각형) 형식으로 지정합니다. 프롬프트에 "square format", "1:1 ratio" 또는 "aspect ratio 1:1"을 명시적으로 포함합니다.
5. 상품은 반드시 정사각형 구도의 정중앙에 배치되어야 하며, 적절한 크기로 표현하여 디테일이 완벽하게 보이도록 합니다.
6. 상품을 이미지의 주요 초점으로 부각시키고, 상품의 비율이 전체 이미지에서 60-70% 이상 차지하도록 합니다.
7. 조명 설명을 매우 구체적으로 해주세요. 예: "soft directional lighting from left side", "dramatic rim lighting", "diffused natural light through windows"
8. 배경의 재질과 질감을 상세히 설명해주세요. 예: "polished marble surface", "rustic wooden table with visible grain", "matte concrete wall with subtle texture"
9. 프롬프트에 다음 요소들을 명시적으로 포함하되, 사용 맥락에 적절하게 변형하세요:
   - "award-winning product photography" 
   - "magazine-worthy commercial product shot"
   - "professional advertising imagery with perfect exposure"
   - "studio lighting with color-accurate rendering"
   - "8K ultra high definition product showcase"
   - "commercial product photography with precise detail rendering"
   - "ultra high definition"
   - "crystal clear details"
10. *** 최우선 지침: 사용자가 제공한 추가 요청사항을 완벽하게 프롬프트에 반영해야 합니다. 사용자의 요청사항 각각을 중요한 요소로 취급하여 프롬프트에 명시적으로 반영하고, 관련 디테일을 구체적으로 설명하세요. 절대로 사용자의 추가 요청사항을 생략하거나 무시하지 마세요. ***
11. 사용자가 제공한 구체적인 배경과 추가 요청사항을 프롬프트에 정확히 반영하고 확장합니다.
12. 프롬프트 끝에 "--ar 1:1 --s 750 --q 2 --v 5.2" 파라미터를 추가하여 미드저니에서 고품질 정사각형 비율을 강제합니다.
13. 매우 중요: 프롬프트 외에 다른 설명이나 메타 텍스트를 포함하지 마세요. 예를 들어 "Here's a prompt for you" 또는 "Let me know if you need adjustments" 같은 메시지나 설명을 포함하지 마세요. 오직 프롬프트 자체만 제공하세요.
"""

def generate_prompt_with_gemini(product_name, background_info, additional_info=""):
    """향상된 프롬프트 생성 함수"""
    GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
    if not GEMINI_API_KEY:
        return "Gemini API 키가 설정되지 않았습니다. 환경 변수 GEMINI_API_KEY를 설정하거나 코드에 직접 입력하세요."
    
    try:
        genai_generative.configure(api_key=GEMINI_API_KEY)
        
        # 추가 요청사항을 더 강조한 프롬프트 요청 템플릿
        prompt_request = f"""
상품명: {product_name}
배경 유형: {background_info.get('english', 'studio')}
배경 카테고리: {background_info.get('category', '')}
배경 이름: {background_info.get('name', '')}
추가 요청사항: {additional_info}

중요 요구사항:
1. 상품(#1)이 이미지 구도에서 중심적인 위치를 차지하며 적절한 크기(이미지의 60-70%)로 표현되도록 프롬프트를 생성해주세요.
2. 이미지는 정확히 1:1 비율(정사각형)이어야 합니다.
3. 상품의 디자인, 색상, 형태, 로고 등 본질적 특성은 절대 수정하지 마세요.
4. 구체적인 조명 기법을 상세히 명시해주세요:
   - 정확한 조명 위치 (예: "45-degree key light from upper left")
   - 조명 품질 (예: "soft diffused light", "hard directional light")
   - 조명 강도와 색온도 (예: "warm tungsten key light with cool blue fill")
   - 반사와 그림자 처리 방식 (예: "controlled specular highlights with soft shadow transitions")
5. 상품을 더 돋보이게 하는 보조 요소(props)를 자연스럽게 활용하되, 상품이 항상 주인공이어야 합니다.
6. 배경 재질과 표면 질감을 구체적으로 설명하고, 상품과의 상호작용 방식을 명시해주세요.
7. 색상 구성(color palette, color harmonies)을 명확히 해주세요.
8. 고급스러운 상업 광고 품질의 이미지가 되도록 프롬프트를 작성해주세요.
9. 프롬프트 끝에 미드저니 파라미터 "--ar 1:1 --s 750 --q 2 --v 5.2"를 추가해주세요.
10. **매우 중요**: 사용자가 제공한 추가 요청사항("{additional_info}")을 프롬프트에 완벽하게 반영해야 합니다. 각 요청사항을 프롬프트의 주요 부분에 명시적으로 포함시키고, 관련 디테일을 자세히 기술해주세요.

한국어 입력 내용을 전문적인 영어로 번역하여 반영해주세요.
"""
        # 추가 요청사항을 강조한 시스템 인스트럭션 생성
        enhanced_system_instruction = generate_enhanced_system_instruction() + f"""
        
특별 지침: 사용자의 추가 요청사항은 최우선 순위로 처리해야 합니다. 추가 요청사항("{additional_info}")의 각 항목을 프롬프트에 명시적으로 포함시키고, 해당 요청에 관련된 구체적인 설명을 추가하세요. 사용자의 요청사항을 무시하거나 생략하지 마세요.
"""
        
        # 더 창의적인 결과를 위한 생성 설정 조정
        model = genai_generative.GenerativeModel(
            'gemini-2.0-flash',
            system_instruction=enhanced_system_instruction
        )
        
        response = model.generate_content(
            prompt_request,
            generation_config=genai_generative.types.GenerationConfig(
                temperature=0.7,      # 정확성을 위해 약간 낮춤
                top_p=0.95,
                top_k=40,
                max_output_tokens=1600,  # 더 상세한 프롬프트 허용
            )
        )
        
        response_text = response.text.strip()
        
        # 추가 요청사항이 프롬프트에 제대로 반영되었는지 확인하는 로직
        if additional_info and not any(keyword.lower() in response_text.lower() for keyword in additional_info.split(',')):
            # 추가 요청사항이 반영되지 않았다면 추가
            logger.warning("추가 요청사항이 프롬프트에 반영되지 않았습니다. 강제로 추가합니다.")
            
            # 요청사항을 직접 프롬프트에 추가 (메타 텍스트 없이)
            if "--ar 1:1" in response_text:
                # 미드저니 파라미터 앞에 추가
                parts = response_text.split("--ar 1:1")
                if len(parts) >= 2:
                    # 추가 요청사항 직접 삽입
                    response_text = parts[0].rstrip(" ,") + ", " + additional_info + ". --ar 1:1" + parts[1]
            else:
                # 미드저니 파라미터가 없는 경우 끝에 추가
                response_text = response_text.rstrip(" .") + ". " + additional_info
            
        # 미드저니 파라미터가 없을 경우 추가
        if "--ar 1:1" not in response_text:
            response_text = response_text.rstrip(" .") + ". --ar 1:1 --s 750 --q 2 --v 5.2"
            
        return response_text
    except Exception as e:
        logger.exception("프롬프트 생성 중 오류 발생:")
        return f"프롬프트 생성 중 오류가 발생했습니다: {str(e)}"

# ------------------- 프롬프트 생성 전용 함수 -------------------
def generate_prompt_only(image, bg_type, simple, studio, nature, indoor, tech, colorful, abstract, jewelry, product_name, additional_info):
    product_name = product_name.strip() or "제품"
    background_info = get_selected_background_info(bg_type, simple, studio, nature, indoor, tech, colorful, abstract, jewelry)
    generated_prompt = generate_prompt_with_gemini(product_name, background_info, additional_info)
    if "Gemini API 키가 설정되지 않았습니다" in generated_prompt:
        warning_msg = (
            "[Gemini API 키 누락]\n"
            "API 키 설정 방법:\n"
            "1. 환경 변수: export GEMINI_API_KEY=\"your-api-key\"\n"
            "2. 코드 내 직접 입력: GEMINI_API_KEY = \"your-api-key\"\n"
            "키 발급: https://makersuite.google.com/"
        )
        return warning_msg
    final_prompt = filter_prompt_only(generated_prompt)
    return final_prompt

# ------------------- Gradio 인터페이스 구성 -------------------
def create_app():
    # 드롭다운 옵션 초기화
    dropdown_options = initialize_dropdowns()
    
    with gr.Blocks(title="고급 상품 이미지 배경 프롬프트 생성") as demo:
        gr.Markdown("# 고급 상품 이미지 배경 프롬프트 생성")
        gr.Markdown(
            "상품 이미지를 업로드하고, 제품명, 배경 옵션, 추가 요청사항을 입력하면 Gemini API를 통해 영어 프롬프트를 생성합니다."
        )
        with gr.Row():
            with gr.Column(scale=1):
                product_name = gr.Textbox(label="상품명 (한국어 입력)", placeholder="예: 스킨케어 튜브, 스마트워치, 향수, 운동화 등", interactive=True)
                # 상품 이미지 업로드를 상품명 아래로 이동
                image_input = gr.Image(label="상품 이미지 업로드", type="pil")
                
                background_type = gr.Radio(
                    choices=["심플 배경", "스튜디오 배경", "자연 환경", "실내 환경", "테크놀로지 배경", "컬러풀 패턴 배경", "추상/특수 배경", "쥬얼리 배경"],
                    label="배경 유형",
                    value="심플 배경"
                )
                simple_dropdown = gr.Dropdown(
                    choices=dropdown_options["simple"],
                    value=dropdown_options["simple"][0] if dropdown_options["simple"] else None,
                    label="심플 배경 선택",
                    visible=True,
                    interactive=True
                )
                studio_dropdown = gr.Dropdown(
                    choices=dropdown_options["studio"],
                    value=dropdown_options["studio"][0] if dropdown_options["studio"] else None,
                    label="스튜디오 배경 선택",
                    visible=False,
                    interactive=True
                )
                nature_dropdown = gr.Dropdown(
                    choices=dropdown_options["nature"],
                    value=dropdown_options["nature"][0] if dropdown_options["nature"] else None,
                    label="자연 환경 선택",
                    visible=False,
                    interactive=True
                )
                indoor_dropdown = gr.Dropdown(
                    choices=dropdown_options["indoor"],
                    value=dropdown_options["indoor"][0] if dropdown_options["indoor"] else None,
                    label="실내 환경 선택",
                    visible=False,
                    interactive=True
                )
                tech_dropdown = gr.Dropdown(
                    choices=dropdown_options["tech"],
                    value=dropdown_options["tech"][0] if dropdown_options["tech"] else None,
                    label="테크놀로지 배경 선택",
                    visible=False,
                    interactive=True
                )
                colorful_dropdown = gr.Dropdown(
                    choices=dropdown_options["colorful"],
                    value=dropdown_options["colorful"][0] if dropdown_options["colorful"] else None,
                    label="컬러풀 패턴 배경 선택",
                    visible=False,
                    interactive=True
                )
                abstract_dropdown = gr.Dropdown(
                    choices=dropdown_options["abstract"],
                    value=dropdown_options["abstract"][0] if dropdown_options["abstract"] else None,
                    label="추상/특수 배경 선택",
                    visible=False,
                    interactive=True
                )
                
                jewelry_dropdown = gr.Dropdown(
                    choices=dropdown_options["jewelry"],
                    value=dropdown_options["jewelry"][0] if dropdown_options["jewelry"] else None,
                    label="쥬얼리 배경 선택",
                    visible=False,
                    interactive=True
                )
        
                additional_info = gr.Textbox(
                    label="추가 요청사항 (선택사항)",
                    placeholder="예: 고급스러운 느낌, 밝은 조명, 자연스러운 보조 객체 등",
                    lines=3,
                    interactive=True
                )
                def update_dropdowns(bg_type):
                    return {
                        simple_dropdown: gr.update(visible=(bg_type == "심플 배경")),
                        studio_dropdown: gr.update(visible=(bg_type == "스튜디오 배경")),
                        nature_dropdown: gr.update(visible=(bg_type == "자연 환경")),
                        indoor_dropdown: gr.update(visible=(bg_type == "실내 환경")),
                        tech_dropdown: gr.update(visible=(bg_type == "테크놀로지 배경")),
                        colorful_dropdown: gr.update(visible=(bg_type == "컬러풀 패턴 배경")),
                        abstract_dropdown: gr.update(visible=(bg_type == "추상/특수 배경")),
                        jewelry_dropdown: gr.update(visible=(bg_type == "쥬얼리 배경"))  # 쥬얼리 드롭다운 추가
                        
                    }
                background_type.change(
                    fn=update_dropdowns,
                    inputs=[background_type],
                    outputs=[simple_dropdown, studio_dropdown, nature_dropdown, indoor_dropdown, tech_dropdown, colorful_dropdown, abstract_dropdown, jewelry_dropdown]
                )
                
                # 프롬프트 생성 버튼으로 변경
                prompt_btn = gr.Button("프롬프트 생성", variant="primary")
            
            with gr.Column(scale=1):
                # 프롬프트 출력 영역만 유지
                prompt_output = gr.Textbox(label="생성된 프롬프트 (영어)", lines=10)
        
        # 프롬프트 생성 함수 연결
        prompt_btn.click(
            fn=generate_prompt_only,
            inputs=[image_input, background_type, simple_dropdown, studio_dropdown, nature_dropdown, indoor_dropdown, tech_dropdown, colorful_dropdown, abstract_dropdown, jewelry_dropdown, product_name, additional_info],
            outputs=prompt_output
        )
    
    return demo

# ------------------- 메인 실행 함수 -------------------
if __name__ == "__main__":
    # 배경 옵션 초기화 - JSON 파일에서 로드
    initialize_backgrounds()
    
    # 앱 생성 및 실행
    app = create_app()
    app.queue()
    app.launch()