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 openai import OpenAI
import replicate
from google import genai
from google.genai import types

# 로깅 설정
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 클라이언트 초기화 (안전하게)
openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", ""))

# Gemini 클라이언트 초기화 - API 키가 있을 때만
gemini_api_key = os.environ.get("GEMINI_API_KEY")
if gemini_api_key and gemini_api_key.strip():
    try:
        gemini_client = genai.Client(api_key=gemini_api_key)
        logger.info("Gemini client initialized successfully")
    except Exception as e:
        logger.error(f"Failed to initialize Gemini client: {e}")
        gemini_client = None
else:
    logger.warning("GEMINI_API_KEY not found or empty, Gemini client not initialized")
    gemini_client = None

# 환경변수에서 배경 프롬프트 로드
BACKGROUNDS_DATA = os.environ.get("BACKGROUNDS_DATA", "")
# 환경변수에서 비밀번호 로드
APP_PASSWORD = os.environ.get("APP_PASSWORD", "")

if not BACKGROUNDS_DATA:
    logger.error("BACKGROUNDS_DATA environment variable not found")
    raise ValueError("BACKGROUNDS_DATA 환경변수가 설정되지 않았습니다.")

if not APP_PASSWORD:
    logger.error("APP_PASSWORD environment variable not found")
    raise ValueError("APP_PASSWORD 환경변수가 설정되지 않았습니다.")

# 환경변수 데이터를 파이썬 코드로 실행하여 배경 딕셔너리들 생성
try:
    exec(BACKGROUNDS_DATA)
    logger.info("Background data loaded from environment variable")
except Exception as e:
    logger.error(f"Failed to load background data: {e}")
    raise ValueError(f"배경 데이터 로드 실패: {e}")

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

# 텍스트 번역 함수 (한국어 → 영어) - Gemini 2.0 Flash 전용
def translate_to_english(text):
    """한국어 텍스트를 영어로 번역 (Gemini 2.0 Flash 사용)"""
    try:
        if not text or not text.strip():
            return ""
        
        # Gemini 클라이언트가 초기화되었는지 확인
        if gemini_client is None:
            logger.warning("Gemini client not available, returning original text")
            return text
        
        # Gemini 2.0 Flash를 사용한 번역
        try:
            response = gemini_client.models.generate_content(
                model="gemini-2.0-flash",
                config=types.GenerateContentConfig(
                    system_instruction="You are a professional translator. Translate the given Korean text to English. Keep the translation natural and contextually appropriate for image generation prompts. If the text is already in English, return it as is. Only return the translated text without any additional explanation.",
                    max_output_tokens=200,
                    temperature=0.1
                ),
                contents=[f"Translate this to English: {text}"]
            )
            
            translated = response.text.strip()
            logger.info(f"Translated '{text}' to '{translated}' using Gemini 2.0 Flash")
            return translated
            
        except Exception as e:
            logger.error(f"Gemini translation error: {e}")
            logger.warning("Translation failed, returning original text")
            return text
            
    except Exception as e:
        logger.error(f"Translation error: {e}")
        return text

