app.py

import os
import sys
import base64
import io
import logging
import tempfile
import traceback
import requests
from PIL import Image
import gradio as gr
from gradio_client import Client, handle_file
from dotenv import load_dotenv

# 환경변수 로드
load_dotenv()

# 로깅 설정 (민감한 정보 제외)
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("app.log"),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger("image-enhancer-app")

# API 클라이언트 초기화 (환경변수에서)
API_ENDPOINT = os.environ.get("API_ENDPOINT")
if not API_ENDPOINT:
    logger.error("API_ENDPOINT environment variable is not set")
    sys.exit(1)

try:
    # 로그에 엔드포인트 정보 출력하지 않음
    api_client = Client(API_ENDPOINT)
    logger.info("API client initialized successfully")
except Exception as e:
    logger.error(f"Failed to initialize API client: {e}")
    sys.exit(1)

# backgrounds.py 전체 내용을 환경변수에서 가져오기
def load_backgrounds_from_env():
    """환경변수에서 backgrounds.py 전체 내용 로드"""
    backgrounds_code = os.environ.get("BACKGROUNDS_DATA", "")
    if not backgrounds_code:
        logger.warning("BACKGROUNDS_DATA environment variable is empty")
        return {}, {}, {}, {}, {}, {}, {}
    
    try:
        # 환경변수에서 가져온 코드 실행하여 딕셔너리들 추출
        local_vars = {}
        exec(backgrounds_code, {}, local_vars)
        
        return (
            local_vars.get("SIMPLE_BACKGROUNDS", {}),
            local_vars.get("STUDIO_BACKGROUNDS", {}),
            local_vars.get("NATURE_BACKGROUNDS", {}),
            local_vars.get("INDOOR_BACKGROUNDS", {}),
            local_vars.get("SPECIAL_BACKGROUNDS", {}),
            local_vars.get("JEWELRY_BACKGROUNDS", {}),
            local_vars.get("SPECIAL_EFFECTS_BACKGROUNDS", {})
        )
    except Exception as e:
        logger.error(f"Failed to parse BACKGROUNDS_DATA: {e}")
        return {}, {}, {}, {}, {}, {}, {}

# 환경변수에서 배경 딕셔너리들 로드
(SIMPLE_BACKGROUNDS, STUDIO_BACKGROUNDS, NATURE_BACKGROUNDS, 
 INDOOR_BACKGROUNDS, SPECIAL_BACKGROUNDS, JEWELRY_BACKGROUNDS, 
 SPECIAL_EFFECTS_BACKGROUNDS) = load_backgrounds_from_env()

# 만약 환경변수가 비어있다면 기본값 사용
if not SIMPLE_BACKGROUNDS:
    SIMPLE_BACKGROUNDS = {
        "화이트 기본": "clean white background with soft lighting",
        "회색 투톤": "light gray background with minimal shadows",
        "라이트 그레이": "soft light gray backdrop with even illumination",
        "그레이 그라데이션 스포트라이트": "gray gradient with spotlight",
        "프리미엄 드라마틱 블랙": "premium dramatic black background",
        "딥블루 유리반사": "deep blue glass reflection",
        "파스텔 그라데이션": "pastel gradient background",
        "스카이블루 파스텔": "sky blue pastel",
        "버터옐로우 파스텔": "butter yellow pastel",
        "블루 원색": "pure blue background",
        "레드 원색": "pure red background"
    }

if not STUDIO_BACKGROUNDS:
    STUDIO_BACKGROUNDS = {
        "연녹색 장미 정원": "soft green rose garden",
        "연분홍 장미 대리석": "pink rose with marble",
        "파스텔 블루 수국": "pastel blue hydrangea",
        "소프트 핑크 벚꽃": "soft pink cherry blossom",
        "파스텔 옐로우 꽃봉오리": "pastel yellow buds",
        "아이보리 로즈 심플": "ivory rose simple",
        "심플 화이트 작은 화분": "simple white small pot",
        "퍼플 장난스러운 심비로운 분위기": "purple playful mood",
        "코지 데스크 플랜트": "cozy desk plant",
        "연핑크 튤립 작은테이블": "light pink tulip small table",
        "오크우드 튤립 샤워 푸프": "oak wood tulip shower puff",
        "크림 코튼": "cream cotton",
        "브라운 거친 페이퍼": "brown rough paper"
    }

