fix: build failure

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import numpy as np
 import math
+import os
+import shutil
 import torch
 
 from PIL import Image, ImageDraw
@@ -13,9 +15,70 @@ config = Config()
 cuda = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
-docscanner = load_docscanner_model(
-    cuda, path_l=config.get_rec_model_path, path_m=config.get_seg_model_path
-)
+_docscanner = None
+
+
+def _is_git_lfs_pointer_file(path: str) -> bool:
+    try:
+        with open(path, "rb") as f:
+            head = f.read(200)
+        return b"version https://git-lfs.github.com/spec/v1" in head
+    except FileNotFoundError:
+        return False
+
+
+def _try_download_from_hf_hub(repo_id: str, filename: str, local_path: str) -> bool:
+    try:
+        from huggingface_hub import hf_hub_download  # type: ignore
+
+        downloaded = hf_hub_download(repo_id=repo_id, filename=filename)
+        os.makedirs(os.path.dirname(local_path), exist_ok=True)
+        shutil.copyfile(downloaded, local_path)
+        return True
+    except Exception:
+        return False
+
+
+def _ensure_weights_present() -> None:
+    seg_path = config.get_seg_model_path
+    rec_path = config.get_rec_model_path
+
+    needs_seg = (not os.path.exists(seg_path)) or _is_git_lfs_pointer_file(seg_path)
+    needs_rec = (not os.path.exists(rec_path)) or _is_git_lfs_pointer_file(rec_path)
+    if not (needs_seg or needs_rec):
+        return
+
+    repo_id = (
+        os.getenv("SPACE_ID")
+        or os.getenv("HF_SPACE_ID")
+        or os.getenv("HF_REPO_ID")
+        or os.getenv("REPO_ID")
+    )
+    if repo_id:
+        if needs_seg:
+            _try_download_from_hf_hub(repo_id, os.path.basename(seg_path), seg_path)
+        if needs_rec:
+            _try_download_from_hf_hub(repo_id, os.path.basename(rec_path), rec_path)
+
+
+def get_docscanner():
+    global _docscanner
+    if _docscanner is not None:
+        return _docscanner
+
+    _ensure_weights_present()
+
+    seg_path = config.get_seg_model_path
+    rec_path = config.get_rec_model_path
+    if _is_git_lfs_pointer_file(seg_path) or _is_git_lfs_pointer_file(rec_path):
+        raise RuntimeError(
+            "Model weight files look like Git LFS pointers. "
+            "Make sure LFS objects are downloaded (e.g. `git lfs pull`) "
+            "or allow the Space to download them from the Hub at runtime."
+        )
+
+    _docscanner = load_docscanner_model(cuda, path_l=rec_path, path_m=seg_path)
+    return _docscanner
 
 # 좌표를 초기화하는 함수
 def reset_points(image, state):
@@ -38,7 +101,8 @@ def rotate_image(image):
     return rotated_image, state
 
 def reset_image(image, state):
-    out_image, msk_np = docscanner_rec(image, docscanner)
+    docscanner = get_docscanner()
+    out_image, msk_np = docscanner_rec(image, docscanner, cuda)
     state = list(get_corner(mask2point(mask=msk_np)))
     
     img = Image.fromarray(image)
@@ -59,6 +123,7 @@ def reset_image(image, state):
     return img, state
 
 def auto_point_detect(image):
+    docscanner = get_docscanner()
     out_image, msk_np = docscanner_rec(image, docscanner, cuda)
     state = list(get_corner(mask2point(mask=msk_np)))
     
@@ -126,6 +191,7 @@ def sort_corners(corners):
 def convert(image, state):
     h,w = image.shape[:2]
     if len(state) < 4:
+        docscanner = get_docscanner()
         out_image, msk_np = docscanner_rec(image, docscanner, cuda)
         out_image = out_image[:,:,::-1]
     elif len(state) ==4:
@@ -151,7 +217,12 @@ with gr.Blocks(css=css) as demo:
     with gr.Row():
         with gr.Column():
             text = gr.Textbox("입력 이미지(코너를 클릭하세요)", show_label=False)
-            image_input = gr.Image(show_label=False, interactive=True, elem_classes="image-container")
+            image_input = gr.Image(
+                show_label=False,
+                interactive=True,
+                elem_classes="image-container",
+                type="numpy",
+            )
             clear_button = gr.Button("Clear Points")
             cutting_button = gr.Button("Cutting Image(need more than 2 points)")
             rotating_button = gr.Button("Rotate Image(clock wise 90 degree)")
@@ -159,13 +230,13 @@ with gr.Blocks(css=css) as demo:
             convert_button = gr.Button("Convert Image")
         with gr.Column():
             text = gr.Textbox("변환될 영역", show_label=False)
-            image_output = gr.Image(show_label=False)
+            image_output = gr.Image(show_label=False, type="pil")
             # state_display = gr.Textbox(label="Current State")
             # coordinates_text = gr.Textbox(label="Coordinates", placeholder="Enter coordinates (x, y) for each point")
             # update_coords_button = gr.Button("Update Coordinates")
         with gr.Column():
             text = gr.Textbox("결과 이미지", show_label=False)
-            result_image = gr.Image(show_label=False, format="png")
+            result_image = gr.Image(show_label=False, format="png", type="numpy")
 
     # # 이미지 위에서 클릭 이벤트 처리
     image_input.select(draw_polygon_on_image, inputs=[image_input,state], outputs=[image_output,state])
@@ -183,4 +254,7 @@ with gr.Blocks(css=css) as demo:
     # 변환 버튼
     convert_button.click(fn=convert, inputs=[image_input,state], outputs=result_image)
 
-demo.launch(share=True)
+is_spaces = bool(
+    os.getenv("SPACE_ID") or os.getenv("HF_SPACE_ID") or os.getenv("SYSTEM") == "spaces"
+)
+demo.launch(share=not is_spaces and bool(os.getenv("GRADIO_SHARE")))