# 프롬프트 생성 함수 (종횡비와 요청사항 통합)
def generate_prompt(background_type, background_name, user_request, aspect_ratio="1:1"):
    # 기본 고정 프롬프트 (종횡비 정보 포함) - 영어로 변경
    fixed_prompt = f"""
## Fixed Prompt (Required)
[Aspect Ratio: {aspect_ratio}]
[Foreground: all uploaded product images, preserve their original proportions and clarity]
[Preserve originals: keep the same random seed; maintain exact shape and aspect ratio; no vertical or horizontal scaling; do not alter any existing logos or text]
[Product sizing: ensure product images maintain at least 50% of their height relative to the background]
[Composition: products must be naturally composited with the background, maintain proper shadows aligned with lighting]
[Product placement: if products already exist in the background prompt, follow their exact arrangement and positioning]
"""
    
    # 배경 프롬프트 선택
    if background_type == "심플 배경":
        background_prompt = SIMPLE_BACKGROUNDS.get(background_name, "")
    elif background_type == "스튜디오 배경":
        background_prompt = STUDIO_BACKGROUNDS.get(background_name, "")
    elif background_type == "자연 환경":
        background_prompt = NATURE_BACKGROUNDS.get(background_name, "")
    elif background_type == "실내 환경":
        background_prompt = INDOOR_BACKGROUNDS.get(background_name, "")
    elif background_type == "특수배경":
        background_prompt = SPECIAL_BACKGROUNDS.get(background_name, "")
    elif background_type == "주얼리":
        background_prompt = JEWELRY_BACKGROUNDS.get(background_name, "")
    elif background_type == "특수효과":
        background_prompt = SPECIAL_EFFECTS_BACKGROUNDS.get(background_name, "")
    else:
        background_prompt = "clean white background with soft even lighting"
    
    # 사용자 요청사항 처리
    if user_request and user_request.strip():
        # 한국어 요청사항을 영어로 번역 (Gemini 2.0 Flash 사용)
        translated_request = translate_to_english(user_request)
        
        # 번역된 요청사항을 배경 프롬프트에 통합
        integrated_background = f"{background_prompt} Additionally, incorporate the following elements naturally into the scene: {translated_request}. Ensure these elements blend harmoniously with the existing background while maintaining the overall aesthetic and lighting."
        
        # 요청 프롬프트 섹션 (번역된 내용 사용)
        request_prompt = f"""
## Request Prompt
{translated_request}
"""
        
        # 배경 프롬프트 섹션
        background_section = f"""
## Background Prompt (Background Settings)
{integrated_background}
"""
    else:
        # 요청사항이 없는 경우
        request_prompt = f"""
## Request Prompt
No specific request
"""
        
        # 요청사항이 없는 경우 기본 배경만 사용
        background_section = f"""
## Background Prompt (Background Settings)
{background_prompt}
"""
    
    # 최종 프롬프트 조합
    final_prompt = fixed_prompt + request_prompt + background_section
    
    return final_prompt

