Update app.py

app.py

@@ -33,22 +33,43 @@ pipe_i2i=pipe_i2i_merged.to("cuda")
 pipe.enable_xformers_memory_efficient_attention()
 pipe_i2i.enable_xformers_memory_efficient_attention()
 
-nfixer_path=hf_hub_download(repo_id=model_id, filename="nfixer.pt", use_auth_token=token)
-embeddings_dict=torch.load(nfixer_path)
+embeddings_path=hf_hub_download(repo_id=model_id, filename="nfixer.pt", use_auth_token=token)
+embeddings_dict=torch.load(embeddings_path)
 print(embeddings_dict)
 if "string_to_param" in embeddings_dict:
     embeddings = next(iter(embeddings_dict['string_to_param'].values()))
-    placeholder_token = ""
+    nfixer = ""
     for i, emb in enumerate(embeddings):
         token = f"_s{i+1}"
-        placeholder_token += token
+        nfixer += token
         pipe.tokenizer.add_tokens(token)
         pipe.text_encoder.resize_token_embeddings(len(pipe.tokenizer))
         token_id = pipe.tokenizer.convert_tokens_to_ids(token)
         pipe.text_encoder.get_input_embeddings().weight.data[token_id] = emb
 else:
-    placeholder_token = list(embeddings_dict.keys())[0]
-    embeddings = embeddings_dict[placeholder_token].to(pipe.text_encoder.get_input_embeddings().weight.dtype)
+    nfixer = list(embeddings_dict.keys())[0]
+    embeddings = embeddings_dict[nfixer].to(pipe.text_encoder.get_input_embeddings().weight.dtype)
+    pipe.tokenizer.add_tokens(placeholder_token)
+    pipe.text_encoder.resize_token_embeddings(len(pipe.tokenizer))
+    placeholder_token_id = pipe.tokenizer.convert_tokens_to_ids(placeholder_token)
+    pipe.text_encoder.get_input_embeddings().weight.data[placeholder_token_id] = embeddings
+
+embeddings_path=hf_hub_download(repo_id=model_id, filename="embellish2.pt", use_auth_token=token)
+embeddings_dict=torch.load(embeddings_path)
+print(embeddings_dict)
+if "string_to_param" in embeddings_dict:
+    embeddings = next(iter(embeddings_dict['string_to_param'].values()))
+    embellish2 = ""
+    for i, emb in enumerate(embeddings):
+        token = f"_s{i+1}"
+        embellish2 += token
+        pipe.tokenizer.add_tokens(token)
+        pipe.text_encoder.resize_token_embeddings(len(pipe.tokenizer))
+        token_id = pipe.tokenizer.convert_tokens_to_ids(token)
+        pipe.text_encoder.get_input_embeddings().weight.data[token_id] = emb
+else:
+    embellish2 = list(embeddings_dict.keys())[0]
+    embeddings = embeddings_dict[embellish2].to(pipe.text_encoder.get_input_embeddings().weight.dtype)
     pipe.tokenizer.add_tokens(placeholder_token)
     pipe.text_encoder.resize_token_embeddings(len(pipe.tokenizer))
     placeholder_token_id = pipe.tokenizer.convert_tokens_to_ids(placeholder_token)
@@ -98,8 +119,8 @@ def auto_prompt_correction(prompt_ui,neg_prompt_ui,disable_auto_prompt_correctio
         return prompt, neg_prompt
 
     if(prompt=="" and neg_prompt==""):
-        prompt="anime, masterpiece, portrait, a girl with flowers, good pupil, 4k, detailed"
-        neg_prompt=f"{placeholder_token},(((deformed))), blurry, ((((bad anatomy)))), bad pupil, disfigured, poorly drawn face, mutation, mutated, (extra limb), (ugly), (poorly drawn hands), bad hands, fused fingers, messy drawing, broken legs censor, low quality, (mutated hands and fingers:1.5), (long body :1.3), (mutation, poorly drawn :1.2), ((bad eyes)), ui, error, missing fingers, fused fingers, one hand with more than 5 fingers, one hand with less than 5 fingers, one hand with more than 5 digit, one hand with less than 5 digit, extra digit, fewer digits, fused digit, missing digit, bad digit, liquid digit, long body, uncoordinated body, unnatural body, lowres, jpeg artifacts, 3d, cg, text"
+        prompt=f"{embellish2},anime, masterpiece, portrait, a girl with flowers, good pupil, 4k, detailed"
+        neg_prompt=f"{nfixer},(((deformed))), blurry, ((((bad anatomy)))), bad pupil, disfigured, poorly drawn face, mutation, mutated, (extra limb), (ugly), (poorly drawn hands), bad hands, fused fingers, messy drawing, broken legs censor, low quality, (mutated hands and fingers:1.5), (long body :1.3), (mutation, poorly drawn :1.2), ((bad eyes)), ui, error, missing fingers, fused fingers, one hand with more than 5 fingers, one hand with less than 5 fingers, one hand with more than 5 digit, one hand with less than 5 digit, extra digit, fewer digits, fused digit, missing digit, bad digit, liquid digit, long body, uncoordinated body, unnatural body, lowres, jpeg artifacts, 3d, cg, text"
         return prompt, neg_prompt
 
     splited_prompt=prompt.replace(","," ").replace("_"," ").split(" ")