Use text_encoder from cache

ForserX · Nov 23, 2023 · 2b56ba2 · 2b56ba2
1 parent e79d0b7
commit 2b56ba2
Show file tree

Hide file tree

Showing 7 changed files with 60 additions and 7 deletions.
diff --git a/data/repo/diffusion_scripts/sd_cuda_safe.py b/data/repo/diffusion_scripts/sd_cuda_safe.py
@@ -20,7 +20,7 @@
 
 PipeDevice = Device(opt.device, fptype)
 
-pipe = PipeDevice.GetPipe(opt.mdlpath, opt.mode, opt.nsfw)
+pipe = PipeDevice.GetPipe(opt.mdlpath, opt.mode, opt.nsfw, "")
 pipe.to(PipeDevice.device, fptype)
 
 if opt.dlora:

diff --git a/data/repo/diffusion_scripts/sd_onnx_safe.py b/data/repo/diffusion_scripts/sd_onnx_safe.py
@@ -13,7 +13,7 @@
 pipe = None
 PipeDevice = Device("onnx", torch.float32)
 
-pipe = PipeDevice.GetPipe(opt.mdlpath, opt.mode, opt.nsfw)
+pipe = PipeDevice.GetPipe(opt.mdlpath, opt.mode, opt.nsfw, opt.textencoder)
 safe_unet = pipe.unet
 
 print("SD: Model preload: done")
@@ -43,7 +43,7 @@
     tokenizer_extract = False
 
     if (data['LoRA'] != old_lora_json) or (data['TI'] != old_te_json) or (data['TINeg'] != old_ten_json):
-        onnx_te_model = onnx.load(opt.mdlpath + "/text_encoder/" + ONNX_MODEL)
+        onnx_te_model = onnx.load(opt.textencoder + ONNX_MODEL)
 
         # Hard reload
         old_lora_json = None

diff --git a/data/repo/diffusion_scripts/sd_xbackend.py b/data/repo/diffusion_scripts/sd_xbackend.py
@@ -69,7 +69,7 @@ def LPW_Path(self):
 
 		return dir_path
 
-	def GetPipe(self, Model: str, Mode: str, NSFW: bool):
+	def GetPipe(self, Model: str, Mode: str, NSFW: bool, TextEnc: str):
 		pipe = None
 		nsfw_pipe = None
 
@@ -87,11 +87,13 @@ def GetPipe(self, Model: str, Mode: str, NSFW: bool):
 					nsfw_pipe = OnnxRuntimeModel.from_pretrained(safety_model, provider=self.prov)
 				print (Mode)    
 				if Mode == "txt2img":
-					pipe = OnnxStableDiffusionPipeline.from_pretrained(Model, custom_pipeline=self.LPW_Path(), provider=self.prov, safety_checker=nsfw_pipe)
+					pipe = OnnxStableDiffusionPipeline.from_pretrained(Model, custom_pipeline=self.LPW_Path(), provider=self.prov, safety_checker=nsfw_pipe, text_encoder=None)
 				if Mode == "img2img":
-					pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(Model, custom_pipeline=self.LPW_Path(), provider=self.prov, safety_checker=nsfw_pipe)
+					pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(Model, custom_pipeline=self.LPW_Path(), provider=self.prov, safety_checker=nsfw_pipe, text_encoder=None)
 				if Mode == "inpaint":
-					pipe = OnnxStableDiffusionInpaintPipeline.from_pretrained(Model, custom_pipeline=self.LPW_Path(), provider=self.prov, safety_checker=nsfw_pipe)
+					pipe = OnnxStableDiffusionInpaintPipeline.from_pretrained(Model, custom_pipeline=self.LPW_Path(), provider=self.prov, safety_checker=nsfw_pipe, text_encoder=None)
+
+				pipe.text_encoder = OnnxRuntimeModel.from_pretrained(TextEnc, provider=self.prov)
 		else:
 			if Mode == "pix2pix":
 				if NSFW:
@@ -399,6 +401,9 @@ def ApplyArg(parser):
 		parser.add_argument(
 			"--model", type=str, help="Path to model checkpoint (.ckpt or .safetensors)", dest='mdlpath',
 		)