# 이미지 편집 및 화질 개선 함수
def edit_and_enhance_image(
    prompt, 
    image,
    quality_level="gpt",
    aspect_ratio="1:1",
    output_format="jpg",
    enable_enhancement=True,
    enhancement_level=2
):
    try:
        logger.info(f"Editing image with prompt: '{prompt[:50]}...' (truncated)")
        logger.info(f"Parameters: quality_level={quality_level}, aspect_ratio={aspect_ratio}, output_format={output_format}")
        logger.info(f"Enhancement requested: {enable_enhancement}, level: {enhancement_level}")
        
        if image is None:
            logger.error("No image provided")
            return None, None, None, "이미지를 업로드해야 합니다."
        
        # 이미지 처리
        processed_image = None
        temp_paths = []  # 나중에 정리할 경로 추적
        
        img_path = save_uploaded_file(image)
        if img_path:
            logger.info(f"Saved image to temp path: {img_path}")
            processed_image = open(img_path, "rb")
            temp_paths.append(img_path)
        else:
            logger.error("Failed to save image")
            return None, None, None, "이미지 처리에 실패했습니다. 다른 이미지를 업로드해 보세요."
        
        # 모델 선택에 따른 처리
        edited_images = []
        usage_info = ""
        error_msg = None
        
        try:
            if quality_level == "gpt":
                # GPT 모델 사용
                if not openai_client.api_key:
                    logger.error("OpenAI API key is not set")
                    return None, None, None, "OpenAI API 키가 설정되지 않았습니다. API 키를 설정해주세요."
                
                # 종횡비를 크기로 변환
                size_mapping = {
                    "1:1": "1024x1024",
                    "3:2": "1536x1024", 
                    "2:3": "1024x1536"
                }
                size = size_mapping.get(aspect_ratio, "1024x1024")
                
                params = {
                    "prompt": prompt,
                    "model": "gpt-image-1",
                    "n": 1,
                    "size": size,
                    "image": processed_image
                }
                
                logger.info(f"Calling OpenAI API for image editing")
                response = openai_client.images.edit(**params)
                logger.info("OpenAI API call successful")
                
                # 결과 처리
                for i, data in enumerate(response.data):
                    logger.info(f"Processing result image {i+1}/{len(response.data)}")
                    
                    if hasattr(data, 'b64_json') and data.b64_json:
                        image_data = base64.b64decode(data.b64_json)
                        image = Image.open(io.BytesIO(image_data))
                    elif hasattr(data, 'url') and data.url:
                        response_url = requests.get(data.url)
                        image = Image.open(io.BytesIO(response_url.content))
                    else:
                        logger.warning(f"No image data found in response item {i+1}")
                        continue
                    
                    # 이미지 형식 변환
                    if output_format.lower() != "png" and image.mode == "RGBA":
                        background = Image.new("RGB", image.size, (255, 255, 255))
                        background.paste(image, mask=image.split()[3])
                        image = background
                    
                    edited_images.append(image)
                
                usage_info = "이미지 편집 완료 (GPT 모델 사용)"
                
            else:  # quality_level == "flux"
                # Flux 모델 사용 (항상 기본 화질개선 1회 적용)
                if not os.environ.get("REPLICATE_API_TOKEN"):
                    logger.error("Replicate API token is not set")
                    return None, None, None, "Replicate API 토큰이 설정되지 않았습니다. API 토큰을 설정해주세요."
                
                logger.info(f"Using Flux model for image editing")
                
                # Flux 모델로 이미지 생성
                output = replicate.run(
                    "black-forest-labs/flux-kontext-pro",
                    input={
                        "prompt": prompt,
                        "input_image": processed_image,
                        "output_format": output_format.lower(),
                        "aspect_ratio": aspect_ratio,
                        "safety_tolerance": 2
                    }
                )
                
                logger.info(f"Flux API response received")
                
                # Flux API 응답 처리
                flux_image = None
                if output:
                    # output이 바이트 스트림인 경우
                    if hasattr(output, 'read'):
                        image_data = output.read()
                        flux_image = Image.open(io.BytesIO(image_data))
                    # output이 URL인 경우  
                    elif isinstance(output, str) and output.startswith('http'):
                        response_url = requests.get(output)
                        flux_image = Image.open(io.BytesIO(response_url.content))
                    # output이 바이너리 데이터인 경우
                    else:
                        flux_image = Image.open(io.BytesIO(output))
                    
                    # 이미지 형식 변환
                    if output_format.lower() != "png" and flux_image.mode == "RGBA":
                        background = Image.new("RGB", flux_image.size, (255, 255, 255))
                        background.paste(flux_image, mask=flux_image.split()[3])
                        flux_image = background
                    
                    # Flux 모델은 항상 첫 번째 화질 개선을 자동 적용
                    try:
                        logger.info("Applying automatic first enhancement for Flux model")
                        
                        # 임시 파일로 저장
                        temp_flux_path = tempfile.mktemp(suffix='.png')
                        flux_image.save(temp_flux_path)
                        temp_paths.append(temp_flux_path)
                        
                        # 첫 번째 화질 향상 (Flux 모델 기본 적용)
                        first_enhanced_output = replicate.run(
                            "philz1337x/clarity-upscaler:dfad41707589d68ecdccd1dfa600d55a208f9310748e44bfe35b4a6291453d5e",
                            input={
                                "image": open(temp_flux_path, "rb"),
                                "scale_factor": 2,
                                "resemblance": 0.8,
                                "creativity": 0.2,
                                "output_format": output_format.lower(),
                                "prompt": prompt,
                                "negative_prompt": "(worst quality, low quality, normal quality:2)"
                            }
                        )
                        
                        if first_enhanced_output and isinstance(first_enhanced_output, list) and len(first_enhanced_output) > 0:
                            first_enhanced_url = first_enhanced_output[0]
                            first_enhanced_response = requests.get(first_enhanced_url)
                            
                            if first_enhanced_response.status_code == 200:
                                first_enhanced_image = Image.open(io.BytesIO(first_enhanced_response.content))
                                
                                # 이미지 형식 변환
                                if output_format.lower() != "png" and first_enhanced_image.mode == "RGBA":
                                    background = Image.new("RGB", first_enhanced_image.size, (255, 255, 255))
                                    background.paste(first_enhanced_image, mask=first_enhanced_image.split()[3])
                                    first_enhanced_image = background
                                
                                edited_images.append(first_enhanced_image)
                                usage_info = "이미지 편집 완료 (Flux 모델 + 기본 화질개선 적용)"
                                logger.info("First enhancement completed for Flux model")
                            else:
                                # 첫 번째 화질개선 실패 시 원본 사용
                                edited_images.append(flux_image)
                                usage_info = "이미지 편집 완료 (Flux 모델 사용, 기본 화질개선 실패)"
                        else:
                            # 첫 번째 화질개선 실패 시 원본 사용
                            edited_images.append(flux_image)
                            usage_info = "이미지 편집 완료 (Flux 모델 사용, 기본 화질개선 실패)"
                            
                    except Exception as e:
                        logger.error(f"Error in first enhancement for Flux: {e}")
                        # 첫 번째 화질개선 실패 시 원본 사용
                        edited_images.append(flux_image)
                        usage_info = f"이미지 편집 완료 (Flux 모델 사용, 기본 화질개선 오류: {str(e)})"
                        
                else:
                    logger.error("No output from Flux API")
                    error_msg = "Flux API에서 응답을 받지 못했습니다."
        
        except Exception as e:
            if quality_level == "gpt":
                logger.error(f"OpenAI API call error: {e}")
                error_msg = f"OpenAI API 호출 오류: {str(e)}"
            else:
                logger.error(f"Flux API call error: {e}")
                error_msg = f"Flux API 호출 오류: {str(e)}"
        
        finally:
            # 임시 파일 정리
            if processed_image and hasattr(processed_image, 'close'):
                processed_image.close()
        
        # 화질 향상 처리 (GPT 모델은 일반적인 화질개선, Flux 모델은 2차 화질개선)
        enhanced_image = None
        temp_image_path = None
        
        if enable_enhancement and edited_images and not error_msg:
            try:
                if quality_level == "gpt":
                    # GPT 모델: 일반적인 화질 개선
                    logger.info(f"Enhancing GPT image with Replicate API, enhancement level: {enhancement_level}")
                    enhancement_info = "화질 개선"
                else:
                    # Flux 모델: 2차 화질 개선 (이미 1차는 적용됨)
                    logger.info(f"Applying second enhancement for Flux image, enhancement level: {enhancement_level}")
                    enhancement_info = "2차 화질 개선"
                
                if not os.environ.get("REPLICATE_API_TOKEN"):
                    logger.error("Replicate API token is not set")
                    usage_info += f" | {enhancement_info} 실패: Replicate API 토큰이 설정되지 않았습니다."
                else:
                    # 임시 파일로 저장
                    temp_image_path = tempfile.mktemp(suffix='.png')
                    edited_images[0].save(temp_image_path)
                    temp_paths.append(temp_image_path)
                    
                    # Replicate API로 화질 향상
                    output = replicate.run(
                        "philz1337x/clarity-upscaler:dfad41707589d68ecdccd1dfa600d55a208f9310748e44bfe35b4a6291453d5e",
                        input={
                            "image": open(temp_image_path, "rb"),
                            "scale_factor": enhancement_level,
                            "resemblance": 0.8,
                            "creativity": 0.2,
                            "output_format": output_format.lower(),
                            "prompt": prompt,
                            "negative_prompt": "(worst quality, low quality, normal quality:2)"
                        }
                    )
                    
                    logger.info(f"Replicate API response: {output}")
                    
                    if output and isinstance(output, list) and len(output) > 0:
                        enhanced_url = output[0]
                        enhanced_response = requests.get(enhanced_url)
                        if enhanced_response.status_code == 200:
                            enhanced_image = Image.open(io.BytesIO(enhanced_response.content))
                            
                            if output_format.lower() != "png" and enhanced_image.mode == "RGBA":
                                background = Image.new("RGB", enhanced_image.size, (255, 255, 255))
                                background.paste(enhanced_image, mask=enhanced_image.split()[3])
                                enhanced_image = background
                            
                            if quality_level == "gpt":
                                usage_info += f" | {enhancement_info} 완료: Replicate Clarity Upscaler 사용"
                            else:
                                usage_info += f" | {enhancement_info} 완료: 총 2회 화질개선 적용"
                        else:
                            usage_info += f" | {enhancement_info} 실패: 이미지 다운로드 오류"
                    else:
                        usage_info += f" | {enhancement_info} 실패: Replicate API 응답 없음"
                
            except Exception as e:
                logger.error(f"Error enhancing image: {e}")
                if quality_level == "gpt":
                    usage_info += f" | 화질 개선 실패: {str(e)}"
                else:
                    usage_info += f" | 2차 화질 개선 실패: {str(e)}"
        
        # 임시 파일 정리
        for path in temp_paths:
            if os.path.exists(path):
                try:
                    os.remove(path)
                    logger.info(f"Removed temp file: {path}")
                except Exception as e:
                    logger.error(f"Error removing temp file {path}: {e}")
        
        # 결과 반환
        if error_msg:
            return None, None, None, error_msg
        elif edited_images:
            if enable_enhancement and enhanced_image:
                return edited_images, [enhanced_image], usage_info, None
            else:
                return edited_images, None, usage_info, None
        else:
            return None, None, None, "이미지 편집에 실패했습니다."
            
    except Exception as e:
        logger.error(f"Error in edit_and_enhance_image function: {e}")
        logger.error(traceback.format_exc())
        return None, None, None, f"에러 발생: {str(e)}\n\n{traceback.format_exc()}"

