(state_dict: dict)
| 455 | |
| 456 | |
| 457 | def build_model(state_dict: dict): |
| 458 | vit = "visual.proj" in state_dict |
| 459 | |
| 460 | if vit: |
| 461 | vision_width = state_dict["visual.conv1.weight"].shape[0] |
| 462 | vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")]) |
| 463 | vision_patch_size = state_dict["visual.conv1.weight"].shape[-1] |
| 464 | grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5) |
| 465 | image_resolution = vision_patch_size * grid_size |
| 466 | else: |
| 467 | counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]] |
| 468 | vision_layers = tuple(counts) |
| 469 | vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0] |
| 470 | output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5) |
| 471 | vision_patch_size = None |
| 472 | assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0] |
| 473 | image_resolution = output_width * 32 |
| 474 | |
| 475 | embed_dim = state_dict["text_projection"].shape[1] |
| 476 | context_length = state_dict["positional_embedding"].shape[0] |
| 477 | vocab_size = state_dict["token_embedding.weight"].shape[0] |
| 478 | transformer_width = state_dict["ln_final.weight"].shape[0] |
| 479 | transformer_heads = transformer_width // 64 |
| 480 | transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks"))) |
| 481 | |
| 482 | model = CLIP( |
| 483 | embed_dim, |
| 484 | image_resolution, vision_layers, vision_width, vision_patch_size, |
| 485 | context_length, vocab_size, transformer_width, transformer_heads, transformer_layers |
| 486 | ) |
| 487 | |
| 488 | for key in ["input_resolution", "context_length", "vocab_size"]: |
| 489 | if key in state_dict: |
| 490 | del state_dict[key] |
| 491 | |
| 492 | convert_weights(model) |
| 493 | try: |
| 494 | model.load_state_dict(state_dict) |
| 495 | except: |
| 496 | missing_keys, _ = model.load_state_dict(state_dict, strict=False) |
| 497 | print('Weights not found for some missing keys: ', missing_keys) |
| 498 | return model.eval() |
no test coverage detected