Commit
·
394463f
1
Parent(s):
fdae70a
Update handler.py
Browse files- 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("
|
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
|
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"
|
88 |
-
self.
|
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
|