# 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
                            )
                            
                            # 드롭다운 변경 함수
                            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)
            
            # 프롬프트만 생성하는 함수 (비밀번호 체크 포함)
            def generate_prompt_with_password_check(password, bg_type, simple, studio, nature, indoor, special, jewelry, special_effects, request_text, aspect_ratio):
                # 비밀번호 확인
                if password != APP_PASSWORD:
                    return "비밀번호가 틀렸습니다. 올바른 비밀번호를 입력해주세요."
                
                # 배경 선택
                background_name = ""
                if bg_type == "심플 배경":
                    background_name = simple
                elif bg_type == "스튜디오 배경":
                    background_name = studio
                elif bg_type == "자연 환경":
                    background_name = nature
                elif bg_type == "실내 환경":
                    background_name = indoor
                elif bg_type == "특수배경":
                    background_name = special
                elif bg_type == "주얼리":
                    background_name = jewelry
                elif bg_type == "특수효과":
                    background_name = special_effects
                
                # 프롬프트 생성 (종횡비 포함)
                prompt = generate_prompt(bg_type, background_name, request_text, aspect_ratio)
                return prompt
            
            # 비밀번호 확인 함수
            def check_password(password, *args):
                if password != APP_PASSWORD:
                    return (
                        [],  # original_output
                        None,  # original_download
                        [],  # enhanced_output
                        None,  # enhanced_download
                        "",  # prompt_output
                        "",  # info
                        "비밀번호가 틀렸습니다. 올바른 비밀번호를 입력해주세요."  # error
                    )
                
                # 이미지 편집 요청 처리
                image, bg_type, simple, studio, nature, indoor, special, jewelry, special_effects, request_text, quality_level, aspect_ratio, output_format, enable_enhancement = args
                
                # 배경 선택
                background_name = ""
                if bg_type == "심플 배경":
                    background_name = simple
                elif bg_type == "스튜디오 배경":
                    background_name = studio
                elif bg_type == "자연 환경":
                    background_name = nature
                elif bg_type == "실내 환경":
                    background_name = indoor
                elif bg_type == "특수배경":
                    background_name = special
                elif bg_type == "주얼리":
                    background_name = jewelry
                elif bg_type == "특수효과":
                    background_name = special_effects
                
                # 프롬프트 생성
                prompt = generate_prompt(bg_type, background_name, request_text, aspect_ratio)
                
                # 이미지 편집 및 화질 개선 실행
                original_images, enhanced_images, usage_info, error_msg = edit_and_enhance_image(
                    prompt, image, quality_level, aspect_ratio, output_format, enable_enhancement, 2
                )
                
                # 이미지 저장 및 다운로드 파일 준비
                original_path = None
                enhanced_path = None
                
                if error_msg:
                    logger.error(f"Error returned from edit_and_enhance_image: {error_msg}")
                    return (
                        [],  # original_output
                        None,  # original_download
                        [],  # enhanced_output
                        None,  # enhanced_download
                        prompt,  # prompt_output
                        "",  # info
                        error_msg  # error
                    )
                else:
                    # 원본 편집 이미지 저장
                    if original_images and len(original_images) > 0:
                        try:
                            original_path = f"original_image.{output_format}"
                            original_images[0].save(original_path)
                            logger.info(f"Saved original image to {original_path}")
                        except Exception as e:
                            logger.error(f"Error saving original image: {e}")
                    
                    # 화질 개선 이미지 저장
                    if enhanced_images and len(enhanced_images) > 0:
                        try:
                            enhanced_path = f"enhanced_image.{output_format}"
                            enhanced_images[0].save(enhanced_path)
                            logger.info(f"Saved enhanced image to {enhanced_path}")
                        except Exception as e:
                            logger.error(f"Error saving enhanced image: {e}")
                    
                    # 결과 반환
                    return (
                        original_images if original_images else [],  # original_output
                        original_path,  # original_download
                        enhanced_images if enhanced_images else [],  # enhanced_output
                        enhanced_path,  # enhanced_download
                        prompt,  # prompt_output
                        usage_info,  # info
                        ""  # error (빈 문자열로 설정)
                    )
            
            # 프롬프트 생성 버튼 클릭 이벤트
            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]
            )
            
            # 편집 버튼 클릭 이벤트
            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
                ]
            )
        
        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())