(p_head_path, w_head_path)
| 156 | |
| 157 | |
| 158 | def convert_safety_checker(p_head_path, w_head_path): |
| 159 | state_dict = {} |
| 160 | |
| 161 | # p head |
| 162 | |
| 163 | p_head = np.load(p_head_path) |
| 164 | |
| 165 | p_head_weights = p_head["weights"] |
| 166 | p_head_weights = torch.from_numpy(p_head_weights) |
| 167 | p_head_weights = p_head_weights.unsqueeze(0) |
| 168 | |
| 169 | p_head_biases = p_head["biases"] |
| 170 | p_head_biases = torch.from_numpy(p_head_biases) |
| 171 | p_head_biases = p_head_biases.unsqueeze(0) |
| 172 | |
| 173 | state_dict["p_head.weight"] = p_head_weights |
| 174 | state_dict["p_head.bias"] = p_head_biases |
| 175 | |
| 176 | # w head |
| 177 | |
| 178 | w_head = np.load(w_head_path) |
| 179 | |
| 180 | w_head_weights = w_head["weights"] |
| 181 | w_head_weights = torch.from_numpy(w_head_weights) |
| 182 | w_head_weights = w_head_weights.unsqueeze(0) |
| 183 | |
| 184 | w_head_biases = w_head["biases"] |
| 185 | w_head_biases = torch.from_numpy(w_head_biases) |
| 186 | w_head_biases = w_head_biases.unsqueeze(0) |
| 187 | |
| 188 | state_dict["w_head.weight"] = w_head_weights |
| 189 | state_dict["w_head.bias"] = w_head_biases |
| 190 | |
| 191 | # vision model |
| 192 | |
| 193 | vision_model = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14") |
| 194 | vision_model_state_dict = vision_model.state_dict() |
| 195 | |
| 196 | for key, value in vision_model_state_dict.items(): |
| 197 | key = f"vision_model.{key}" |
| 198 | state_dict[key] = value |
| 199 | |
| 200 | # full model |
| 201 | |
| 202 | config = CLIPConfig.from_pretrained("openai/clip-vit-large-patch14") |
| 203 | safety_checker = IFSafetyChecker(config) |
| 204 | |
| 205 | safety_checker.load_state_dict(state_dict) |
| 206 | |
| 207 | return safety_checker |
| 208 | |
| 209 | |
| 210 | def create_unet_diffusers_config(original_unet_config, class_embed_type=None): |
no test coverage detected