if not NATURE_BACKGROUNDS:
    NATURE_BACKGROUNDS = {
        "작은 파도가 있는 해변": "beach with small waves",
        "펭귄이 있는 빙산": "iceberg with penguin",
        "눈내리는 산악 설원": "snowy mountain field",
        "열대해변": "tropical beach",
        "일출 직전의 바위산": "rocky mountain before sunrise",
        "초근접 봄꽃 들판": "close-up spring flower field",
        "물가 바위": "waterside rocks",
        "작은 돌이 깔린 얕은 물가": "shallow water with small stones",
        "나무 테이블 숲속 계곡": "wooden table forest valley",
        "안개 낀 숲": "foggy forest",
        "일몰 바다": "sunset sea",
        "별이 빛나는 캠핑": "starry camping"
    }

if not INDOOR_BACKGROUNDS:
    INDOOR_BACKGROUNDS = {
        "기본 책상": "basic desk",
        "빛이 비치는 책상": "desk with light",
        "빛이 비치는 거실": "living room with light",
        "스튜디어 거실": "studio living room",
        "화분이 있는 거실": "living room with plants",
        "전체적인 거실모습": "overall living room view",
        "포인트 거실": "point living room",
        "중앙 거실": "central living room",
        "서랍테이블": "drawer table",
        "침실 탁자위": "bedroom table top",
        "원목 테이블 블록": "wood table block",
        "테이블 심플": "simple table",
        "분위기 있는 침실": "atmospheric bedroom",
        "나무테이블위": "on wooden table",
        "책위 거실테이블": "living room table with books",
        "유리테이블위": "on glass table",
        "침실 옆 작은탁자": "small bedside table",
        "분위기 있는 원형 테이블": "atmospheric round table",
        "포커스 탁자": "focus table",
        "투톤 벽면": "two-tone wall",
        "아늑한 테이블위 액세서리": "cozy table accessories",
        "아늑한 테이블 보 위": "on cozy tablecloth",
        "유리 바닥": "glass floor",
        "빛치는 식탁위": "dining table with light",
        "나늑한 식탁위 작업대": "cozy dining workspace"
    }

if not SPECIAL_BACKGROUNDS:
    SPECIAL_BACKGROUNDS = {
        "네이비 빈티지 플로럴 벽지": "navy vintage floral wallpaper",
        "빈티지 꽃무늬 패브릭 배경": "vintage floral fabric background",
        "팝 도트와 과일": "pop dots and fruits",
        "블루 마블 잉크 텍스처 배경": "blue marble ink texture",
        "오렌지블루 워터컬러 배경": "orange blue watercolor",
        "레드브러시 페인팅 아트 배경": "red brush painting art",
        "블루오렌지 페인팅 아트": "blue orange painting art"
    }

if not JEWELRY_BACKGROUNDS:
    JEWELRY_BACKGROUNDS = {
        "화이트 미러 스팟 라이트": "white mirror spotlight",
        "그레이 그라데이션 미러": "gray gradient mirror",
        "네이비 벨벳": "navy velvet",
        "블랙 미러 시네마틱": "black mirror cinematic",
        "화이트 마블 프리미엄": "white marble premium",
        "파스텔 블루 큐브 플랫폼": "pastel blue cube platform",
        "내추럴 그라스": "natural grass",
        "소프트 베이직 패브릭": "soft basic fabric",
        "마이크로텍스쳐 프리미엄": "micro texture premium"
    }

if not SPECIAL_EFFECTS_BACKGROUNDS:
    SPECIAL_EFFECTS_BACKGROUNDS = {
        "블루블랙 큰 물방울 효과": "blue black water drop effect",
        "크리스탈 버블 물속 장면": "crystal bubble underwater scene",
        "잔 물결 수면 위 장면": "gentle ripple water surface scene",
        "컬러 스모크 효과": "color smoke effect",
        "자옥한 안개 효과": "dense fog effect",
        "가습기 수중기 효과": "humidifier mist effect"
    }

