Checkpoint key example: text_model.embeddings.token_embedding.weight Expected token_embedding sum: -4.9096 Actual token_embedding sum: -0.0497 # ← random init, not checkpoint value

AssertionError: Weights were NOT loaded! expected=-4.9096, got=-0.0497

TL;DR

The issue can be fixed by updating the checkpoint loading logic to handle the changed model structure after refactoring.

Guidance

• The root cause is the mismatch between the old checkpoint keys (e.g., text_model.embeddings.token_embedding.weight) and the new model's state dict keys (e.g., embeddings.token_embedding.weight).
• To verify the fix, compare the expected and actual sums of the token_embedding.weight tensor after loading the checkpoint.
• A potential workaround is to manually update the checkpoint keys to match the new model structure before loading the checkpoint.
• Consider adding a compatibility layer to handle old checkpoints and ensure a smooth transition to the new model structure.

Example

# Manually update the checkpoint keys
updated_sd = {}
for key, value in sd.items():
    if key.startswith("text_model."):
        updated_sd[key[11:]] = value
    else:
        updated_sd[key] = value

# Load the updated checkpoint
te = CLIPTextModel.from_pretrained(
    model_path, subfolder="text_encoder", state_dict=updated_sd
)

Notes

This fix assumes that the only change is the removal of the self.text_model and self.vision_model wrappers. If there are other structural changes, additional updates may be necessary.

Recommendation

Apply a workaround by manually updating the checkpoint keys or adding a compatibility layer to handle old checkpoints, as upgrading to a fixed version is not mentioned in the issue. This approach ensures that existing checkpoints can still be loaded correctly.