marutitecblic commited on
Commit
7dc9973
·
verified ·
1 Parent(s): cf63769

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +3 -3
handler.py CHANGED
@@ -58,7 +58,7 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
58
  # response = postprocess(model_outputs)
59
  # return response
60
 
61
- class PreTrainedPipeline():
62
  def __init__(self ):
63
  # Load processor and model
64
  self.PROCESSOR = AutoProcessor.from_pretrained(
@@ -102,7 +102,7 @@ class PreTrainedPipeline():
102
  return {"generated_text": generated_text}
103
  # return {"text":prediction[0]}
104
 
105
- @classmethod
106
  def convert_to_rgb(self, image):
107
  if image.mode == "RGB":
108
  return image
@@ -111,7 +111,7 @@ class PreTrainedPipeline():
111
  alpha_composite = Image.alpha_composite(background, image_rgba)
112
  alpha_composite = alpha_composite.convert("RGB")
113
  return alpha_composite
114
- @classmethod
115
  def custom_transform(self, x):
116
  x = self.convert_to_rgb(x)
117
  x = to_numpy_array(x)
 
58
  # response = postprocess(model_outputs)
59
  # return response
60
 
61
+ class EndpointHandler():
62
  def __init__(self ):
63
  # Load processor and model
64
  self.PROCESSOR = AutoProcessor.from_pretrained(
 
102
  return {"generated_text": generated_text}
103
  # return {"text":prediction[0]}
104
 
105
+ # @classmethod
106
  def convert_to_rgb(self, image):
107
  if image.mode == "RGB":
108
  return image
 
111
  alpha_composite = Image.alpha_composite(background, image_rgba)
112
  alpha_composite = alpha_composite.convert("RGB")
113
  return alpha_composite
114
+ # @classmethod
115
  def custom_transform(self, x):
116
  x = self.convert_to_rgb(x)
117
  x = to_numpy_array(x)