+		parser.add_argument(
+			"--textencoder", type=str, help="Path to model checkpoint (.onnx)", dest='textencoder',
+		)
 		parser.add_argument(
 			"--workdir", default=None, type=str, help="Path to working directory", dest='workdir',
 		)

diff --git a/ui-src/MainWindow.xaml.cs b/ui-src/MainWindow.xaml.cs
@@ -223,9 +223,13 @@ private void Button_ClickBreak(object sender, RoutedEventArgs e)
             if (!Directory.Exists(WorkingPath))
             {
                 if (GlobalVariables.Mode == Helper.ImplementMode.DiffCUDA)
+                {
                     Install.CheckAndInstallCUDA();
+                }
                 else
+                {
                     Install.CheckAndInstallONNX();
+                }
 
                 return;
             }
@@ -264,6 +268,7 @@ private void btnONNX_Click(object sender, RoutedEventArgs e)
 
                 GlobalVariables.Mode = Helper.ImplementMode.ONNX;
                 Install.CheckAndInstallONNX();
+                Task.Run(() => TextEncoderONNX.CheckEncoder());
 
                 Brush Safe = new SolidColorBrush(Colors.Black);
 

diff --git a/ui-src/Utils/ModelCMD.cs b/ui-src/Utils/ModelCMD.cs
@@ -1,8 +1,10 @@
 using Newtonsoft.Json;
+using SD_FXUI.Utils.Models;
 using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Text;
+using System.Text.Encodings.Web;
 using System.Threading.Tasks;
 using System.Windows.Xps.Serialization;
 
@@ -73,6 +75,7 @@ public void PreStart(string StartModel, string StartMode, bool StartNSFW, bool I
                     {
                         CmdLine += " --nsfw=True ";
                     }
+                    CmdLine += $" --textencoder={TextEncoderONNX.GetModelDir()} ";
 
                     Process = new Host(FS.GetWorkingDir(), "repo/" + PythonEnv.GetPy(Helper.VENV.DiffONNX));
                     Process.Start("./repo/diffusion_scripts/sd_onnx_safe.py " + CmdLine);

diff --git a/ui-src/Utils/Models/TextEncoderONNX.cs b/ui-src/Utils/Models/TextEncoderONNX.cs
@@ -0,0 +1,28 @@
+using System.IO;
+
+namespace SD_FXUI.Utils.Models
+{
+    internal class TextEncoderONNX
+    {
+        static public void CheckEncoder()
+        {
+            if (GlobalVariables.Mode == Helper.ImplementMode.ONNX)
+            {
+                if (!File.Exists(GetModel()))
+                {
+                    WGetDownloadModels.DownloadTextEncoder();
+                }
+            }
+        }
+
+        static public string GetModel()
+        {
+            return FS.GetModelDir() + "text_encoder/model.onnx";
+        }
+        static public string GetModelDir()
+        {
+            return FS.GetModelDir() + "text_encoder/";
+        }
+
+    }
+}
diff --git a/ui-src/Utils/Models/WGetDownloadModels.cs b/ui-src/Utils/Models/WGetDownloadModels.cs
@@ -50,6 +50,18 @@ public static void DownloadSDPoser()
             FileDownloader.DownloadFileAsync("https://huggingface.co/lllyasviel/control_v11p_sd15_openpose/raw/main/config.json", WorkingDir + @"config.json");
             FileDownloader.DownloadFileAsync("https://huggingface.co/lllyasviel/control_v11p_sd15_openpose/resolve/main/diffusion_pytorch_model.fp16.bin", WorkingDir + @"diffusion_pytorch_model.bin");
         }
+        public static void DownloadTextEncoder()
+        {
+            string WorkingDir = FS.GetModelDir() + "text_encoder/";
+
+            if (Directory.Exists(WorkingDir))
+            {
+                return;
+            }
+            Directory.CreateDirectory(WorkingDir);
+
+            FileDownloader.DownloadFileAsync("https://huggingface.co/ForserX/TextEncoderBackupONNX/resolve/main/model.onnx", WorkingDir + @"model.onnx");
+        }
 
         public static void DownloadSDFacegen()
         {