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
import replicate

# 로깅 설정
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("enhancer-app")

# 임시 파일 저장 함수
def save_uploaded_file(uploaded_file, suffix='.png'):
    try:
        if uploaded_file is None:
            return None
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
            temp_filename = temp_file.name
        
        if isinstance(uploaded_file, str):
            return uploaded_file
        
        if isinstance(uploaded_file, Image.Image):
            uploaded_file.save(temp_filename, format="PNG")
            return temp_filename
        
        with open(temp_filename, "wb") as f:
            if hasattr(uploaded_file, "read"):
                content = uploaded_file.read()
                f.write(content)
            else:
                f.write(uploaded_file)
        
        return temp_filename
    except Exception as e:
        logger.error(f"Error saving uploaded file: {e}")
        logger.error(traceback.format_exc())
        return None

# 이미지 화질 개선 함수
def enhance_image(
    image,
    output_format="jpg",
    enhancement_level=2
):
    try:
        if image is None:
            return None, None, "이미지를 업로드해주세요."
        
        # 이미지 처리
        temp_paths = []
        img_path = save_uploaded_file(image)
        if not img_path:
            return None, None, "이미지 처리에 실패했습니다."
        
        temp_paths.append(img_path)
        
        # Replicate API 토큰 확인
        if not os.environ.get("REPLICATE_API_TOKEN"):
            return None, None, "API 토큰이 설정되지 않았습니다."
        
        try:
            # Replicate API로 화질 향상
            output = replicate.run(
                "philz1337x/clarity-upscaler:dfad41707589d68ecdccd1dfa600d55a208f9310748e44bfe35b4a6291453d5e",
                input={
                    "image": open(img_path, "rb"),
                    "scale_factor": enhancement_level,
                    "resemblance": 0.8,
                    "creativity": 0.2,
                    "output_format": output_format.lower(),
                    "prompt": "",
                    "negative_prompt": "(worst quality, low quality, normal quality:2)"
                }
            )
            
            # Replicate API 응답 처리
            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
                    
                    # 결과 이미지 저장
                    enhanced_path = tempfile.mktemp(suffix=f'.{output_format}')
                    enhanced_image.save(enhanced_path)
                    temp_paths.append(enhanced_path)
                    
                    return enhanced_image, enhanced_path, ""
                else:
                    return None, None, "이미지 다운로드 오류"
            else:
                return None, None, "API 응답 없음"
                
        except Exception as e:
            logger.error(f"Error enhancing image: {e}")
            logger.error(traceback.format_exc())
            return None, None, f"{str(e)}"
        
        finally:
            # 임시 파일 정리 (결과 이미지는 유지)
            for path in temp_paths[:-1]:  # 마지막 파일(결과)은 남김
                if os.path.exists(path):
                    try:
                        os.remove(path)
                    except Exception as e:
                        logger.error(f"Error removing temp file {path}: {e}")
            
    except Exception as e:
        logger.error(f"Error in enhance_image function: {e}")
        logger.error(traceback.format_exc())
        return None, None, f"{str(e)}"

# Gradio 인터페이스 구성
def create_gradio_interface():
    try:
        with gr.Blocks() as app:
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type="pil")
                    
                    with gr.Row():
                        output_format = gr.Dropdown(
                            choices=["jpg", "png", "webp"],
                            value="jpg",
                            visible=False
                        )
                        enhancement_level = gr.Slider(
                            minimum=1, maximum=4, value=2, step=1
                        )
                    
                    process_btn = gr.Button("실행")
                
                with gr.Column():
                    output_image = gr.Image()
                    output_download = gr.File(interactive=False)
            
            def on_process(img, format, level):
                if img is None:
                    return None, None
                
                enhanced_img, download_path, error = enhance_image(
                    img, format, level
                )
                
                if error:
                    return None, None
                
                return enhanced_img, download_path
            
            process_btn.click(
                on_process,
                inputs=[input_image, output_format, enhancement_level],
                outputs=[output_image, output_download]
            )
        
        return app
    except Exception as e:
        logger.error(f"Error creating Gradio interface: {e}")
        logger.error(traceback.format_exc())
        raise

# 앱 실행
if __name__ == "__main__":
    try:
        app = create_gradio_interface()
        app.launch(share=True)
    except Exception as e:
        logger.error(f"Error running app: {e}")
        logger.error(traceback.format_exc())