does not work for custom tokens/embeddings

[kaibioinfo]

Whenever I want to look at the heatmap for custom tokens (trained with textual inversion) no heatmap is shown for them.

I wanted to use this extension to check if the custom token solely refers to a subject or is somehow leaking information into the background or style.

[torridgristle]

I think it might be as simple as changing all the tokenizers to shared.sd_model.cond_stage_model.wrapped.tokenizer instead of the PreTrainedTokenizer from transformers and adding from modules import shared.