app.py

import os
import logging
import replicate
import gradio as gr
from PIL import Image
import tempfile
import uuid

# 로깅 설정
logger = logging.getLogger(__name__)

# 환경변수에서 설정값들 가져오기
REPLICATE_API_TOKEN = os.environ.get("REPLICATE_API_TOKEN", "")
APP_PASSWORD = os.environ.get("APP_PASSWORD", "")
REPLICATE_MODEL = os.environ.get("REPLICATE_MODEL", "")

# 환경변수 검증
if not REPLICATE_API_TOKEN:
    logger.warning("REPLICATE_API_TOKEN이 설정되지 않았습니다. 환경 변수를 설정해주세요.")

if not APP_PASSWORD:
    logger.warning("APP_PASSWORD가 설정되지 않았습니다. 환경 변수를 설정해주세요.")

if not REPLICATE_MODEL:
    logger.warning("REPLICATE_MODEL이 설정되지 않았습니다. 환경 변수를 설정해주세요.")

# 임시 디렉토리 생성
TEMP_DIR = os.path.join(tempfile.gettempdir(), "image_processor")
os.makedirs(TEMP_DIR, exist_ok=True)
logger.info(f"임시 디렉토리 생성: {TEMP_DIR}")

def upscale_image(
    image, 
    scale_factor=2, 
    output_format="jpg",
    sd_model="juggernaut_reborn.safetensors [338b85bc4f]",
    resemblance=0.6,
    creativity=0.35,
    prompt="masterpiece, best quality, highres, <lora:more_details:0.5> <lora:SDXLrender_v2.0:1>",
    negative_prompt="(worst quality, low quality, normal quality:2) JuggernautNegative-neg",
    seed=1337,
    dynamic=6,
    sharpen=0
):
    """
    Clarity Upscaler를 사용하여 이미지 화질을 개선합니다.
    
    Args:
        image: 입력 이미지
        scale_factor: 확대 비율 (기본값: 2)
        output_format: 출력 형식 (기본값: jpg)
        sd_model: 사용할 SD 모델
        resemblance: 원본과의 유사도 (0.0-1.0)
        creativity: 창의성 수준 (0.0-1.0)
        prompt: 업스케일링 가이드 프롬프트
        negative_prompt: 네거티브 프롬프트
        seed: 랜덤 시드
        dynamic: 다이나믹 임계값 (1-10)
        sharpen: 선명도 (0-2)
    
    Returns:
        개선된 이미지
    """
    if not REPLICATE_API_TOKEN:
        raise ValueError("REPLICATE_API_TOKEN이 설정되지 않았습니다.")
    
    if not REPLICATE_MODEL:
        raise ValueError("REPLICATE_MODEL이 설정되지 않았습니다.")
    
    try:
        # 이미지 저장
        temp_input_path = os.path.join(TEMP_DIR, f"input_{uuid.uuid4()}.png")
        image.save(temp_input_path)
        logger.info(f"입력 이미지 저장: {temp_input_path}")
        
        # Replicate 설정
        replicate.client.Client(api_token=REPLICATE_API_TOKEN)
        
        # 로그 정보 출력
        logger.info(f"SD 모델: {sd_model}")
        logger.info(f"스케일 팩터: {scale_factor}")
        logger.info(f"출력 형식: {output_format}")
        logger.info(f"유사도: {resemblance}")
        logger.info(f"창의성: {creativity}")
        logger.info(f"다이나믹: {dynamic}")
        logger.info(f"선명도: {sharpen}")
        logger.info(f"시드: {seed}")
        
        # 환경변수에서 가져온 모델로 API 호출
        output = replicate.run(
            REPLICATE_MODEL,
            input={
                "seed": seed,
                "image": open(temp_input_path, "rb"),
                "prompt": prompt,
                "dynamic": dynamic,
                "handfix": "disabled",
                "pattern": False,
                "sharpen": sharpen,
                "sd_model": sd_model,
                "scheduler": "DPM++ 3M SDE Karras",
                "creativity": creativity,
                "lora_links": "",
                "downscaling": False,
                "resemblance": resemblance,
                "scale_factor": scale_factor,
                "tiling_width": 112,
                "output_format": output_format,
                "tiling_height": 144,
                "custom_sd_model": "",
                "negative_prompt": negative_prompt,
                "num_inference_steps": 18,
                "downscaling_resolution": 768
            }
        )
        
        logger.info(f"Replicate API 응답: {output}")
        
        # 결과 이미지 다운로드 및 반환
        if output and isinstance(output, list) and len(output) > 0:
            import requests
            from io import BytesIO
            
            response = requests.get(output[0])
            if response.status_code == 200:
                result_image = Image.open(BytesIO(response.content))
                
                # JPG로 변환 (output_format이 jpg인 경우)
                if output_format.lower() == "jpg":
                    temp_output = BytesIO()
                    if result_image.mode == 'RGBA':
                        result_image = result_image.convert('RGB')
                    result_image.save(temp_output, format='JPEG', quality=95)
                    temp_output.seek(0)
                    result_image = Image.open(temp_output)
                
                logger.info("이미지 화질 개선 완료!")
                return result_image
            else:
                logger.error(f"결과 이미지 다운로드 실패: {response.status_code}")
                return None
        else:
            logger.error("API 응답에서 이미지를 찾을 수 없습니다.")
            return None
            
    except Exception as e:
        logger.error(f"화질 개선 중 오류 발생: {str(e)}")
        raise
    finally:
        # 임시 파일 정리
        if os.path.exists(temp_input_path):
            os.remove(temp_input_path)

