Update app.py

app.py

@@ -3,7 +3,6 @@ import replicate
 import os
 import requests
 import tempfile
-from pathlib import Path
 import logging
 import base64
 
@@ -16,77 +15,54 @@ logger = logging.getLogger(__name__)
 def process_image(password, input_image):
     # 비밀번호 검증
     if password != "1089":
-        raise ValueError("잘못된 입력력입니다.")
-    
+        raise ValueError("잘못된 비밀번호입니다.")
+    if not os.getenv("REPLICATE_API_TOKEN"):
+        logger.error("REPLICATE_API_TOKEN이 설정되지 않았습니다.")
+        return None, None
+    if input_image is None:
+        logger.error("입력 이미지가 없습니다.")
+        return None, None
+
     try:
-        logger.info("=== 처리 시작 ===")
-        
-        api_token = os.getenv("REPLICATE_API_TOKEN")
-        if not api_token:
-            logger.error("REPLICATE_API_TOKEN이 설정되지 않았습니다.")
-            return None
-        
-        logger.info(f"API 토큰 확인 완료: {api_token[:5]}...")
-        
-        if input_image is None:
-            logger.error("입력 이미지가 없습니다.")
-            return None
-        
-        logger.info(f"입력 이미지 경로: {input_image}")
-        logger.info(f"입력 이미지 타입: {type(input_image)}")
-        
-        try:
-            with open(input_image, "rb") as image_file:
-                image_data = base64.b64encode(image_file.read()).decode('utf-8')
-                image_uri = f"data:image/png;base64,{image_data}"
-                
-            logger.info("이미지 변환 완료")
-            logger.info(f"Data URI 길이: {len(image_uri)}")
-            
-            logger.info("모델 실행 시작")
-            
-            output = replicate.run(
-                "851-labs/background-remover:a029dff38972b5fda4ec5d75d7d1cd25aeff621d2cf4946a41055d7db66b80bc",
-                input={"image": image_uri}
-            )
-            
-            logger.info("모델 실행 완료")
-            logger.info(f"결과 타입: {type(output)}")
-            
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
-                # replicate.run() 반환 타입이 file‐like 객체일 경우
-                if hasattr(output, "read"):
-                    tmp_file.write(output.read())
-                # 또는 URL 리스트인 경우 다운로드
-                elif isinstance(output, (list, tuple)) and output and isinstance(output[0], str):
-                    resp = requests.get(output[0])
-                    resp.raise_for_status()
-                    tmp_file.write(resp.content)
-                tmp_file_path = tmp_file.name
-            
-            logger.info(f"결과 이미지 저장 완료: {tmp_file_path}")
-            return tmp_file_path
-        
-        except Exception as file_error:
-            logger.error(f"파일 처리 오류: {str(file_error)}")
-            raise file_error
-    
+        # base64로 변환
+        with open(input_image, "rb") as f:
+            data = base64.b64encode(f.read()).decode()
+        image_uri = f"data:image/png;base64,{data}"
+
+        # 모델 실행
+        output = replicate.run(
+            "851-labs/background-remover:a029dff38972b5fda4ec5d75d7d1cd25aeff621d2cf4946a41055d7db66b80bc",
+            input={"image": image_uri}
+        )
+
+        # 결과 저장
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
+            if hasattr(output, "read"):
+                tmp.write(output.read())
+            elif isinstance(output, (list, tuple)) and output:
+                resp = requests.get(output[0])
+                resp.raise_for_status()
+                tmp.write(resp.content)
+            out_path = tmp.name
+
+        return out_path, out_path
+
     except replicate.exceptions.ReplicateError as re:
-        logger.error(f"API 오류: {str(re)}")
-        return None
-    
+        logger.error(f"API 오류: {re}")
+        return None, None
     except Exception as e:
-        logger.error(f"예상치 못한 오류: {str(e)}")
-        return None
+        logger.error(f"예상치 못한 오류: {e}")
+        return None, None
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# Background Remover (Password Protected)")
+
     with gr.Row():
         with gr.Column():
-            # 비밀번호 입력 필드
             password_input = gr.Textbox(
-                label="정보", 
-                type="password", 
-                placeholder="정보를 입력하세요"
+                label="비밀번호",
+                type="password",
+                placeholder="비밀번호를 입력하세요"
             )
             input_image = gr.Image(
                 type="filepath",
@@ -94,21 +70,24 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 interactive=True
             )
             process_btn = gr.Button("실행", variant="primary")
-        
+
         with gr.Column():
             output_image = gr.Image(
                 type="filepath",
-                label="출력 이미지",
+                label="결과 이미지",
                 interactive=False
             )
-    
-    # 이벤트 핸들러: 비밀번호와 이미지 경로 두 개를 인풋으로
+            download_btn = gr.DownloadButton(
+                label="이미지 다운로드"
+            )
+
+    # click 시 (표시용, 다운로드용) 두 개 리턴
     process_btn.click(
         fn=process_image,
         inputs=[password_input, input_image],
-        outputs=[output_image]
+        outputs=[output_image, download_btn]
     )
-    
+
     gr.Markdown("---")
 
 if __name__ == "__main__":