Upload app.py

app.py

@@ -21,32 +21,37 @@ from transformers import (
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"🖥️ Using compute device: {device}")
 
-# Load models
-print("⏳ Loading SAM3 Models permanently into memory...")
-try:
-    # 오프라인 모드로 캐시에서 로드 시도
-    print("   ... Loading from local cache (offline mode)")
-    IMG_MODEL = Sam3Model.from_pretrained("DiffusionWave/sam3", local_files_only=True, device_map="cpu", torch_dtype=torch.float32)
-    IMG_PROCESSOR = Sam3Processor.from_pretrained("DiffusionWave/sam3", local_files_only=True)
-
-    TRK_MODEL = Sam3TrackerModel.from_pretrained("DiffusionWave/sam3", local_files_only=True, device_map="cpu", torch_dtype=torch.float32)
-    TRK_PROCESSOR = Sam3TrackerProcessor.from_pretrained("DiffusionWave/sam3", local_files_only=True)
-
-    print("✅ All Models loaded successfully from local cache!")
-except Exception as e:
-    print(f"❌ Cache loading failed: {e}")
-    print("   Trying online loading...")
+# Models will be loaded lazily in functions to avoid build timeouts
+IMG_MODEL = None
+IMG_PROCESSOR = None
+TRK_MODEL = None
+TRK_PROCESSOR = None
+
+@spaces.GPU
+def load_models():
+    """Lazy load models when needed"""
+    global IMG_MODEL, IMG_PROCESSOR, TRK_MODEL, TRK_PROCESSOR
+
+    if IMG_MODEL is not None:
+        return True
+
+    print("⏳ Loading SAM3 Models...")
     try:
-        IMG_MODEL = Sam3Model.from_pretrained("DiffusionWave/sam3", device_map="cpu", torch_dtype=torch.float32)
+        # GPU가 사용 가능하면 GPU로 로드
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        dtype = torch.float16 if device == "cuda" else torch.float32
+
+        IMG_MODEL = Sam3Model.from_pretrained("DiffusionWave/sam3", device_map=device, torch_dtype=dtype)
         IMG_PROCESSOR = Sam3Processor.from_pretrained("DiffusionWave/sam3")
 
-        TRK_MODEL = Sam3TrackerModel.from_pretrained("DiffusionWave/sam3", device_map="cpu", torch_dtype=torch.float32)
+        TRK_MODEL = Sam3TrackerModel.from_pretrained("DiffusionWave/sam3", device_map=device, torch_dtype=dtype)
         TRK_PROCESSOR = Sam3TrackerProcessor.from_pretrained("DiffusionWave/sam3")
 
-        print("✅ All Models loaded successfully (CPU mode)!")
-    except Exception as e2:
-        print(f"❌ Online loading also failed: {e2}")
-        IMG_MODEL = IMG_PROCESSOR = TRK_MODEL = TRK_PROCESSOR = None
+        print(f"✅ All Models loaded successfully on {device}!")
+        return True
+    except Exception as e:
+        print(f"❌ Model loading failed: {e}")
+        return False
 
 # ============ LAYER MANAGEMENT ============
 class LayerManager:
@@ -248,9 +253,12 @@ def draw_points_on_image(image, layer_manager):
 
 # ============ UI FUNCTIONS ============
 def update_layer_selector_choices(manager):
-    """레이어 선택 드롭다운의 choices 업데이트"""
-    choices = [(layer['name'], lid) for lid, layer in manager.layers.items()]
-    return gr.Dropdown(choices=choices, interactive=True, value=manager.current_layer_id)
+    """레이어 선택 라디오 버튼의 choices 업데이트"""
+    choices = [layer['name'] for layer in manager.layers.values()]
+    current_value = None
+    if manager.current_layer_id and manager.current_layer_id in manager.layers:
+        current_value = manager.layers[manager.current_layer_id]['name']
+    return gr.Radio(choices=choices, interactive=True, value=current_value)
 
 def create_new_layer(name, current_manager):
     """새 레이어 생성"""
@@ -384,6 +392,11 @@ def segment_all_layers(current_manager, image, opacity, border_width):
             print(f"\n[segment_all_layers] Processing layer: {layer_name}")
             print(f"[segment_all_layers] Points: {len(layer['points'])}, Labels: {layer['point_labels']}")
             
+            # Load models if needed
+            if not load_models():
+                print(f"[segment_all_layers] Failed to load models for layer: {layer_name}")
+                continue
+
             # SAM3 Tracker로 세그멘테이션
             points_list = layer['points']
             labels_list = layer['point_labels']
@@ -391,7 +404,9 @@ def segment_all_layers(current_manager, image, opacity, border_width):
             input_points = [[points_list]]
             input_labels = [[labels_list]]
 
-            inputs = TRK_PROCESSOR(images=image, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(device)
+            # Use the same device as the model
+            model_device = next(TRK_MODEL.parameters()).device
+            inputs = TRK_PROCESSOR(images=image, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(model_device)
 
             with torch.no_grad():
                 outputs = TRK_MODEL(**inputs, multimask_output=False)
@@ -485,8 +500,8 @@ with gr.Blocks() as demo:
                 gr.Markdown("### Layers Status")
                 layer_buttons_html = gr.HTML("<div style='padding: 10px; text-align: center; color: #888;'>No layers created</div>")
                 
-                # 레이어 선택
-                layer_selector = gr.Dropdown(label="Select Layer to Add Points", choices=[], interactive=True, value=None)
+                # 레이어 선택 (라디오 버튼으로 변경)
+                layer_selector = gr.Radio(label="Select Layer to Add Points", choices=[], interactive=True)
                 
                 # 포인트 모드 선택
                 gr.Markdown("### Point Mode")
@@ -537,13 +552,21 @@ with gr.Blocks() as demo:
         )
         
         # 레이어 선택
-        def on_layer_select(layer_id, mgr):
+        def on_layer_select(selected_name, mgr):
             if mgr is None:
                 mgr = LayerManager()
 
-            if layer_id:
-                mgr.set_current_layer(layer_id)
-                return mgr, create_layer_status_html(mgr), f"Layer '{mgr.layers[layer_id]['name']}' selected"
+            if selected_name:
+                # 이름으로 layer_id 찾기
+                layer_id = None
+                for lid, layer in mgr.layers.items():
+                    if layer['name'] == selected_name:
+                        layer_id = lid
+                        break
+
+                if layer_id:
+                    mgr.set_current_layer(layer_id)
+                    return mgr, create_layer_status_html(mgr), f"Layer '{selected_name}' selected"
             return mgr, create_layer_status_html(mgr), "Please select a layer"
 
         layer_selector.change(