The text_encoder supports inference using onnx & axmodel

.gitattributes

@@ -43,3 +43,5 @@ models/unet.axmodel filter=lfs diff=lfs merge=lfs -text
 models/vae_decoder.axmodel filter=lfs diff=lfs merge=lfs -text
 models/vae_encoder.axmodel filter=lfs diff=lfs merge=lfs -text
 models/7ffcf62c-d292-11ef-bb2a-9d527016cd35 filter=lfs diff=lfs merge=lfs -text
+models/text_encoder/sd15_text_encoder_sim.onnx filter=lfs diff=lfs merge=lfs -text
+models/text_encoder/sd15_text_encoder_sim.axmodel filter=lfs diff=lfs merge=lfs -text

models/text_encoder/sd15_text_encoder_sim.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82cde36ef8130294b1e908a60ba75a266e5543b7f606ba73998f3f7ee5cac243
+size 175200001

run_img2img_axe_infer.py

@@ -5,7 +5,7 @@ import axengine
 import torch
 from PIL import Image
 from transformers import CLIPTokenizer, CLIPTextModel, PreTrainedTokenizer, CLIPTextModelWithProjection
-
+import os
 import time
 import argparse
 from diffusers.utils import load_image
@@ -353,9 +353,7 @@ def _maybe_convert_prompt(prompt: str, tokenizer: "PreTrainedTokenizer"):  # noq
 
 def get_embeds(prompt = "Portrait of a pretty girl", tokenizer_dir = "./models/tokenizer", text_encoder_dir = "./models/text_encoder"):
     tokenizer = CLIPTokenizer.from_pretrained(tokenizer_dir)
-    text_encoder = CLIPTextModel.from_pretrained(text_encoder_dir,
-                                                 torch_dtype=torch.float32,
-                                                 variant="fp16")
+
     text_inputs = tokenizer(
         prompt,
         padding="max_length",
@@ -364,9 +362,16 @@ def get_embeds(prompt = "Portrait of a pretty girl", tokenizer_dir = "./models/t
         return_tensors="pt",
     )
     text_input_ids = text_inputs.input_ids
-    prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=None)
 
-    prompt_embeds_npy = prompt_embeds[0].detach().numpy()
+    text_encoder = axengine.InferenceSession(
+        os.path.join(
+            text_encoder_dir,
+            "sd15_text_encoder_sim.axmodel"
+        ),
+    )
+    text_encoder_onnx_out = text_encoder.run(None, {"input_ids": text_input_ids.to("cpu").numpy().astype(np.int32)})[0]
+
+    prompt_embeds_npy = text_encoder_onnx_out
     return prompt_embeds_npy
 
 

run_img2img_onnx_infer.py

@@ -7,6 +7,7 @@ from transformers import CLIPTokenizer, CLIPTextModel, PreTrainedTokenizer, CLIP
 
 # import axengine as axe
 import time
+import os
 import argparse
 from diffusers.utils import load_image
 import PIL.Image
@@ -353,9 +354,7 @@ def _maybe_convert_prompt(prompt: str, tokenizer: "PreTrainedTokenizer"):  # noq
 
 def get_embeds(prompt = "Portrait of a pretty girl", tokenizer_dir = "./models/tokenizer", text_encoder_dir = "./models/text_encoder"):
     tokenizer = CLIPTokenizer.from_pretrained(tokenizer_dir)
-    text_encoder = CLIPTextModel.from_pretrained(text_encoder_dir,
-                                                 torch_dtype=torch.float32,
-                                                 variant="fp16")
+
     text_inputs = tokenizer(
         prompt,
         padding="max_length",
@@ -364,9 +363,17 @@ def get_embeds(prompt = "Portrait of a pretty girl", tokenizer_dir = "./models/t
         return_tensors="pt",
     )
     text_input_ids = text_inputs.input_ids
-    prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=None)
 
-    prompt_embeds_npy = prompt_embeds[0].detach().numpy()
+    text_encoder = onnxruntime.InferenceSession(
+        os.path.join(
+            text_encoder_dir,
+            "sd15_text_encoder_sim.onnx"
+        ),
+        providers=["CPUExecutionProvider"]
+    )
+    text_encoder_onnx_out = text_encoder.run(None, {"input_ids": text_input_ids.to("cpu").numpy()})[0]
+
+    prompt_embeds_npy = text_encoder_onnx_out
     return prompt_embeds_npy
 
 
@@ -431,9 +438,9 @@ if __name__ == '__main__':
 
     # load unet model and vae model
     start = time.time()
-    vae_encoder = onnxruntime.InferenceSession(vae_encoder_model)
-    unet_session_main = onnxruntime.InferenceSession(unet_model)
-    vae_decoder = onnxruntime.InferenceSession(vae_decoder_model)
+    vae_encoder = onnxruntime.InferenceSession(vae_encoder_model, providers=["CPUExecutionProvider"])
+    unet_session_main = onnxruntime.InferenceSession(unet_model, providers=["CPUExecutionProvider"])
+    vae_decoder = onnxruntime.InferenceSession(vae_decoder_model, providers=["CPUExecutionProvider"])
     print(f"load models take {(1000 * (time.time() - start)):.1f}ms")
 
     # load time input file

run_txt2img_axe_infer.py

@@ -5,7 +5,7 @@ import axengine
 import torch
 from PIL import Image
 from transformers import CLIPTokenizer, CLIPTextModel, PreTrainedTokenizer, CLIPTextModelWithProjection
