app.py

import os
import sys
import logging
import tempfile
import traceback
from typing import List, Tuple, Optional, Any
import gradio as gr
from gradio_client import Client, handle_file
from dotenv import load_dotenv
import time

# .env 파일 로드
load_dotenv()

# 로깅 설정 (API 엔드포인트 정보는 로그에 남기지 않음)
class SafeFormatter(logging.Formatter):
    """API 엔드포인트 정보를 로그에서 제거하는 안전한 포매터"""
    def format(self, record):
        # API_ENDPOINT 관련 정보 필터링
        msg = super().format(record)
        # 민감한 정보가 포함된 경우 마스킹
        if "API_ENDPOINT" in msg or "happydoggg" in msg or "49493h" in msg:
            return msg.replace(os.environ.get("API_ENDPOINT", ""), "[API_ENDPOINT_HIDDEN]")
        return msg

# 로그 핸들러 설정
file_handler = logging.FileHandler("app.log")
console_handler = logging.StreamHandler(sys.stdout)
formatter = SafeFormatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler.setFormatter(formatter)
console_handler.setFormatter(formatter)

logger = logging.getLogger("image-enhancer-control-tower")
logger.setLevel(logging.INFO)
logger.addHandler(file_handler)
logger.addHandler(console_handler)

class GradioClientController:
    """허깅페이스 그라디오 API 클라이언트 컨트롤러"""
    
    def __init__(self):
        self.client = None
        self.api_endpoint = None
        self.background_options = {}  # 캐시된 배경 옵션
        self._initialize_client()
        
    def _initialize_client(self):
        """클라이언트 초기화 (엔드포인트 정보는 로그에 남기지 않음)"""
        try:
            self.api_endpoint = os.environ.get("API_ENDPOINT")
            if not self.api_endpoint:
                logger.error("API_ENDPOINT 환경 변수가 설정되지 않았습니다.")
                raise ValueError("API_ENDPOINT 환경 변수가 필요합니다.")
            
            # 클라이언트 초기화 시도 (재시도 로직 포함)
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    self.client = Client(self.api_endpoint)
                    logger.info("API 클라이언트가 성공적으로 초기화되었습니다.")
                    # 초기화 성공 시 배경 옵션 로드
                    self._load_background_options()
                    return
                except Exception as e:
                    if attempt < max_retries - 1:
                        logger.warning(f"클라이언트 초기화 재시도 {attempt + 1}/{max_retries}")
                        time.sleep(2)
                    else:
                        raise e
                        
        except Exception as e:
            logger.error(f"클라이언트 초기화 실패: {str(e)}")
            self.client = None
    
    def _load_background_options(self):
        """배경 옵션을 미리 로드하여 캐시"""
        try:
            if self.client:
                # 각 배경 타입별로 옵션 로드
                background_types = ["심플 배경", "스튜디오 배경", "자연 환경", "실내 환경", "특수배경", "주얼리", "특수효과"]
                for bg_type in background_types:
                    result = self.client.predict(bg_type, api_name="/update_dropdowns")
                    if isinstance(result, (list, tuple)) and len(result) >= 7:
                        self.background_options[bg_type] = result
                        logger.info(f"배경 옵션 로드 완료: {bg_type}")
                    time.sleep(0.1)  # API 호출 간격
        except Exception as e:
            logger.warning(f"배경 옵션 미리 로드 실패: {str(e)}")
            # 실패해도 계속 진행
    
    def _ensure_client(self) -> bool:
        """클라이언트 연결 상태 확인 및 재연결"""
        if self.client is None:
            logger.info("클라이언트 재연결 시도")
            self._initialize_client()
        return self.client is not None
    
    def update_dropdowns(self, bg_type: str) -> Tuple:
        """배경 유형에 따른 드롭다운 업데이트"""
        try:
            if not self._ensure_client():
                logger.error("클라이언트 연결 실패")
                # 캐시된 데이터가 있으면 사용
                if bg_type in self.background_options:
                    result = self.background_options[bg_type]
                    logger.info(f"캐시된 드롭다운 데이터 사용: {bg_type}")
                else:
                    return tuple([gr.update() for _ in range(7)])
            else:
                # 실시간으로 API 호출
                result = self.client.predict(bg_type, api_name="/update_dropdowns")
                logger.info(f"실시간 드롭다운 업데이트 완료: {bg_type}")
            
            # 결과를 Gradio 업데이트 형식으로 변환
            if isinstance(result, (list, tuple)) and len(result) >= 7:
                updates = []
                visibility_map = ["심플 배경", "스튜디오 배경", "자연 환경", "실내 환경", "특수배경", "주얼리", "특수효과"]
                
                for i, choices in enumerate(result[:7]):
                    is_visible = (bg_type == visibility_map[i])
                    updates.append(gr.update(
                        visible=is_visible, 
                        choices=choices if choices else [], 
                        value=choices[0] if choices and len(choices) > 0 else None
                    ))
                
                return tuple(updates)
            else:
                logger.warning("API에서 예상과 다른 형식의 결과를 반환했습니다.")
                return tuple([gr.update() for _ in range(7)])
            
        except Exception as e:
            logger.error(f"드롭다운 업데이트 오류: {str(e)}")
            # 캐시된 데이터 시도
            if bg_type in self.background_options:
                try:
                    result = self.background_options[bg_type]
                    updates = []
                    visibility_map = ["심플 배경", "스튜디오 배경", "자연 환경", "실내 환경", "특수배경", "주얼리", "특수효과"]
                    
                    for i, choices in enumerate(result[:7]):
                        is_visible = (bg_type == visibility_map[i])
                        updates.append(gr.update(
                            visible=is_visible, 
                            choices=choices if choices else [], 
                            value=choices[0] if choices and len(choices) > 0 else None
                        ))
                    
                    logger.info(f"캐시된 데이터로 드롭다운 업데이트: {bg_type}")
                    return tuple(updates)
                except:
                    pass
                    
            return tuple([gr.update() for _ in range(7)])
    
    def generate_prompt_only(self, password: str, bg_type: str, simple: str, studio: str, 
                           nature: str, indoor: str, special: str, jewelry: str, 
                           special_effects: str, request_text: str, aspect_ratio: str) -> str:
        """프롬프트만 생성"""
        try:
            if not self._ensure_client():
                return "클라이언트 연결 오류가 발생했습니다. 잠시 후 다시 시도해주세요."
            
            result = self.client.predict(
                password,
                bg_type,
                simple or "",
                studio or "",
                nature or "",
                indoor or "",
                special or "",
                jewelry or "",
                special_effects or "",
                request_text or "",
                aspect_ratio,
                api_name="/generate_prompt_with_password_check"
            )
            
            logger.info("프롬프트 생성 완료")
            return result if result else "프롬프트 생성에 실패했습니다."
            
        except Exception as e:
            logger.error(f"프롬프트 생성 오류: {str(e)}")
            return f"프롬프트 생성 오류: {str(e)}"
    
    def process_image(self, password: str, image, bg_type: str, simple: str, studio: str,
                     nature: str, indoor: str, special: str, jewelry: str, special_effects: str,
                     request_text: str, quality_level: str, aspect_ratio: str, 
                     output_format: str, enable_enhancement: bool) -> Tuple:
        """이미지 처리 (편집 및 화질 개선) - API 문서와 정확히 매칭"""
        temp_path = None
        try:
            if not self._ensure_client():
                return ([], None, [], None, "", "", "클라이언트 연결 오류가 발생했습니다. 잠시 후 다시 시도해주세요.")
            
            # 이미지 검증
            if image is None:
                return ([], None, [], None, "", "", "이미지를 업로드해주세요.")
            
            # 이미지를 임시 파일로 저장
            temp_path = tempfile.mktemp(suffix='.png')
            try:
                image.save(temp_path, format='PNG')
                logger.info(f"이미지를 임시 파일로 저장: {os.path.basename(temp_path)}")
            except Exception as e:
                logger.error(f"이미지 저장 실패: {str(e)}")
                return ([], None, [], None, "", "", "이미지 저장에 실패했습니다. 다른 이미지를 시도해보세요.")
            
            # API 호출 - 매개변수 순서를 API 문서와 정확히 맞춤
            logger.info("이미지 처리 API 호출 시작")
            result = self.client.predict(
                password,                    # password (str)
                handle_file(temp_path),      # param_1 (image file)
                bg_type,                     # param_2 (배경 유형)
                simple or "",                # param_3 (심플 배경)
                studio or "",                # param_4 (스튜디오 배경)
                nature or "",                # param_5 (자연 환경)
                indoor or "",                # param_6 (실내 환경)
                special or "",               # param_7 (특수배경)
                jewelry or "",               # param_8 (주얼리)
                special_effects or "",       # param_9 (특수효과)
                request_text or "",          # param_10 (요청사항)
                quality_level,               # param_11 (품질 레벨)
                aspect_ratio,                # param_12 (종횡비)
                output_format,               # param_13 (이미지 형식)
                enable_enhancement,          # param_14 (추가 화질 개선)
                api_name="/check_password"
            )
            
            logger.info("이미지 처리 API 호출 완료")
            
            # 결과 검증 및 반환 - API 문서에 따르면 tuple of 7 elements
            if isinstance(result, (list, tuple)) and len(result) >= 7:
                # [0] original_gallery, [1] original_file, [2] enhanced_gallery, 
                # [3] enhanced_file, [4] prompt_text, [5] info_text, [6] error_text
                return result[:7]  # 정확히 7개 요소만 반환
            else:
                logger.warning(f"API 응답 형식 이상: type={type(result)}, length={len(result) if hasattr(result, '__len__') else 'N/A'}")
                return ([], None, [], None, "", "", "API에서 올바르지 않은 응답을 받았습니다.")
            
        except Exception as e:
            error_msg = str(e)
            logger.error(f"이미지 처리 오류: {error_msg}")
            
            # 일반적인 오류 메시지 개선
            if "timeout" in error_msg.lower():
                user_error = "요청 처리 시간이 초과되었습니다. 잠시 후 다시 시도해주세요."
            elif "connection" in error_msg.lower():
                user_error = "서버 연결에 문제가 있습니다. 네트워크 상태를 확인해주세요."
            elif "unauthorized" in error_msg.lower() or "password" in error_msg.lower():
                user_error = "인증에 실패했습니다. 비밀번호를 확인해주세요."
            else:
                user_error = f"이미지 처리 중 오류가 발생했습니다: {error_msg}"
                
            return ([], None, [], None, "", "", user_error)
        
        finally:
            # 임시 파일 정리
            if temp_path and os.path.exists(temp_path):
                try:
                    os.remove(temp_path)
                    logger.info("임시 파일 정리 완료")
                except Exception as e:
                    logger.warning(f"임시 파일 정리 실패: {str(e)}")

