Fix stable diffusion engine initialization

Author: gblong1

gimpopenvino/plugins/openvino_utils/tools/model_manager.py

@@ -515,8 +515,6 @@ def is_model_installed(self, model_id):
                                 "best performance" : ["GPU","GPU","GPU"]
                         }
 
-
-
                 npu_is_available = self._npu_is_available
                 npu_arch = self._npu_arch
 

gimpopenvino/plugins/openvino_utils/tools/openvino_common/models_ov/stable_diffusion_engine.py

@@ -455,15 +455,15 @@ def __init__(
             self,
             model="bes-dev/stable-diffusion-v1-4-openvino",
             tokenizer="openai/clip-vit-large-patch14",
-            device=["CPU","CPU","CPU","CPU"]):
+            device=["CPU","CPU","CPU","CPU"], model_name="fp16"):
 
         self.core = Core()
         self.core.set_property({'CACHE_DIR': os.path.join(model, 'cache')})
 
         batch_size = 2 if device[1] == device[2] and device[1] == "GPU" else 1
 
         # if 'int8' is in model, then we are using unet_int8a16 model, and for this we will always use batch size 1.
-        if "int8" in model:
+        if "int8" in model_name:
             batch_size = 1
 
         self.batch_size = batch_size
@@ -477,21 +477,24 @@ def __init__(
             self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer)
             self.tokenizer.save_pretrained(model)
 
-        print("Loading models... ")
+    
 
         with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
             text_future = executor.submit(self.load_model, model, "text_encoder", device[0])
             vae_de_future = executor.submit(self.load_model, model, "vae_decoder", device[3])
             vae_en_future = executor.submit(self.load_model, model, "vae_encoder", device[3])
 
             if self.batch_size == 1:
-                if "int8" not in model:
-                    unet_future = executor.submit(self.load_model, model, "unet_bs1", device[1])
-                    unet_neg_future = executor.submit(self.load_model, model, "unet_bs1", device[2]) if device[1] != device[2] else None
-                else:
-                    unet_future = executor.submit(self.load_model, model, "unet_int8a16", device[1])
+                if "int8a16" in model_name:
+                    print("Loading models ... int8a16")
+                    unet_future     = executor.submit(self.load_model, model, "unet_int8a16", device[1])
                     unet_neg_future = executor.submit(self.load_model, model, "unet_int8a16", device[2]) if device[1] != device[2] else None
+                else:
+                    print("Loading models ... fp16 bs1")
+                    unet_future     = executor.submit(self.load_model, model, "unet_bs1", device[1])
+                    unet_neg_future = executor.submit(self.load_model, model, "unet_bs1", device[2]) if device[1] != device[2] else None
             else:
+                print("Loading models ... fp16")
                 unet_future = executor.submit(self.load_model, model, "unet", device[1])
                 unet_neg_future = None
 

gimpopenvino/plugins/openvino_utils/tools/stable_diffusion_ov_server.py

@@ -217,7 +217,7 @@ def initialize_engine(model_name, model_path, device_list):
         return controlnet_openpose.ControlNetOpenPose(model=model_path, device=device_list)
     if model_name == "controlnet_referenceonly":
         return stable_diffusion_engine.StableDiffusionEngineReferenceOnly(model=model_path, device=device_list)
-    return stable_diffusion_engine.StableDiffusionEngine(model=model_path, device=device_list)
+    return stable_diffusion_engine.StableDiffusionEngine(model=model_path, device=device_list, model_name=model_name)
 
 def handle_client_data(data, conn, engine, model_name, model_path, scheduler):
     if data.decode() == "kill":