-
+import os
 import time
 import argparse
 
@@ -55,9 +55,7 @@ def _maybe_convert_prompt(prompt: str, tokenizer: "PreTrainedTokenizer"):  # noq
 
 def get_embeds(prompt = "Portrait of a pretty girl", tokenizer_dir = "./models/tokenizer", text_encoder_dir = "./models/text_encoder"):
     tokenizer = CLIPTokenizer.from_pretrained(tokenizer_dir)
-    text_encoder = CLIPTextModel.from_pretrained(text_encoder_dir,
-                                                 torch_dtype=torch.float32,
-                                                 variant="fp16")
+
     text_inputs = tokenizer(
         prompt,
         padding="max_length",
@@ -66,9 +64,16 @@ def get_embeds(prompt = "Portrait of a pretty girl", tokenizer_dir = "./models/t
         return_tensors="pt",
     )
     text_input_ids = text_inputs.input_ids
-    prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=None)
 
-    prompt_embeds_npy = prompt_embeds[0].detach().numpy()
+    text_encoder = axengine.InferenceSession(
+        os.path.join(
+            text_encoder_dir,
+            "sd15_text_encoder_sim.axmodel"
+        ),
+    )
+    text_encoder_onnx_out = text_encoder.run(None, {"input_ids": text_input_ids.to("cpu").numpy().astype(np.int32)})[0]
+
+    prompt_embeds_npy = text_encoder_onnx_out
     return prompt_embeds_npy
 
 

run_txt2img_axe_infer_new.py

@@ -57,9 +57,7 @@ def _maybe_convert_prompt(prompt: str, tokenizer: "PreTrainedTokenizer"):  # noq
 
 def get_embeds(prompt = "Portrait of a pretty girl", tokenizer_dir = "./models/tokenizer", text_encoder_dir = "./models/text_encoder"):
     tokenizer = CLIPTokenizer.from_pretrained(tokenizer_dir)
-    text_encoder = CLIPTextModel.from_pretrained(text_encoder_dir,
-                                                 torch_dtype=torch.float32,
-                                                 variant="fp16")
+
     text_inputs = tokenizer(
         prompt,
         padding="max_length",
@@ -68,9 +66,16 @@ def get_embeds(prompt = "Portrait of a pretty girl", tokenizer_dir = "./models/t
         return_tensors="pt",
     )
     text_input_ids = text_inputs.input_ids
-    prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=None)
 
-    prompt_embeds_npy = prompt_embeds[0].detach().numpy()
+    text_encoder = axengine.InferenceSession(
+        os.path.join(
+            text_encoder_dir,
+            "sd15_text_encoder_sim.axmodel"
+        ),
+    )
+    text_encoder_onnx_out = text_encoder.run(None, {"input_ids": text_input_ids.to("cpu").numpy().astype(np.int32)})[0]
+
+    prompt_embeds_npy = text_encoder_onnx_out
     return prompt_embeds_npy
 
 

run_txt2img_onnx_infer.py

@@ -5,7 +5,7 @@ import onnxruntime
 import torch
 from PIL import Image
 from transformers import CLIPTokenizer, CLIPTextModel, PreTrainedTokenizer, CLIPTextModelWithProjection
-
+import os
 import time
 import argparse
 
@@ -55,9 +55,7 @@ def _maybe_convert_prompt(prompt: str, tokenizer: "PreTrainedTokenizer"):  # noq
 
 def get_embeds(prompt = "Portrait of a pretty girl", tokenizer_dir = "./models/tokenizer", text_encoder_dir = "./models/text_encoder"):
     tokenizer = CLIPTokenizer.from_pretrained(tokenizer_dir)
-    text_encoder = CLIPTextModel.from_pretrained(text_encoder_dir,
-                                                 torch_dtype=torch.float32,
-                                                 variant="fp16")
+
     text_inputs = tokenizer(
         prompt,
         padding="max_length",
@@ -66,9 +64,17 @@ def get_embeds(prompt = "Portrait of a pretty girl", tokenizer_dir = "./models/t
         return_tensors="pt",
     )
     text_input_ids = text_inputs.input_ids
-    prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=None)
 
-    prompt_embeds_npy = prompt_embeds[0].detach().numpy()
+    text_encoder = onnxruntime.InferenceSession(
+        os.path.join(
+            text_encoder_dir,
+            "sd15_text_encoder_sim.onnx"
+        ),
+        providers=["CPUExecutionProvider"]
+    )
+    text_encoder_onnx_out = text_encoder.run(None, {"input_ids": text_input_ids.to("cpu").numpy()})[0]
+
+    prompt_embeds_npy = text_encoder_onnx_out
     return prompt_embeds_npy
 
 
@@ -116,8 +122,8 @@ if __name__ == '__main__':
     
     # load unet model and vae model
     start = time.time()    
-    unet_session_main = onnxruntime.InferenceSession(unet_model)
-    vae_decoder = onnxruntime.InferenceSession(vae_decoder_model)
+    unet_session_main = onnxruntime.InferenceSession(unet_model, providers=["CPUExecutionProvider"])
+    vae_decoder = onnxruntime.InferenceSession(vae_decoder_model, providers=["CPUExecutionProvider"])
     print(f"load models take {(1000 * (time.time() - start)):.1f}ms")
     
     # load time input file