Update main.py

main.py

@@ -15,31 +15,47 @@ PY_MODULES = {
     "texture.py": "TextureDetector",
     "skin_tone.py": "SkinToneDetector",
     "oiliness.py": "OilinessDetector", 
-    "wrinkle_unet.py": "WrinkleDetector",
+    "wrinkle.py": "WrinkleDetector",
     "age.py": "AgePredictor"
 }
-
 def load_model(token):
-    repo_id = "IFMedTech/Skin-Analysis"
-    filename = "model/wrinkles_unet_v1.pth"
-    # token = os.environ.get("HUGGINGFACE_HUB_TOKEN")  # Set this env var with your token
+    """Download and load ONNX model"""
+    model_path = hf_hub_download(
+        repo_id=REPO_ID,
+        filename="model/wrinkle_model.onnx",  # Adjust path if needed
+        token=token
+    )
     
-    # if not token:
-    #     raise ValueError("HUGGINGFACE_HUB_TOKEN environment variable is required for private repo access.")
+    # Create ONNX Runtime session
+    session = ort.InferenceSession(
+        model_path,
+        providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] 
+        if torch.cuda.is_available() else ['CPUExecutionProvider']
+    )
     
-    model_path = hf_hub_download(repo_id=repo_id, filename=filename, token=token)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    return session, device
+# def load_model(token):
+#     repo_id = "IFMedTech/Skin-Analysis"
+#     filename = "model/wrinkles_unet_v1.pth"
+#     # token = os.environ.get("HUGGINGFACE_HUB_TOKEN")  # Set this env var with your token
     
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = smp.Unet(
-        encoder_name="resnet34",
-        encoder_weights=None,
-        in_channels=3,
-        classes=1
-    )
-    model.load_state_dict(torch.load(model_path, map_location=device))
-    model.to(device)
-    model.eval()
-    return model, device
+#     # if not token:
+#     #     raise ValueError("HUGGINGFACE_HUB_TOKEN environment variable is required for private repo access.")
+    
+#     model_path = hf_hub_download(repo_id=repo_id, filename=filename, token=token)
+    
+#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#     model = smp.Unet(
+#         encoder_name="resnet34",
+#         encoder_weights=None,
+#         in_channels=3,
+#         classes=1
+#     )
+#     model.load_state_dict(torch.load(model_path, map_location=device))
+#     model.to(device)
+#     model.eval()
+#     return model, device
 
 def dynamic_import(module_path, class_name):
     spec = importlib.util.spec_from_file_location(class_name, module_path)