tests

app.py

@@ -25,31 +25,31 @@ pipeline = StableDiffusionPipeline.from_pretrained(
 # Load in the new concepts.
 CONCEPT_PATH = pathlib.Path("learned_embeddings_dict.pt")
 learned_embeddings_dict = torch.load(CONCEPT_PATH)
-#
-# concept_to_dummy_tokens_map = {}
-# for concept_token, embedding_dict in learned_embeddings_dict.items():
-#     initializer_tokens = embedding_dict["initializer_tokens"]
-#     learned_embeddings = embedding_dict["learned_embeddings"]
-#     (
-#         initializer_ids,
-#         dummy_placeholder_ids,
-#         dummy_placeholder_tokens,
-#     ) = utils.add_new_tokens_to_tokenizer(
-#         concept_token=concept_token,
-#         initializer_tokens=initializer_tokens,
-#         tokenizer=pipeline.tokenizer,
-#     )
-#     pipeline.text_encoder.resize_token_embeddings(len(pipeline.tokenizer))
-#     token_embeddings = pipeline.text_encoder.get_input_embeddings().weight.data
-#     for d_id, tensor in zip(dummy_placeholder_ids, learned_embeddings):
-#         token_embeddings[d_id] = tensor
-#     concept_to_dummy_tokens_map[concept_token] = dummy_placeholder_tokens
-#
-#
-# def replace_concept_tokens(text: str):
-#     for concept_token, dummy_tokens in concept_to_dummy_tokens_map.items():
-#         text = text.replace(concept_token, dummy_tokens)
-#     return text
+
+concept_to_dummy_tokens_map = {}
+for concept_token, embedding_dict in learned_embeddings_dict.items():
+    initializer_tokens = embedding_dict["initializer_tokens"]
+    learned_embeddings = embedding_dict["learned_embeddings"]
+    (
+        initializer_ids,
+        dummy_placeholder_ids,
+        dummy_placeholder_tokens,
+    ) = utils.add_new_tokens_to_tokenizer(
+        concept_token=concept_token,
+        initializer_tokens=initializer_tokens,
+        tokenizer=pipeline.tokenizer,
+    )
+    pipeline.text_encoder.resize_token_embeddings(len(pipeline.tokenizer))
+    token_embeddings = pipeline.text_encoder.get_input_embeddings().weight.data
+    for d_id, tensor in zip(dummy_placeholder_ids, learned_embeddings):
+        token_embeddings[d_id] = tensor
+    concept_to_dummy_tokens_map[concept_token] = dummy_placeholder_tokens
+
+
+def replace_concept_tokens(text: str):
+    for concept_token, dummy_tokens in concept_to_dummy_tokens_map.items():
+        text = text.replace(concept_token, dummy_tokens)
+    return text
 
 
 # def inference(
@@ -101,7 +101,7 @@ with gr.Blocks() as demo:
         interactive=True,
     )
     output = gr.Textbox(
-        label="output", placeholder=use_auth_token[:10], interactive=False
+        label="output", placeholder=use_auth_token[:5], interactive=False
     )
     gr.Button("test").click(
         lambda s: replace_concept_tokens(s), inputs=[prompt], outputs=output