# 전역 클라이언트 컨트롤러 인스턴스
controller = GradioClientController()

def create_gradio_interface():
    """Gradio 인터페이스 생성"""
    try:
        logger.info("Gradio 인터페이스 생성 시작")
        
        # 배경 유형 선택지
        background_types = ["심플 배경", "스튜디오 배경", "자연 환경", "실내 환경", "특수배경", "주얼리", "특수효과"]
        
        # 커스텀 CSS 추가 (원본과 동일한 스타일 유지)
        css = """
        .gradio-container {
            max-width: 1200px !important;
        }
        """
        
        with gr.Blocks(title="AI 이미지 편집 및 화질 개선", css=css, theme=gr.themes.Soft()) 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=background_types,
                                label="배경 유형",
                                value="심플 배경"
                            )
                            
                            # 드롭다운 컴포넌트들 (API에서 동적으로 로드)
                            simple_dropdown = gr.Dropdown(
                                choices=[],
                                label="심플 배경 선택",
                                visible=True,
                                interactive=True
                            )
                            
                            studio_dropdown = gr.Dropdown(
                                choices=[],
                                label="스튜디오 배경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            nature_dropdown = gr.Dropdown(
                                choices=[],
                                label="자연 환경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            indoor_dropdown = gr.Dropdown(
                                choices=[],
                                label="실내 환경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            special_dropdown = gr.Dropdown(
                                choices=[],
                                label="특수배경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            jewelry_dropdown = gr.Dropdown(
                                choices=[],
                                label="주얼리 배경 선택",
                                visible=False,
                                interactive=True
                            )
                            
                            special_effects_dropdown = gr.Dropdown(
                                choices=[],
                                label="특수효과 배경 선택",
                                visible=False,
                                interactive=True
                            )
                    
                    # 요청사항 입력
                    request_text = gr.Textbox(
                        label="요청사항", 
                        placeholder="상품 이미지에 적용할 스타일, 분위기, 특별 요청사항 등을 입력하세요.", 
                        lines=3
                    )
                    
                    # 새로운 옵션들
                    quality_level = gr.Radio(
                        label="품질 레벨", 
                        choices=["gpt", "flux"],
                        value="flux",
                        info="GPT: GPT 모델 (고품질), Flux: 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회)"
                    )
                    
                    # 프롬프트 생성 버튼
                    generate_prompt_btn = gr.Button("프롬프트만 생성", variant="secondary")
                    
                    # 편집 버튼
                    edit_btn = gr.Button("이미지 편집 및 화질 개선", variant="primary")
                
                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)
            
            # 드롭다운 변경 이벤트
            background_type.change(
                fn=controller.update_dropdowns,
                inputs=[background_type],
                outputs=[simple_dropdown, studio_dropdown, nature_dropdown, 
                        indoor_dropdown, special_dropdown, jewelry_dropdown, special_effects_dropdown]
            )
            
            # 프롬프트 생성 버튼 클릭 이벤트
            generate_prompt_btn.click(
                fn=controller.generate_prompt_only,
                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=controller.process_image,
                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
                ]
            )
            
            # 앱 로드 시 초기 드롭다운 설정
            app.load(
                fn=controller.update_dropdowns,
                inputs=[background_type],
                outputs=[simple_dropdown, studio_dropdown, nature_dropdown, 
                        indoor_dropdown, special_dropdown, jewelry_dropdown, special_effects_dropdown]
            )
        
        logger.info("Gradio 인터페이스 생성 완료")
        return app
        
    except Exception as e:
        logger.error(f"Gradio 인터페이스 생성 오류: {str(e)}")
        logger.error(traceback.format_exc())
        raise

