Update app.py

app.py

@@ -13,6 +13,7 @@ from PIL import Image
 import requests
 import yaml
 import numpy as np
+import gc
 
 from src.core import YAMLConfig
 
@@ -108,7 +109,7 @@ def download_weights(model_name):
     print(f"Downloaded weights to: {weights_path}")
     return weights_path
 
-
+@torch.no_grad()
 def process_image_for_gradio(model, device, image, model_name, threshold=0.4):
     """Process image function for Gradio interface"""
     if isinstance(image, np.ndarray):
@@ -185,10 +186,37 @@ class ModelWrapper(nn.Module):
         return outputs
 
 
+# YAMLConfig 클래스의 내부 상태를 초기화하는 함수 추가
+def reset_yaml_config():
+    """YAMLConfig 클래스의 내부 상태를 초기화"""
+    # 클래스 내부에 캐싱된 정보가 있다면 삭제
+    if hasattr(YAMLConfig, '_instances'):
+        YAMLConfig._instances = {}
+    if hasattr(YAMLConfig, '_configs'):
+        YAMLConfig._configs = {}
+    
+    # 가능한 다른 모든 모듈 캐시 리셋
+    import importlib
+    for module_name in list(sys.modules.keys()):
+        if module_name.startswith('src.'):
+            try:
+                importlib.reload(sys.modules[module_name])
+            except:
+                pass
+
 def load_model(model_name):
+    # 모델 로드 전에 CUDA 캐시와 가비지 컬렉션 정리
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
+    
+    # YAMLConfig 내부 상태 초기화
+    reset_yaml_config()
+    
     cfgfile = model_configs[model_name]["cfgfile"]
     weights_path = download_weights(model_name)
     
+    # 완전히 새로운 YAMLConfig 인스턴스 생성
     cfg = YAMLConfig(cfgfile, resume=weights_path)
 
     if "HGNetv2" in cfg.yaml_cfg:
@@ -197,7 +225,11 @@ def load_model(model_name):
     checkpoint = torch.load(weights_path, map_location="cpu")
     state = checkpoint["ema"]["module"] if "ema" in checkpoint else checkpoint["model"]
     
-    cfg.model.load_state_dict(state)
+    # 모델 생성 전 한번 더 확인
+    torch.cuda.empty_cache()
+    gc.collect()
+    
+    cfg.model.load_state_dict(state, strict=False)
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = ModelWrapper(cfg).to(device)
@@ -205,26 +237,34 @@ def load_model(model_name):
 
     return model, device
 
-
-# Dictionary to store loaded models
-loaded_models = {}
-
 @spaces.GPU
 def process_image(image, model_name, confidence_threshold):
     """Main processing function for Gradio interface"""
-    global loaded_models
     
-    # Load model if not already loaded
-    if model_name not in loaded_models:
+    # 모든 사용 가능한 CUDA 장치 메모리 확보
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    
+    # 모든 Python 객체 가비지 컬렉션
+    gc.collect()
+    
+    try:
         print(f"Loading model: {model_name}")
         model, device = load_model(model_name)
-        loaded_models[model_name] = (model, device)
-    else:
-        print(f"Using cached model: {model_name}")
-        model, device = loaded_models[model_name]
+        
+        # 이미지 처리
+        result = process_image_for_gradio(model, device, image, model_name, confidence_threshold)
+        
+        # 모델 객체 및 관련 데이터 명시적 제거
+        del model
+        
+    finally:
+        # 항상 메모리 정리 보장
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
     
-    # Process the image
-    return process_image_for_gradio(model, device, image, model_name, confidence_threshold)
+    return result
 
 
 # Create Gradio interface
@@ -256,4 +296,6 @@ demo = gr.Interface(
     ]
 )
 
-demo.launch(debug=True)
+if __name__ == "__main__":
+    # Launch the Gradio app
+    demo.launch(share=True)