philschmid HF staff commited on
Commit
394463f
1 Parent(s): fdae70a

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +4 -4
handler.py CHANGED
@@ -74,18 +74,18 @@ class EndpointHandler():
74
  :param data: A dictionary contains `inputs` and optional `image` field.
75
  :return: A dictionary with `image` field contains image in base64.
76
  """
77
- prompt = data.pop("prompt", None)
78
  image = data.pop("image", None)
79
  controlnet_type = data.pop("controlnet_type", None)
80
 
81
  # Check if neither prompt nor image is provided
82
- if prompt is None and encoded_image is None:
83
  return {"error": "Please provide a prompt and base64 encoded image."}
84
 
85
  # Check if a new controlnet is provided
86
  if controlnet_type is not None and controlnet_type != self.control_type:
87
- print(f"Loading {controlnet_type} controlnet...")
88
- self.controlnet_type = controlnet_type
89
  self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"],
90
  torch_dtype=dtype).to(device)
91
  self.pipe.controlnet = self.controlnet
74
  :param data: A dictionary contains `inputs` and optional `image` field.
75
  :return: A dictionary with `image` field contains image in base64.
76
  """
77
+ prompt = data.pop("inputs", None)
78
  image = data.pop("image", None)
79
  controlnet_type = data.pop("controlnet_type", None)
80
 
81
  # Check if neither prompt nor image is provided
82
+ if prompt is None and image is None:
83
  return {"error": "Please provide a prompt and base64 encoded image."}
84
 
85
  # Check if a new controlnet is provided
86
  if controlnet_type is not None and controlnet_type != self.control_type:
87
+ print(f"changing controlnet from {self.control_type} to {controlnet_type} using {CONTROLNET_MAPPING[controlnet_type]['model_id']} model")
88
+ self.control_type = controlnet_type
89
  self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"],
90
  torch_dtype=dtype).to(device)
91
  self.pipe.controlnet = self.controlnet