# 임시 파일 저장 함수 (원본과 동일)
def save_uploaded_file(uploaded_file, suffix='.png'):
    try:
        logger.info(f"Processing uploaded file: {type(uploaded_file)}")
        
        if uploaded_file is None:
            logger.warning("Uploaded file is None")
            return None
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
            temp_filename = temp_file.name
            logger.info(f"Created temporary file: {temp_filename}")
        
        # Gradio 업로드 파일 처리
        if isinstance(uploaded_file, str):  # 이미 파일 경로인 경우
            logger.info(f"Uploaded file is already a path: {uploaded_file}")
            return uploaded_file
        
        # PIL Image 처리
        if isinstance(uploaded_file, Image.Image):
            logger.info("Uploaded file is a PIL Image")
            uploaded_file.save(temp_filename, format="PNG")
            return temp_filename
        
        # 바이너리 데이터 처리
        with open(temp_filename, "wb") as f:
            if hasattr(uploaded_file, "read"):  # 파일 객체인 경우
                logger.info("Processing file object")
                content = uploaded_file.read()
                f.write(content)
                logger.info(f"Wrote {len(content)} bytes to {temp_filename}")
            else:  # 바이너리 데이터인 경우
                logger.info("Processing binary data")
                f.write(uploaded_file)
                logger.info(f"Wrote data to {temp_filename}")
        
        return temp_filename
    except Exception as e:
        logger.error(f"Error saving uploaded file: {e}")
        logger.error(traceback.format_exc())
        return None

