Commit
·
75c898e
1
Parent(s):
4933de5
init
Browse files- handler.py +99 -100
handler.py
CHANGED
@@ -54,116 +54,115 @@ class EndpointHandler():
|
|
54 |
hf_hub_download(repo_id="InstantX/InstantID", filename="ip-adapter.bin", local_dir="./checkpoints")
|
55 |
|
56 |
print("Model dir: ", model_dir)
|
57 |
-
|
58 |
-
|
59 |
-
|
60 |
-
|
61 |
-
|
62 |
-
|
63 |
-
|
64 |
-
|
65 |
-
|
66 |
-
|
67 |
-
|
68 |
-
|
69 |
-
|
70 |
-
|
71 |
-
|
72 |
-
|
73 |
-
|
74 |
-
|
75 |
-
|
76 |
-
|
77 |
-
|
78 |
-
|
79 |
-
|
80 |
-
|
81 |
-
|
82 |
-
|
83 |
-
|
84 |
-
|
85 |
-
|
86 |
-
|
87 |
-
|
88 |
-
|
89 |
-
|
90 |
-
|
91 |
-
|
92 |
-
|
93 |
-
|
94 |
-
|
95 |
-
|
96 |
-
|
97 |
-
|
98 |
-
|
99 |
-
|
100 |
-
|
101 |
-
|
102 |
-
|
103 |
-
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
|
110 |
-
|
111 |
-
|
112 |
-
|
113 |
-
|
114 |
-
|
115 |
-
|
116 |
-
|
117 |
-
|
118 |
-
|
119 |
-
|
120 |
-
|
121 |
-
|
122 |
-
# return Image.fromarray(edges, "L")
|
123 |
|
124 |
-
|
125 |
|
126 |
-
|
127 |
|
128 |
-
|
129 |
|
130 |
-
|
131 |
-
|
132 |
|
133 |
-
|
134 |
-
|
135 |
|
136 |
-
|
137 |
-
|
138 |
|
139 |
-
|
140 |
|
141 |
-
|
142 |
|
143 |
-
|
144 |
|
145 |
-
|
146 |
-
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
152 |
-
|
153 |
-
|
154 |
-
|
155 |
-
|
156 |
-
|
157 |
-
|
158 |
-
|
159 |
-
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
167 |
|
168 |
def __call__(self, param):
|
169 |
print("Param: ", param)
|
|
|
54 |
hf_hub_download(repo_id="InstantX/InstantID", filename="ip-adapter.bin", local_dir="./checkpoints")
|
55 |
|
56 |
print("Model dir: ", model_dir)
|
57 |
+
face_adapter = f"./checkpoints/ip-adapter.bin"
|
58 |
+
controlnet_path = f"./checkpoints/ControlNetModel"
|
59 |
+
|
60 |
+
transform = Compose([
|
61 |
+
Resize(
|
62 |
+
width=518,
|
63 |
+
height=518,
|
64 |
+
resize_target=False,
|
65 |
+
keep_aspect_ratio=True,
|
66 |
+
ensure_multiple_of=14,
|
67 |
+
resize_method='lower_bound',
|
68 |
+
image_interpolation_method=cv2.INTER_CUBIC,
|
69 |
+
),
|
70 |
+
NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
71 |
+
PrepareForNet(),
|
72 |
+
])
|
73 |
+
|
74 |
+
self.controlnet_identitynet = ControlNetModel.from_pretrained(
|
75 |
+
controlnet_path, torch_dtype=dtype
|
76 |
+
)
|
77 |
+
|
78 |
+
pretrained_model_name_or_path = "wangqixun/YamerMIX_v8"
|
79 |
+
|
80 |
+
self.pipe = StableDiffusionXLInstantIDPipeline.from_pretrained(
|
81 |
+
pretrained_model_name_or_path,
|
82 |
+
controlnet=[self.controlnet_identitynet],
|
83 |
+
torch_dtype=dtype,
|
84 |
+
safety_checker=None,
|
85 |
+
feature_extractor=None,
|
86 |
+
).to(device)
|
87 |
+
|
88 |
+
|
89 |
+
self.pipe.scheduler = diffusers.EulerDiscreteScheduler.from_config(
|
90 |
+
self.pipe.scheduler.config
|
91 |
+
)
|
92 |
+
|
93 |
+
# load and disable LCM
|
94 |
+
self.pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl")
|
95 |
+
self.pipe.disable_lora()
|
96 |
+
|
97 |
+
self.pipe.cuda()
|
98 |
+
self.pipe.load_ip_adapter_instantid(face_adapter)
|
99 |
+
self.pipe.image_proj_model.to("cuda")
|
100 |
+
self.pipe.unet.to("cuda")
|
101 |
+
|
102 |
+
|
103 |
+
# controlnet-pose/canny/depth
|
104 |
+
controlnet_pose_model = "thibaud/controlnet-openpose-sdxl-1.0"
|
105 |
+
controlnet_canny_model = "diffusers/controlnet-canny-sdxl-1.0"
|
106 |
+
controlnet_depth_model = "diffusers/controlnet-depth-sdxl-1.0-small"
|
107 |
+
|
108 |
+
controlnet_pose = ControlNetModel.from_pretrained(
|
109 |
+
controlnet_pose_model, torch_dtype=dtype
|
110 |
+
).to(device)
|
111 |
+
controlnet_canny = ControlNetModel.from_pretrained(
|
112 |
+
controlnet_canny_model, torch_dtype=dtype
|
113 |
+
).to(device)
|
114 |
+
controlnet_depth = ControlNetModel.from_pretrained(
|
115 |
+
controlnet_depth_model, torch_dtype=dtype
|
116 |
+
).to(device)
|
117 |
+
|
118 |
+
def get_canny_image(image, t1=100, t2=200):
|
119 |
+
image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
|
120 |
+
edges = cv2.Canny(image, t1, t2)
|
121 |
+
return Image.fromarray(edges, "L")
|
|
|
122 |
|
123 |
+
def get_depth_map(image):
|
124 |
|
125 |
+
image = np.array(image) / 255.0
|
126 |
|
127 |
+
h, w = image.shape[:2]
|
128 |
|
129 |
+
image = transform({'image': image})['image']
|
130 |
+
image = torch.from_numpy(image).unsqueeze(0).to("cuda")
|
131 |
|
132 |
+
with torch.no_grad():
|
133 |
+
depth = depth_anything(image)
|
134 |
|
135 |
+
depth = F.interpolate(depth[None], (h, w), mode='bilinear', align_corners=False)[0, 0]
|
136 |
+
depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
|
137 |
|
138 |
+
depth = depth.cpu().numpy().astype(np.uint8)
|
139 |
|
140 |
+
depth_image = Image.fromarray(depth)
|
141 |
|
142 |
+
return depth_image
|
143 |
|
144 |
+
self.controlnet_map = {
|
145 |
+
"pose": controlnet_pose,
|
146 |
+
"canny": get_canny_image,
|
147 |
+
"depth": controlnet_depth,
|
148 |
+
}
|
149 |
+
|
150 |
+
openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
|
151 |
+
depth_anything = DepthAnything.from_pretrained('LiheYoung/depth_anything_vitl14').to(device).eval()
|
152 |
+
|
153 |
+
|
154 |
+
self.controlnet_map_fn = {
|
155 |
+
"pose": openpose,
|
156 |
+
"canny": get_canny_image,
|
157 |
+
"depth": get_depth_map,
|
158 |
+
}
|
159 |
+
|
160 |
+
self.app = FaceAnalysis(
|
161 |
+
name="antelopev2",
|
162 |
+
root="./",
|
163 |
+
providers=["CPUExecutionProvider"],
|
164 |
+
)
|
165 |
+
self.app.prepare(ctx_id=0, det_size=(640, 640))
|
166 |
|
167 |
def __call__(self, param):
|
168 |
print("Param: ", param)
|