amaye15
commited on
Commit
·
7e994bd
1
Parent(s):
a44d3b7
Debug
Browse files- handler.py +5 -5
handler.py
CHANGED
@@ -25,7 +25,7 @@ class EndpointHandler:
|
|
25 |
device: The device (CPU or CUDA) used to run model inference.
|
26 |
"""
|
27 |
|
28 |
-
def __init__(self):
|
29 |
"""
|
30 |
Initializes the EndpointHandler with the model and processor from the current directory.
|
31 |
"""
|
@@ -40,10 +40,10 @@ class EndpointHandler:
|
|
40 |
# Load the model and processor from the current directory
|
41 |
self.logger.info("Loading model and processor from the current directory.")
|
42 |
try:
|
43 |
-
self.processor = AutoImageProcessor.from_pretrained(
|
44 |
-
self.model = AutoModel.from_pretrained(
|
45 |
-
|
46 |
-
)
|
47 |
self.logger.info("Model and processor loaded successfully.")
|
48 |
except Exception as e:
|
49 |
self.logger.error(f"Failed to load model or processor: {e}")
|
|
|
25 |
device: The device (CPU or CUDA) used to run model inference.
|
26 |
"""
|
27 |
|
28 |
+
def __init__(self, path: str = ""):
|
29 |
"""
|
30 |
Initializes the EndpointHandler with the model and processor from the current directory.
|
31 |
"""
|
|
|
40 |
# Load the model and processor from the current directory
|
41 |
self.logger.info("Loading model and processor from the current directory.")
|
42 |
try:
|
43 |
+
self.processor = AutoImageProcessor.from_pretrained(path)
|
44 |
+
self.model = AutoModel.from_pretrained(path, trust_remote_code=True).to(
|
45 |
+
self.device
|
46 |
+
)
|
47 |
self.logger.info("Model and processor loaded successfully.")
|
48 |
except Exception as e:
|
49 |
self.logger.error(f"Failed to load model or processor: {e}")
|