Update main.py

main.py

@@ -16,47 +16,47 @@ PY_MODULES = {
     "texture.py": "TextureDetector",
     "skin_tone.py": "SkinToneDetector",
     "oiliness.py": "OilinessDetector", 
-    "wrinkle.py": "WrinkleDetector",
+    "wrinkle_unet.py": "WrinkleDetector",
     "age.py": "AgePredictor"
 }
-def load_model(token):
-    """Download and load ONNX model"""
-    model_path = hf_hub_download(
-        repo_id=REPO_ID,
-        filename="model/wrinkle_model.onnx.data",  # Adjust path if needed
-        token=token
-    )
+# def load_model(token):
+#     """Download and load ONNX model"""
+#     model_path = hf_hub_download(
+#         repo_id=REPO_ID,
+#         filename="model/wrinkle_model.onnx.data",  # Adjust path if needed
+#         token=token
+#     )
     
-    # Create ONNX Runtime session
-    session = ort.InferenceSession(
-        model_path,
-        providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] 
-        if torch.cuda.is_available() else ['CPUExecutionProvider']
-    )
+#     # Create ONNX Runtime session
+#     session = ort.InferenceSession(
+#         model_path,
+#         providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] 
+#         if torch.cuda.is_available() else ['CPUExecutionProvider']
+#     )
     
-    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 = "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
     
-#     # if not token:
-#     #     raise ValueError("HUGGINGFACE_HUB_TOKEN environment variable is required for private repo access.")
+    # 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)
+    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
+    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)