# Gradio 인터페이스 구성 (원본과 완전 동일)
def create_gradio_interface():
    try:
        logger.info("Creating Gradio interface")
        with gr.Blocks(title="AI 이미지 편집 및 화질 개선") as app:
            gr.Markdown("# AI 이미지 편집 및 화질 개선 도구")
            
            # 비밀번호 입력 필드
            password_box = gr.Textbox(
                label="비밀번호", 
                type="password", 
                placeholder="사용하려면 비밀번호를 입력하세요",
                interactive=True
            )
            
            # 이미지 편집 및 화질 개선 인터페이스
            with gr.Row():
                with gr.Column():
                    # 상품 이미지 업로드
                    image = gr.Image(label="상품 이미지 업로드", type="pil")
                    
                    with gr.Row():
                        with gr.Column():
                            background_type = gr.Radio(
                                choices=["심플 배경", "스튜디오 배경", "자연 환경", "실내 환경", "특수배경", "주얼리", "특수효과"],
                                label="배경 유형",
                                value="심플 배경"
                            )
                            
                            # 드롭다운 컴포넌트들
                            simple_dropdown = gr.Dropdown(
                                choices=list(SIMPLE_BACKGROUNDS.keys()),
                                value=list(SIMPLE_BACKGROUNDS.keys())[0] if SIMPLE_BACKGROUNDS else None,
                                label="심플 배경 선택",
                                visible=True,
                                interactive=True
                            )
                            
                            studio_dropdown = gr.Dropdown(
                                choices=list(STUDIO_BACKGROUNDS.keys()),
                                value=list(STUDIO_BACKGROUNDS.keys())[0] if STUDIO_BACKGROUNDS else None,
                                label="스튜디오 배경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            nature_dropdown = gr.Dropdown(
                                choices=list(NATURE_BACKGROUNDS.keys()),
                                value=list(NATURE_BACKGROUNDS.keys())[0] if NATURE_BACKGROUNDS else None,
                                label="자연 환경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            indoor_dropdown = gr.Dropdown(
                                choices=list(INDOOR_BACKGROUNDS.keys()),
                                value=list(INDOOR_BACKGROUNDS.keys())[0] if INDOOR_BACKGROUNDS else None,
                                label="실내 환경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            special_dropdown = gr.Dropdown(
                                choices=list(SPECIAL_BACKGROUNDS.keys()),
                                value=list(SPECIAL_BACKGROUNDS.keys())[0] if SPECIAL_BACKGROUNDS else None,
                                label="특수배경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            jewelry_dropdown = gr.Dropdown(
                                choices=list(JEWELRY_BACKGROUNDS.keys()),
                                value=list(JEWELRY_BACKGROUNDS.keys())[0] if JEWELRY_BACKGROUNDS else None,
                                label="주얼리 배경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            special_effects_dropdown = gr.Dropdown(
                                choices=list(SPECIAL_EFFECTS_BACKGROUNDS.keys()),
                                value=list(SPECIAL_EFFECTS_BACKGROUNDS.keys())[0] if SPECIAL_EFFECTS_BACKGROUNDS else None,
                                label="특수효과 배경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            # 드롭다운 변경 함수 (원본과 동일, API 엔드포인트 생성)
                            def update_dropdowns(bg_type):
                                try:
                                    # API 호출
                                    result = api_client.predict(bg_type=bg_type, api_name="/update_dropdowns")
                                    return result
                                except:
                                    # 백업용 로컬 로직
                                    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 == "실내 환경")),
                                        special_dropdown: gr.update(visible=(bg_type == "특수배경")),
                                        jewelry_dropdown: gr.update(visible=(bg_type == "주얼리")),
                                        special_effects_dropdown: gr.update(visible=(bg_type == "특수효과"))
                                    }
                            
                            background_type.change(
                                fn=update_dropdowns,
                                inputs=[background_type],
                                outputs=[simple_dropdown, studio_dropdown, nature_dropdown, indoor_dropdown, special_dropdown, jewelry_dropdown, special_effects_dropdown],
                                api_name="update_dropdowns"
                            )
                    
                    # 요청사항 입력
                    request_text = gr.Textbox(
                        label="요청사항", 
                        placeholder="상품 이미지에 적용할 스타일, 분위기, 특별 요청사항 등을 입력하세요.", 
                        lines=3
                    )
                    
                    # 새로운 옵션들
                    quality_level = gr.Radio(
                        label="품질 레벨", 
                        choices=["gpt", "flux"],
                        value="flux",
                        info="GPT: GPT 모델 (고품질), 일반: Flux 모델 (빠른 처리 + 기본 화질개선)"
                    )
                    
                    aspect_ratio = gr.Dropdown(
                        label="종횡비", 
                        choices=["1:1", "3:2", "2:3"],
                        value="1:1"
                    )
                    
                    output_format = gr.Dropdown(
                        label="이미지 형식", 
                        choices=["jpg", "png"],
                        value="jpg"
                    )
                    
                    # 화질 개선 옵션
                    enable_enhancement = gr.Checkbox(
                        label="추가 화질 개선", 
                        value=False,
                        info="GPT: 1회 화질개선, Flux: 2차 화질개선 (기본 1회 + 추가 1회)"
                    )
                    enhancement_level = gr.Slider(label="화질 개선 레벨", minimum=1, maximum=4, value=2, step=1, visible=False)
                    
                    # 프롬프트 생성 버튼
                    generate_prompt_btn = gr.Button("프롬프트만 생성")
                    
                    # 편집 버튼
                    edit_btn = gr.Button("이미지 편집 및 화질 개선")
                
                with gr.Column():
                    with gr.Row():
                        with gr.Column():
                            gr.Markdown("## 편집된 이미지")
                            original_output = gr.Gallery(label="편집 결과", preview=True)
                            original_download = gr.File(label="편집 이미지 다운로드", interactive=False)
                        
                        with gr.Column():
                            gr.Markdown("## 화질 개선된 이미지")
                            enhanced_output = gr.Gallery(label="화질 개선 결과", preview=True)
                            enhanced_download = gr.File(label="개선 이미지 다운로드", interactive=False)
                    
                    # 프롬프트 출력
                    prompt_output = gr.Textbox(label="생성된 프롬프트", lines=10, interactive=False)
                    info = gr.Textbox(label="처리 정보", interactive=False)
                    error = gr.Textbox(label="오류 메시지", interactive=False, visible=True)
            
            # 프롬프트만 생성하는 함수 (원본과 동일한 엔드포인트, API 호출)
            def generate_prompt_with_password_check(password, bg_type, simple, studio, nature, indoor, special, jewelry, special_effects, request_text, aspect_ratio):
                try:
                    return api_client.predict(
                        password=password,
                        bg_type=bg_type,
                        simple=simple,
                        studio=studio,
                        nature=nature,
                        indoor=indoor,
                        special=special,
                        jewelry=jewelry,
                        special_effects=special_effects,
                        request_text=request_text,
                        aspect_ratio=aspect_ratio,
                        api_name="/generate_prompt_with_password_check"
                    )
                except Exception as e:
                    return f"API 호출 오류: {str(e)}"
            
            # 비밀번호 확인 함수 (원본과 동일한 엔드포인트, API 호출)
            def check_password(password, *args):
                try:
                    image, bg_type, simple, studio, nature, indoor, special, jewelry, special_effects, request_text, quality_level, aspect_ratio, output_format, enable_enhancement = args
                    
                    if image is None:
                        return ([], None, [], None, "", "", "이미지를 업로드해야 합니다.")
                    
                    # 이미지를 임시 파일로 저장
                    temp_path = save_uploaded_file(image)
                    if temp_path is None:
                        return ([], None, [], None, "", "", "이미지 처리에 실패했습니다.")
                    
                    try:
                        # API 호출 (원본과 동일한 파라미터)
                        result = api_client.predict(
                            password=password,
                            param_1=handle_file(temp_path),
                            param_2=bg_type,
                            param_3=simple,
                            param_4=studio,
                            param_5=nature,
                            param_6=indoor,
                            param_7=special,
                            param_8=jewelry,
                            param_9=special_effects,
                            param_10=request_text,
                            param_11=quality_level,
                            param_12=aspect_ratio,
                            param_13=output_format,
                            param_14=enable_enhancement,
                            api_name="/check_password"
                        )
                        return result
                    finally:
                        # 임시 파일 정리
                        if os.path.exists(temp_path):
                            try:
                                os.remove(temp_path)
                            except:
                                pass
                                
                except Exception as e:
                    logger.error(f"Error in check_password: {e}")
                    return ([], None, [], None, "", "", f"오류 발생: {str(e)}")
            
            # 프롬프트 생성 버튼 클릭 이벤트
            generate_prompt_btn.click(
                fn=generate_prompt_with_password_check,
                inputs=[
                    password_box,
                    background_type,
                    simple_dropdown, studio_dropdown, nature_dropdown, indoor_dropdown, special_dropdown,
                    jewelry_dropdown, special_effects_dropdown,
                    request_text, aspect_ratio
                ],
                outputs=[prompt_output],
                api_name="generate_prompt_with_password_check"
            )
            
            # 편집 버튼 클릭 이벤트
            edit_btn.click(
                fn=check_password,
                inputs=[
                    password_box, 
                    image, background_type, 
                    simple_dropdown, studio_dropdown, nature_dropdown, indoor_dropdown, special_dropdown,
                    jewelry_dropdown, special_effects_dropdown,
                    request_text, quality_level, aspect_ratio, output_format, enable_enhancement
                ],
                outputs=[
                    original_output, original_download, 
                    enhanced_output, enhanced_download,
                    prompt_output, info, error
                ],
                api_name="check_password"
            )
        
        logger.info("Gradio interface created successfully")
        return app
    except Exception as e:
        logger.error(f"Error creating Gradio interface: {e}")
        logger.error(traceback.format_exc())
        raise

# 앱 실행
if __name__ == "__main__":
    try:
        logger.info("Starting application")
        
        # imgs 디렉토리 확인/생성
        os.makedirs("imgs", exist_ok=True)
        logger.info("이미지 디렉토리 준비 완료")
        
        app = create_gradio_interface()
        logger.info("Launching Gradio app")
        app.launch(share=True)
    except Exception as e:
        logger.error(f"Error running app: {e}")
        logger.error(traceback.format_exc())