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

# .env 파일 로드
load_dotenv()

# 로깅 설정 (API 엔드포인트 정보는 로그에 남기지 않음)
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")

class GradioClientController:
    """허깅페이스 그라디오 API 클라이언트 컨트롤러"""
    
    def __init__(self):
        self.client = None
        self._initialize_client()
        
    def _initialize_client(self):
        """클라이언트 초기화 (엔드포인트 정보는 로그에 남기지 않음)"""
        try:
            api_endpoint = os.environ.get("API_ENDPOINT")
            if not api_endpoint:
                logger.error("API_ENDPOINT 환경 변수가 설정되지 않았습니다.")
                raise ValueError("API_ENDPOINT 환경 변수가 필요합니다.")
            
            self.client = Client(api_endpoint)
            logger.info("API 클라이언트가 성공적으로 초기화되었습니다.")
            
        except Exception as e:
            logger.error(f"클라이언트 초기화 실패: {e}")
            self.client = None
    
    def update_dropdowns(self, bg_type: str) -> Tuple:
        """배경 유형에 따른 드롭다운 업데이트"""
        try:
            if not self.client:
                logger.error("클라이언트가 초기화되지 않았습니다.")
                return tuple([gr.update() for _ in range(7)])
            
            result = self.client.predict(
                bg_type=bg_type,
                api_name="/update_dropdowns"
            )
            
            logger.info(f"드롭다운 업데이트 완료: {bg_type}")
            
            # 결과를 Gradio 업데이트 형식으로 변환
            updates = []
            for i, choices in enumerate(result):
                if i == 0:  # 심플 배경
                    updates.append(gr.update(visible=(bg_type == "심플 배경"), choices=choices, value=choices[0] if choices else None))
                elif i == 1:  # 스튜디오 배경
                    updates.append(gr.update(visible=(bg_type == "스튜디오 배경"), choices=choices, value=choices[0] if choices else None))
                elif i == 2:  # 자연 환경
                    updates.append(gr.update(visible=(bg_type == "자연 환경"), choices=choices, value=choices[0] if choices else None))
                elif i == 3:  # 실내 환경
                    updates.append(gr.update(visible=(bg_type == "실내 환경"), choices=choices, value=choices[0] if choices else None))
                elif i == 4:  # 특수배경
                    updates.append(gr.update(visible=(bg_type == "특수배경"), choices=choices, value=choices[0] if choices else None))
                elif i == 5:  # 주얼리
                    updates.append(gr.update(visible=(bg_type == "주얼리"), choices=choices, value=choices[0] if choices else None))
                elif i == 6:  # 특수효과
                    updates.append(gr.update(visible=(bg_type == "특수효과"), choices=choices, value=choices[0] if choices else None))
            
            return tuple(updates)
            
        except Exception as e:
            logger.error(f"드롭다운 업데이트 오류: {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.client:
                logger.error("클라이언트가 초기화되지 않았습니다.")
                return "클라이언트 연결 오류"
            
            result = self.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"
            )
            
            logger.info("프롬프트 생성 완료")
            return result
            
        except Exception as e:
            logger.error(f"프롬프트 생성 오류: {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:
        """이미지 처리 (편집 및 화질 개선)"""
        try:
            if not self.client:
                logger.error("클라이언트가 초기화되지 않았습니다.")
                return ([], None, [], None, "", "", "클라이언트 연결 오류")
            
            # 이미지를 임시 파일로 저장
            temp_path = None
            if image is not None:
                temp_path = tempfile.mktemp(suffix='.png')
                image.save(temp_path)
            
            # API 호출
            result = self.client.predict(
                password=password,
                param_1=handle_file(temp_path) if temp_path else None,
                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"
            )
            
            # 임시 파일 정리
            if temp_path and os.path.exists(temp_path):
                try:
                    os.remove(temp_path)
                except:
                    pass
            
            logger.info("이미지 처리 완료")
            
            # 결과 반환: (original_output, original_download, enhanced_output, enhanced_download, prompt_output, info, error)
            return result
            
        except Exception as e:
            logger.error(f"이미지 처리 오류: {e}")
            # 임시 파일 정리
            if 'temp_path' in locals() and temp_path and os.path.exists(temp_path):
                try:
                    os.remove(temp_path)
                except:
                    pass
            return ([], None, [], None, "", "", f"이미지 처리 오류: {str(e)}")

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

def create_gradio_interface():
    """Gradio 인터페이스 생성"""
    try:
        logger.info("Gradio 인터페이스 생성 시작")
        
        # 배경 유형 선택지 (하드코딩)
        background_types = ["심플 배경", "스튜디오 배경", "자연 환경", "실내 환경", "특수배경", "주얼리", "특수효과"]
        
        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=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
                            )
                            
                            # 드롭다운 변경 이벤트
                            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]
                            )
                    
                    # 요청사항 입력
                    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회)"
                    )
                    
                    # 프롬프트 생성 버튼
                    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)
            
            # 프롬프트 생성 버튼 클릭 이벤트
            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 인터페이스 생성 오류: {e}")
        logger.error(traceback.format_exc())
        raise

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