amaye15 commited on
Commit
7e994bd
·
1 Parent(s): a44d3b7
Files changed (1) hide show
  1. 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(os.getcwd())
44
- self.model = AutoModel.from_pretrained(
45
- os.getcwd(), trust_remote_code=True
46
- ).to(self.device)
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}")