Update app.py

app.py

@@ -8,12 +8,13 @@ from torchvision import transforms
 from PIL import Image
 import os
 
-# GPU 설정을 CPU로 변경
+# 모델 로드 및 CPU로 설정
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
 birefnet.to("cpu")  # GPU -> CPU로 변경
 
+# 이미지 전처리
 transform_image = transforms.Compose(
     [
         transforms.Resize((1024, 1024)),
@@ -22,6 +23,18 @@ transform_image = transforms.Compose(
     ]
 )
 
+def process(image):
+    image_size = image.size
+    input_images = transform_image(image).unsqueeze(0).to("cpu")  # CPU로 변경
+    # 예측 수행
+    with torch.no_grad():
+        preds = birefnet(input_images)[-1].sigmoid().cpu()
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    mask = pred_pil.resize(image_size)
+    image.putalpha(mask)
+    return image
+
 def fn(image):
     im = load_img(image, output_type="pil")
     im = im.convert("RGB")
@@ -34,19 +47,7 @@ def fn(image):
     jpg_path = "output.jpg"
     jpg_image.save(jpg_path, format="JPEG")
     
-    return processed_image, jpg_path
-
-def process(image):
-    image_size = image.size
-    input_images = transform_image(image).unsqueeze(0).to("cpu")  # GPU -> CPU로 변경
-    # Prediction
-    with torch.no_grad():
-        preds = birefnet(input_images)[-1].sigmoid().cpu()
-    pred = preds[0].squeeze()
-    pred_pil = transforms.ToPILImage()(pred)
-    mask = pred_pil.resize(image_size)
-    image.putalpha(mask)
-    return image
+    return [processed_image], jpg_path  # ImageSlider는 리스트를 기대함
 
 def process_file(f):
     name_path = f.rsplit(".", 1)[0] + ".png"
@@ -56,13 +57,15 @@ def process_file(f):
     transparent.save(name_path)
     return name_path
 
+# Gradio 컴포넌트 정의
 slider1 = ImageSlider(label="Processed Image", type="pil")
 image_upload = gr.Image(label="Upload an image")
 output_download = gr.File(label="Download JPG File")
 
-# 새로운 샘플 이미지 추가
+# 새로운 샘플 이미지 추가 (app.py와 동일한 폴더에 위치해야 함)
 sample_images = ["1.png", "2.jpg", "3.png"]
 
+# Gradio 인터페이스 설정
 tab1 = gr.Interface(
     fn=fn, 
     inputs=image_upload,