Update app.py

app.py

@@ -1,192 +1,248 @@
 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
+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.FileHandler("app.log"),
-        logging.StreamHandler(sys.stdout)
-    ]
+    handlers=[logging.StreamHandler(sys.stdout)]
 )
-logger = logging.getLogger("enhancer-app")
+logger = logging.getLogger("clarity-upscaler")
 
-# 임시 파일 저장 함수
-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
+# 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 enhance_image(
-    image,
+def upscale_image(
+    image, 
+    scale_factor=2, 
     output_format="jpg",
-    enhancement_level=2  # Fixed enhancement level
+    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:
-        if image is None:
-            return None, None, "이미지를 업로드해주세요."
+        # 이미지 저장
+        temp_input_path = os.path.join(TEMP_DIR, f"input_{uuid.uuid4()}.png")
+        image.save(temp_input_path)
+        logger.info(f"입력 이미지 저장: {temp_input_path}")
         
-        # 이미지 처리
-        temp_paths = []
-        img_path = save_uploaded_file(image)
-        if not img_path:
-            return None, None, "이미지 처리에 실패했습니다."
+        # Replicate 설정
+        replicate.client.Client(api_token=REPLICATE_API_TOKEN)
         
-        temp_paths.append(img_path)
+        # 로그 정보 출력
+        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}")
         
-        # Replicate API 토큰 확인
-        if not os.environ.get("REPLICATE_API_TOKEN"):
-            return None, None, "API 토큰이 설정되지 않았습니다."
+        # 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
+            }
+        )
         
-        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)"
-                }
-            )
+        logger.info(f"Replicate API 응답: {output}")
+        
+        # 결과 이미지 다운로드 및 반환
+        if output and isinstance(output, list) and len(output) > 0:
+            import requests
+            from io import BytesIO
             
-            # Replicate API 응답 처리
-            if output and isinstance(output, list) and len(output) > 0:
-                enhanced_url = output[0]
+            response = requests.get(output[0])
+            if response.status_code == 200:
+                result_image = Image.open(BytesIO(response.content))
                 
-                # 향상된 이미지 다운로드
-                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 응답 없음"
+                # 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)
                 
-        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}")
+                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"Error in enhance_image function: {e}")
-        logger.error(traceback.format_exc())
-        return None, None, f"{str(e)}"
+        logger.error(f"화질 개선 중 오류 발생: {str(e)}")
+        raise
+    finally:
+        # 임시 파일 정리
+        if os.path.exists(temp_input_path):
+            os.remove(temp_input_path)
 
-# Gradio 인터페이스 구성
-def create_gradio_interface():
-    try:
-        with gr.Blocks() as app:
-            with gr.Row():
-                with gr.Column():
-                    input_image = gr.Image(type="pil")
-                    
-                    # Hidden parameters with fixed values
-                    output_format = gr.Dropdown(
-                        choices=["jpg", "png", "webp"],
-                        value="jpg",
-                        visible=False
-                    )
-                    enhancement_level = gr.Number(
-                        value=2,
-                        visible=False
+# 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 모델"
                     )
-                    
-                    process_btn = gr.Button("실행")
+                    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)
                 
-                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
+                submit_btn = gr.Button("화질 개선 시작", variant="primary")
                 
-                enhanced_img, download_path, error = enhance_image(
-                    img, format, level
-                )
+            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)
                 
-                if error:
-                    return None, None
+            def flush(self):
+                pass
                 
-                return enhanced_img, download_path
+            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, "오류: 이미지를 업로드해주세요."
             
-            process_btn.click(
-                on_process,
-                inputs=[input_image, output_format, enhancement_level],
-                outputs=[output_image, output_download]
-            )
+            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()
         
-        return app
-    except Exception as e:
-        logger.error(f"Error creating Gradio interface: {e}")
-        logger.error(traceback.format_exc())
-        raise
+        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__":
-    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())
+    demo = create_interface()
+    demo.launch(share=True)