def process_and_download(
    password, image, scale_factor, output_format, sd_model,
    resemblance, creativity, dynamic, sharpen, seed
):
    # 환경변수에서 비밀번호 검증
    if not APP_PASSWORD:
        raise ValueError("서버 설정 오류입니다. APP_PASSWORD가 설정되지 않았습니다.")
    
    if password != APP_PASSWORD:
        raise ValueError("잘못된 비밀번호입니다.")
        
    if image is None:
        return None, None

    # 이미지 업스케일
    result_img = upscale_image(
        image=image,
        scale_factor=scale_factor,
        output_format=output_format,
        sd_model=sd_model,
        resemblance=resemblance,
        creativity=creativity,
        dynamic=dynamic,
        sharpen=sharpen,
        seed=int(seed)
    )
    if result_img is None:
        return None, None

    # 파일로 저장
    ext = "jpg" if output_format.lower() == "jpg" else "png"
    # JPG로 저장할 땐 RGB 모드로 변환
    if ext == "jpg" and result_img.mode == "RGBA":
        result_img = result_img.convert("RGB")

    filename = f"upscaled_{uuid.uuid4()}.{ext}"
    filepath = os.path.join(TEMP_DIR, filename)
    result_img.save(filepath, format="JPEG" if ext=="jpg" else "PNG")

    # 리턴: (이미지, 파일 경로)
    return result_img, filepath

def create_interface():
    with gr.Blocks(title="Clarity Upscaler") as demo:
        with gr.Row():
            with gr.Column():
                password_input = gr.Textbox(
                    label="비밀번호", type="password", placeholder="비밀번호를 입력하세요"
                )
                input_image = gr.Image(label="원본 이미지", type="pil")

                with gr.Accordion("고급 설정", open=False):
                    scale_factor  = gr.Slider(1, 4,  value=2, step=0.5, label="확대 비율")
                    output_format = gr.Radio(["jpg","png"], value="jpg",    label="출력 형식")
                    sd_model      = gr.Dropdown(
                        choices=[
                            "juggernaut_reborn.safetensors [338b85bc4f]",
                            "sd_xl_base_1.0.safetensors [39d4e625d]",
                            "sdxl_base_v0.9-9961.safetensors [3048045c]",
                            "realisticVisionV51_v51VAE.safetensors [5ecbfa0ac]",
                        ],
                        value="juggernaut_reborn.safetensors [338b85bc4f]",
                        label="SD 모델"
                    )
                    resemblance = gr.Slider(0.0,1.0, value=0.6, step=0.05, label="원본 유사도")
                    creativity  = gr.Slider(0.0,1.0, value=0.35,step=0.05, label="창의성")
                    dynamic     = gr.Slider(1,10,  value=6, step=1,    label="다이나믹 임계값")
                    sharpen     = gr.Slider(0,2,   value=0, step=0.1,   label="선명도")
                    seed        = gr.Number(value=1337, label="시드", precision=0)

                submit_btn  = gr.Button("실행", variant="primary")

            with gr.Column():
                output_image = gr.Image(label="개선된 이미지")
                download_btn = gr.DownloadButton(label="이미지 다운로드")

        # 클릭 시 두 개의 출력(이미지, 파일 경로)을 반환
        submit_btn.click(
            fn=process_and_download,
            inputs=[
                password_input,
                input_image, scale_factor, output_format, sd_model,
                resemblance, creativity, dynamic, sharpen, seed
            ],
            outputs=[output_image, download_btn]
        )

    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(share=True)