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._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 클라이언트가 성공적으로 초기화되었습니다.")
                    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 _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("클라이언트 연결 실패")
                return tuple([gr.update() for _ in range(7)])
            
            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)}")
            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:
        """이미지 처리 (편집 및 화질 개선)"""
        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 호출
            logger.info("이미지 처리 API 호출 시작")
            result = self.client.predict(
                password,
                handle_file(temp_path),
                bg_type,
                simple or "",
                studio or "",
                nature or "",
                indoor or "",
                special or "",
                jewelry or "",
                special_effects or "",
                request_text or "",
                quality_level,
                aspect_ratio,
                output_format,
                enable_enhancement,
                api_name="/check_password"
            )
            
            logger.info("이미지 처리 API 호출 완료")
            
            # 결과 검증 및 반환
            if isinstance(result, (list, tuple)) and len(result) >= 7:
                return result
            else:
                logger.warning("API에서 예상과 다른 형식의 결과를 반환했습니다.")
                return ([], None, [], None, "", "", "API에서 올바르지 않은 응답을 받았습니다.")
            
        except Exception as e:
            logger.error(f"이미지 처리 오류: {str(e)}")
            return ([], None, [], None, "", "", f"이미지 처리 중 오류가 발생했습니다: {str(e)}")
        
        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="심플 배경"
                            )
                            
                            # 드롭다운 컴포넌트들 (초기에는 빈 상태)
                            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 파일을 확인해주세요.")