# 앱 실행
if __name__ == "__main__":
    try:
        logger.info("애플리케이션 시작")
        
        # imgs 디렉토리 확인/생성
        os.makedirs("imgs", exist_ok=True)
        logger.info("이미지 디렉토리 준비 완료")
        
        # API 엔드포인트 확인
        if not os.environ.get("API_ENDPOINT"):
            logger.error("API_ENDPOINT 환경 변수가 설정되지 않았습니다.")
            print("\n❌ 오류: API_ENDPOINT 환경 변수가 설정되지 않았습니다.")
            print("환경 변수를 설정하거나 .env 파일에 API_ENDPOINT를 추가해주세요.")
            sys.exit(1)
        
        app = create_gradio_interface()
        logger.info("Gradio 앱 시작")
        app.launch(
            share=True,
            server_name="0.0.0.0",
            server_port=7860,
            show_error=True,
            quiet=False
        )
        
    except KeyboardInterrupt:
        logger.info("사용자에 의해 애플리케이션이 중단되었습니다.")
    except Exception as e:
        logger.error(f"앱 실행 오류: {str(e)}")
        logger.error(traceback.format_exc())
        print(f"\n❌ 심각한 오류가 발생했습니다: {str(e)}")
        print("자세한 내용은 app.log 파일을 확인해주세요.")