tests

app.py

@@ -1,25 +1,55 @@
+import pathlib
 import os
 
 import gradio as gr
 import torch
 from diffusers import StableDiffusionPipeline
 
-use_auth_token = os.environ["HF_AUTH_TOKEN"]
-
-# pipeline = StableDiffusionPipeline.from_pretrained(
-#     pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4",
-#     use_auth_token=use_auth_token,
-#     revision="fp16",
-#     torch_dtype=torch.float16,
-# ).to("cuda")
-
-concept_to_dummy_tokens_map = torch.load("concept_to_dummy_tokens_map.pt")
+import utils
 
+use_auth_token = os.environ["HF_AUTH_TOKEN"]
 
-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
+# Instantiate the pipeline.
+device, revision, torch_dtype = (
+    ("cuda", "fp16", torch.float16)
+    if torch.cuda.is_available()
+    else ("cpu", "main", torch.float32)
+)
+pipeline = StableDiffusionPipeline.from_pretrained(
+    pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4",
+    use_auth_token=use_auth_token,
+    revision=revision,
+    torch_dtype=torch_dtype,
+).to(device)
+
+# 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
 
 
 # def inference(

learned_embeddings_dict.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:73ab240e6ef7b16a70e14b4625882d8f63050f1d96ffc0eef6e0e0caa2844109
+size 16235