Skip to content

Add plain PyTorch loader for the CSD checkpoint - #15

Open
fahmeed-nabi wants to merge 1 commit into
learn2phoenix:mainfrom
fahmeed-nabi:add-pytorch-loader
Open

Add plain PyTorch loader for the CSD checkpoint#15
fahmeed-nabi wants to merge 1 commit into
learn2phoenix:mainfrom
fahmeed-nabi:add-pytorch-loader

Conversation

@fahmeed-nabi

Copy link
Copy Markdown

Summary

Adds a plain PyTorch loader for the released CSD checkpoint, answering #8 ("How to load the model as a PyTorch module?").

The checkpoint has no transformers AutoModel support and config.json carries no architecture info ({"model_type": "custom"}), so it can't be loaded with from_pretrained. models/csd_clip.py reconstructs the architecture directly from the checkpoint's state_dict key names and shapes:

  • The backbone is a stock CLIP ViT-L/14 (quickgelu) visual transformer, matching OpenAI's original CLIP activation since training started from OpenAI's CLIP weights rather than open_clip's plain-GELU variant.
  • The backbone's own contrastive proj is unused (no proj key in the checkpoint), and the raw pooled feature feeds forward instead.
  • last_layer_style / last_layer_content are separate (1024, 768) matrices applied via matrix multiply, not nn.Linear (the checkpoint stores them as raw (in, out) matrices, not nn.Linear.weight's (out, in) shape).

Loading with strict=True against this class produces zero missing and zero unexpected keys against the released weights.

Also adds a short README section with a runnable example (load the checkpoint, embed an image, get back style/content embeddings).

Why

This has come up a few times (#8, and implicitly in how people are using the HF checkpoint) with no canonical answer in the repo. This closes that gap with a minimal, dependency-light loader that doesn't require guessing the architecture from the state_dict, which is what we had to do to use CSD in our own project.

Testing

  • Verified strict=True state_dict load produces zero missing / zero unexpected keys against tomg-group-umd/CSD-ViT-L.
  • Ran the README example end-to-end: loads the checkpoint, embeds an image, produces style and content embeddings of the expected shape.

Closes #8.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

How to load the model as a PyTorch module?

2 participants