diff options
Diffstat (limited to 'infer.py')
-rw-r--r-- | infer.py | 6 |
1 files changed, 3 insertions, 3 deletions
@@ -28,7 +28,7 @@ from transformers import CLIPTextModel | |||
28 | from models.clip.embeddings import patch_managed_embeddings | 28 | from models.clip.embeddings import patch_managed_embeddings |
29 | from models.clip.tokenizer import MultiCLIPTokenizer | 29 | from models.clip.tokenizer import MultiCLIPTokenizer |
30 | from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion | 30 | from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion |
31 | from common import load_config, load_embeddings_from_dir | 31 | from util import load_config, load_embeddings_from_dir |
32 | 32 | ||
33 | 33 | ||
34 | torch.backends.cuda.matmul.allow_tf32 = True | 34 | torch.backends.cuda.matmul.allow_tf32 = True |
@@ -192,12 +192,12 @@ def save_args(basepath, args, extra={}): | |||
192 | 192 | ||
193 | 193 | ||
194 | def load_embeddings(pipeline, embeddings_dir): | 194 | def load_embeddings(pipeline, embeddings_dir): |
195 | added_tokens = load_embeddings_from_dir( | 195 | added_tokens, added_ids = load_embeddings_from_dir( |
196 | pipeline.tokenizer, | 196 | pipeline.tokenizer, |
197 | pipeline.text_encoder.text_model.embeddings, | 197 | pipeline.text_encoder.text_model.embeddings, |
198 | Path(embeddings_dir) | 198 | Path(embeddings_dir) |
199 | ) | 199 | ) |
200 | print(f"Added {len(added_tokens)} tokens from embeddings dir: {added_tokens}") | 200 | print(f"Added {len(added_tokens)} tokens from embeddings dir: {zip(added_tokens, added_ids)}") |
201 | 201 | ||
202 | 202 | ||
203 | def create_pipeline(model, dtype): | 203 | def create_pipeline(model, dtype): |