app.py

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

# 로깅 설정
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger("clarity-upscaler")

# Replicate API 토큰 설정
REPLICATE_API_TOKEN = os.environ.get("REPLICATE_API_TOKEN", "")
if not REPLICATE_API_TOKEN:
    logger.warning("REPLICATE_API_TOKEN이 설정되지 않았습니다. 환경 변수를 설정해주세요.")

# 임시 디렉토리 생성
TEMP_DIR = os.path.join(tempfile.gettempdir(), "clarity_upscaler")
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 REPLICATE_API_TOKEN == "":
        raise ValueError("REPLICATE_API_TOKEN이 설정되지 않았습니다.")
    
    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("=== Clarity Upscaler 실행 정보 ===")
        logger.info(f"모델: philz1337x/clarity-upscaler")
        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(
            "philz1337x/clarity-upscaler:dfad41707589d68ecdccd1dfa600d55a208f9310748e44bfe35b4a6291453d5e",
            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)

# Gradio 인터페이스 생성
def create_interface():
    with gr.Blocks(title="Clarity Upscaler - 이미지 화질 개선") as demo:
        gr.Markdown("# Clarity Upscaler - 이미지 화질 개선 도구")
        gr.Markdown("Replicate의 Clarity Upscaler 모델을 사용하여 이미지의 화질을 개선합니다.")
        
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="원본 이미지", type="pil")
                
                with gr.Accordion("고급 설정", open=False):
                    scale_factor = gr.Slider(minimum=1, maximum=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(minimum=0.0, maximum=1.0, value=0.6, step=0.05, label="원본 유사도")
                    creativity = gr.Slider(minimum=0.0, maximum=1.0, value=0.35, step=0.05, label="창의성")
                    dynamic = gr.Slider(minimum=1, maximum=10, value=6, step=1, label="다이나믹 임계값")
                    sharpen = gr.Slider(minimum=0, maximum=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="개선된 이미지")
                log_output = gr.Textbox(label="로그", lines=10)
        
        # 로그 캡처 함수
        class LogCapture:
            def __init__(self):
                self.logs = []
                self.handler = self.create_handler()
                
            def create_handler(self):
                handler = logging.StreamHandler(stream=self)
                handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
                return handler
                
            def write(self, message):
                self.logs.append(message)
                
            def flush(self):
                pass
                
            def get_logs(self):
                return "".join(self.logs)
                
            def clear(self):
                self.logs.clear()
        
        log_capture = LogCapture()
        logger.addHandler(log_capture.handler)
        
        # 제출 버튼 클릭 시 동작
        def process_image(image, scale_factor, output_format, sd_model, resemblance, creativity, dynamic, sharpen, seed):
            if image is None:
                return None, "오류: 이미지를 업로드해주세요."
            
            log_capture.clear()
            
            try:
                result = 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)
                )
                return result, log_capture.get_logs()
            except Exception as e:
                logger.error(f"이미지 처리 중 오류 발생: {str(e)}")
                return None, log_capture.get_logs()
        
        submit_btn.click(
            process_image,
            inputs=[
                input_image, scale_factor, output_format, sd_model, 
                resemblance, creativity, dynamic, sharpen, seed
            ],
            outputs=[output_image, log_output]
        )
        
    return demo

# 메인 실행
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(share=True)