Add embeddings

Author: kornesh

README.md

@@ -44,6 +44,10 @@ envsubst < script/download-weights > $MODEL_NAME/script/download-weights
 sudo chmod +x $MODEL_NAME/script/download-weights
 envsubst < cog.yaml > $MODEL_NAME/cog.yaml
 
+# download negative embeddings
+mkdir embeddings/ && cd embeddings/ && wget -O FastNegativeV2.pt "https://civitai.com/api/download/models/94057?type=Model&format=PickleTensor" && wget -O BadDream.pt "https://civitai.com/api/download/models/77169?type=Model&format=PickleTensor" && wget -O UnrealisticDream.pt "https://civitai.com/api/download/models/77173?type=Model&format=PickleTensor" && wget -O EasyNegative.pt "https://civitai.com/api/download/models/9536?type=Model&format=PickleTensor&size=full&fp=fp16" && wget -O ng_deepnegative_v1_75t.pt "https://civitai.com/api/download/models/5637?type=Model&format=PickleTensor&size=full&fp=fp16" && wget -O negative_hand-neg.pt "https://civitai.com/api/download/models/60938?type=Negative&format=Other" && wget -O realisticvision-negative-embedding.pt "https://civitai.com/api/download/models/42247?type=Model&format=Other"
+
+
 # cog login
 cd $MODEL_NAME
 git clone https://huggingface.co/pagebrain/$MODEL_NAME model

predict.py

@@ -42,9 +42,11 @@ def from_config(config):
     "KLMS": LMSDiscreteScheduler,
 }
 
+EMBEDDINGS = [(x.split(".")[0], "/embeddings/"+x) for x in os.listdir("/embeddings/")]
+EMBEDDING_TOKENS = [x[0] for x in EMBEDDINGS]
+EMBEDDING_PATHS = [x[1] for x in EMBEDDINGS]
 
-
-MODEL_ID = "pagebrain/$MODEL_NAME"
+MODEL_ID = "pagebrain/realistic-vision-v5-1"
 MODEL_CACHE = "diffusers-cache"
 SAFETY_MODEL_ID = "CompVis/stable-diffusion-safety-checker"
 
@@ -67,23 +69,24 @@ def run_safety_checker(self, image):
     def setup(self):
         """Load the model into memory to make running multiple predictions efficient"""
         self.safety_checker = StableDiffusionSafetyChecker.from_pretrained(
-            "./model/safety_checker",
+            "/model/safety_checker",
             # cache_dir=MODEL_CACHE,
             local_files_only=True,
             torch_dtype=torch.float16,
         ).to("cuda")
 
-        self.feature_extractor = CLIPImageProcessor.from_pretrained("./model/feature_extractor")
+        self.feature_extractor = CLIPImageProcessor.from_pretrained("/model/feature_extractor")
 
         print("Loading txt2img pipeline...")
         self.txt2img_pipe = StableDiffusionPipeline.from_pretrained(
-            "./model/",
+            "/model/",
             # cache_dir=MODEL_CACHE,
             local_files_only=True,
             use_safetensors=True,
             safety_checker=None,
             requires_safety_checker=False
         )
+        self.txt2img_pipe.load_textual_inversion(EMBEDDING_PATHS, token=EMBEDDING_TOKENS, local_files_only=True)
         self.txt2img_pipe.to("cuda")
 
         print("Loading img2img pipeline...")
@@ -124,7 +127,7 @@ def predict(
             default=None,
         ),
         negative_prompt: str = Input(
-            description="Specify things to not see in the output",
+            description="Specify things to not see in the output. Supported embeddings: " + ", ".join(EMBEDDING_TOKENS),
             default=None,
         ),
         image: Path = Input(

script/download-weights

@@ -8,6 +8,11 @@ from diffusers.pipelines.stable_diffusion.safety_checker import (
 )
 import torch
 
+EMBEDDINGS = [(x.split(".")[0], "/embeddings/"+x) for x in os.listdir("/embeddings/")]
+EMBEDDING_TOKENS = [x[0] for x in EMBEDDINGS]
+EMBEDDING_PATHS = [x[1] for x in EMBEDDINGS]
+print(EMBEDDING_TOKENS, EMBEDDING_PATHS)
+
 MODEL_ID = "pagebrain/$MODEL_NAME"
 MODEL_CACHE = "diffusers-cache"
 
@@ -16,15 +21,17 @@ MODEL_CACHE = "diffusers-cache"
 # os.makedirs(MODEL_CACHE, exist_ok=True)
 
 saftey_checker = StableDiffusionSafetyChecker.from_pretrained(
-    "./model/safety_checker",
+    "/model/safety_checker",
     # cache_dir=MODEL_CACHE,
     # use_safetensors=True
     torch_dtype=torch.float16,
 )
 
 pipe = StableDiffusionPipeline.from_pretrained(
-    "./model",
+    "/model",
     # cache_dir=MODEL_CACHE,
     use_safetensors=True,
     safety_checker=None
 )
+
+pipe.load_textual_inversion(EMBEDDING_PATHS, token=EMBEDDING_TOKENS, local_files_only=True)