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
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-control-tower")

# 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:
    # 로그에 엔드포인트 정보 출력하지 않음
    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")
        
        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

# API 호출 함수들 (클라이언트를 통해 원본 API 호출)
def api_update_dropdowns(bg_type):
    """드롭다운 업데이트 API 호출"""
    try:
        result = client.predict(bg_type=bg_type, api_name="/update_dropdowns")
        return result
    except Exception as e:
        logger.error(f"Error calling update_dropdowns API: {e}")
        return None

def api_generate_prompt_with_password_check(password, bg_type, simple, studio, nature, indoor, special, jewelry, special_effects, request_text, aspect_ratio):
    """프롬프트 생성 API 호출"""
    try:
        result = 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"
        )
        return result
    except Exception as e:
        logger.error(f"Error calling generate_prompt API: {e}")
        return f"API 호출 오류: {str(e)}"

def api_check_password(password, image_file, bg_type, simple, studio, nature, indoor, special, jewelry, special_effects, request_text, quality_level, aspect_ratio, output_format, enable_enhancement):
    """이미지 처리 API 호출"""
    try:
        result = client.predict(
            password=password,
            param_1=image_file,
            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
    except Exception as e:
        logger.error(f"Error calling check_password API: {e}")
        return None, None, None, None, "", "", f"API 호출 오류: {str(e)}"

# Gradio 인터페이스 구성 (완전히 동일한 UI)
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()) if SIMPLE_BACKGROUNDS else ["화이트 기본"],
                                value=list(SIMPLE_BACKGROUNDS.keys())[0] if SIMPLE_BACKGROUNDS else "화이트 기본",
                                label="심플 배경 선택",
                                visible=True,
                                interactive=True
                            )
                            
                            studio_dropdown = gr.Dropdown(
                                choices=list(STUDIO_BACKGROUNDS.keys()) if STUDIO_BACKGROUNDS else ["연녹색 장미 정원"],
                                value=list(STUDIO_BACKGROUNDS.keys())[0] if STUDIO_BACKGROUNDS else "연녹색 장미 정원",
                                label="스튜디오 배경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            nature_dropdown = gr.Dropdown(
                                choices=list(NATURE_BACKGROUNDS.keys()) if NATURE_BACKGROUNDS else ["작은 파도가 있는 해변"],
                                value=list(NATURE_BACKGROUNDS.keys())[0] if NATURE_BACKGROUNDS else "작은 파도가 있는 해변",
                                label="자연 환경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            indoor_dropdown = gr.Dropdown(
                                choices=list(INDOOR_BACKGROUNDS.keys()) if INDOOR_BACKGROUNDS else ["기본 책상"],
                                value=list(INDOOR_BACKGROUNDS.keys())[0] if INDOOR_BACKGROUNDS else "기본 책상",
                                label="실내 환경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            special_dropdown = gr.Dropdown(
                                choices=list(SPECIAL_BACKGROUNDS.keys()) if SPECIAL_BACKGROUNDS else ["네이비 빈티지 플로럴 벽지"],
                                value=list(SPECIAL_BACKGROUNDS.keys())[0] if SPECIAL_BACKGROUNDS else "네이비 빈티지 플로럴 벽지",
                                label="특수배경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            jewelry_dropdown = gr.Dropdown(
                                choices=list(JEWELRY_BACKGROUNDS.keys()) if JEWELRY_BACKGROUNDS else ["화이트 미러 스팟 라이트"],
                                value=list(JEWELRY_BACKGROUNDS.keys())[0] if JEWELRY_BACKGROUNDS else "화이트 미러 스팟 라이트",
                                label="주얼리 배경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            special_effects_dropdown = gr.Dropdown(
                                choices=list(SPECIAL_EFFECTS_BACKGROUNDS.keys()) if SPECIAL_EFFECTS_BACKGROUNDS else ["블루블랙 큰 물방울 효과"],
                                value=list(SPECIAL_EFFECTS_BACKGROUNDS.keys())[0] if SPECIAL_EFFECTS_BACKGROUNDS else "블루블랙 큰 물방울 효과",
                                label="특수효과 배경 선택",
                                visible=False,
                                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 == "실내 환경")),
                                    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]
                            )
                    
                    # 요청사항 입력
                    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_wrapper(password, bg_type, simple, studio, nature, indoor, special, jewelry, special_effects, request_text, aspect_ratio):
                return api_generate_prompt_with_password_check(
                    password, bg_type, simple, studio, nature, indoor, special, jewelry, special_effects, request_text, aspect_ratio
                )
            
            # 이미지 처리 함수 (API 호출)
            def process_image_wrapper(password, image, bg_type, simple, studio, nature, indoor, special, jewelry, special_effects, request_text, quality_level, aspect_ratio, output_format, enable_enhancement):
                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_check_password(
                        password, temp_path, bg_type, simple, studio, nature, indoor, special, jewelry, special_effects, 
                        request_text, quality_level, aspect_ratio, output_format, enable_enhancement
                    )
                    return result
                finally:
                    # 임시 파일 정리
                    if os.path.exists(temp_path):
                        try:
                            os.remove(temp_path)
                        except:
                            pass
            
            # 프롬프트 생성 버튼 클릭 이벤트
            generate_prompt_btn.click(
                fn=generate_prompt_wrapper,
                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]
            )
            
            # 편집 버튼 클릭 이벤트
            edit_btn.click(
                fn=process_image_wrapper,
                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
                ]
            )
        
        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 control tower 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()), None, None, None, None, None, None
    
    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 None, None, None, None, None, None, None

# 환경변수에서 배경 딕셔너리들 로드
(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"
    }

if not STUDIO_BACKGROUNDS:
    STUDIO_BACKGROUNDS = {
        "연녹색 장미 정원": "soft green background with rose garden elements",
        "연분홍 장미 대리석": "pink rose with marble surface backdrop"
    }

if not NATURE_BACKGROUNDS:
    NATURE_BACKGROUNDS = {
        "작은 파도가 있는 해변": "serene beach with gentle waves",
        "열대해변": "tropical beach setting with palm trees"
    }

if not INDOOR_BACKGROUNDS:
    INDOOR_BACKGROUNDS = {
        "기본 책상": "clean modern desk setup",
        "빛이 비치는 거실": "bright living room with natural light"
    }

if not SPECIAL_BACKGROUNDS:
    SPECIAL_BACKGROUNDS = {
        "네이비 빈티지 플로럴 벽지": "navy vintage floral wallpaper background"
    }

if not JEWELRY_BACKGROUNDS:
    JEWELRY_BACKGROUNDS = {
        "화이트 미러 스팟 라이트": "white mirror surface with spotlight"
    }

if not SPECIAL_EFFECTS_BACKGROUNDS:
    SPECIAL_EFFECTS_BACKGROUNDS = {
        "블루블랙 큰 물방울 효과": "blue-black background with water droplet effects"
    }

# 임시 파일 저장 함수
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