Update app.py

app.py

@@ -50,6 +50,9 @@ def load_img(source, output_type="pil"):
 
 def process(image):
     image_size = image.size
+    # RGBA 이미지를 RGB로 변환
+    if image.mode == 'RGBA':
+        image = image.convert('RGB')
     input_images = transform_image(image).unsqueeze(0).to("cpu")  # GPU -> CPU로 변경
     # Prediction
     with torch.no_grad():
@@ -57,8 +60,10 @@ def process(image):
     pred = preds[0].squeeze()
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
-    image.putalpha(mask)
-    return image
+    # 결과 이미지에 알파 채널 추가
+    result_image = image.copy()
+    result_image.putalpha(mask)
+    return result_image
 
 def fn(image):
     im = load_img(image, output_type="pil")