fixed progress bar update in controlnet-pose, added model caching for other plugins and updated model setup

Author: arisha07

choose_sd_model.py

@@ -14,19 +14,20 @@
 
 
 
-print("=========Chose SD-1.5 models to download and convert=========")
-print("1 - Square (512x512 output image) ")
-print("2 - Landscape (640x360 output image, 16:9) ")	
-print("3 - Portrait (360x640 output image, 9:16) ")
-print("4 - Portrait_512x768 (512x768 output image), will take time since model is large ")
-print("5 - Landscape_768x512 (768x512 output image),  will take time since model is large ")
-print("6 - SD-1.5 Inapinting model (512x512 output image) ")
-print("7 - SD-1.5 Controlnet-Openpose model (512x512 output image) ")
-print("8 - ALL the above SD-1.5 models ")
-print("9 - Only Square, Landscape, Portrait")
-print("10 - Skip All SD-1.5 Model setup ")   
-    
+  
 while True:
+
+    print("=========Chose SD-1.5 models to download and convert=========")
+    print("1 - Square (512x512 output image) ")
+    print("2 - Landscape (640x360 output image, 16:9) ")	
+    print("3 - Portrait (360x640 output image, 9:16) ")
+    print("4 - Portrait_512x768 (512x768 output image), will take time since model is large ")
+    print("5 - Landscape_768x512 (768x512 output image),  will take time since model is large ")
+    print("6 - SD-1.5 Inapinting model (512x512 output image) ")
+    print("7 - SD-1.5 Controlnet-Openpose model (512x512 output image) ")
+    print("8 - ALL the above SD-1.5 models ")
+    print("9 - Only Square, Landscape, Portrait")
+    print("10 - Skip All SD-1.5 Model setup ") 
     choice = input("Enter the Number for the model you want to download & convert: ")
 
     # setup all the SD1.5 models
@@ -53,7 +54,7 @@
 
         subprocess.call([sd_python_path, chose_model_controlnet, choice]) 
 
-        break
+        #break
 
     else:
         print("Wrong option selected.")

gimpopenvino/tools/openvino_common/models_ov/controlnet_openpose.py

@@ -163,27 +163,22 @@ def __init__(
             tokenizer="openai/clip-vit-large-patch14",
             device=["CPU","CPU","CPU"],
             ):
+            
+        super().__init__()    
+            
+        self.set_progress_bar_config(disable=False)    
 
         try:
             self.tokenizer = CLIPTokenizer.from_pretrained(model,local_files_only=True)
         except:
             self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer)
             self.tokenizer.save_pretrained(model)
 
-        # self.swap = swap
-        # models
-        #controlnet = ControlNetModel.from_pretrained("C:\\Users\\lab_admin\\openvino-ai-plugins-gimp\\weights\\stable-diffusion-ov\\controlnet-openpose\\models--lllyasviel--control_v11p_sd15_openpose", torch_dtype=torch.float32) #lllyasviel/control_v11p_sd15_openpose"
-        #pipe = StableDiffusionControlNetPipeline.from_pretrained(
-        #    "runwayml/stable-diffusion-v1-5", controlnet=controlnet
-        #) 
-        
-        #scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-        #scheduler.save_config("C:\\Users\\lab_admin\\openvino-ai-plugins-gimp\\weights\\stable-diffusion-ov\\controlnet-openpose\\scheduler_config")
+
 
         self.scheduler =   UniPCMultistepScheduler.from_pretrained(os.path.join(model,"..","UniPCMultistepScheduler_config"))
 
-        #"C:\\Users\\lab_admin\\openvino-ai-plugins-gimp\\weights\\stable-diffusion-ov\\controlnet-openpose\\scheduler_config")
-        #del pipe
+    
 
         self.core = Core()
         self.core.set_property({'CACHE_DIR': os.path.join(model, 'cache')}) #adding caching to reduce init time
@@ -193,25 +188,24 @@ def __init__(
         OPENPOSE_OV_PATH = os.path.join(model, "openpose.xml")
         self.pose_estimator = OpenposeDetector.from_pretrained(os.path.join(model, "lllyasviel_ControlNet"))
 
-        #"C:\\Users\\lab_admin\\openvino-ai-plugins-gimp\\weights\\stable-diffusion-ov\\controlnet-openpose\\models--lllyasviel--ControlNet") #"lllyasviel/ControlNet")
+    
 
         ov_openpose = OpenPoseOVModel(self.core, OPENPOSE_OV_PATH, device="CPU")
         self.pose_estimator.body_estimation.model = ov_openpose
 
 
 
 
-        #self.vae_scale_factor = 8
-        # self.scheduler = scheduler
+   
         controlnet = os.path.join(model, "controlnet-pose.xml")
         text_encoder = os.path.join(model, "text_encoder.xml")
         unet = os.path.join(model, "unet_controlnet.xml")
-        #unet_neg = os.path.join(model, "unet_controlnet.xml")
+
         vae_decoder = os.path.join(model, "vae_decoder.xml")
 
         ####################
         self.load_models(self.core, device, controlnet, text_encoder, unet, vae_decoder)
-        # self.set_progress_bar_config(disable=True)
+        
 
         # encoder
         self.vae_encoder = None
@@ -221,17 +215,7 @@ def __init__(
         self.height = self.unet.input(0).shape[2] * 8
         self.width = self.unet.input(0).shape[3] * 8    
 
-        #if self.unet.input("sample").shape[1] == 4:
-        #    self.height = self.unet.input("sample").shape[2] * 8
-        #    self.width = self.unet.input("sample").shape[3] * 8
-        #else:
-
-        #    self.height = self.unet.input("sample").shape[1] * 8
-        #    self.width = self.unet.input("sample").shape[2] * 8
-
-        #self.infer_request_neg = self.unet_neg.create_infer_request()
-        #self.infer_request = self.unet.create_infer_request()
-
+ 
     def load_models(self, core: Core, device: str, controlnet:Model, text_encoder: Model, unet: Model, vae_decoder: Model):
         """
         Function for loading models on device using OpenVINO
@@ -352,9 +336,7 @@ def __call__(
                 latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
                 #print("latent_model_input", latent_model_input)
 
-                #print("TTTTTTTTTTTTTTTTTTTTTTTT :", t)
-                #print("text_embeddings:", text_embeddings)
-                #print("pose-------",pose)
+             
 
                 result = self.controlnet([latent_model_input, t, text_embeddings, pose])
                 #print("result", result)
@@ -376,6 +358,10 @@ def __call__(
 
                 if create_gif:
                     frames.append(latents)
+                    
+                # update progress
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()                    
 
         if callback:
               callback(num_inference_steps, callback_userdata)

gimpopenvino/tools/openvino_common/pipelines/async_pipeline.py

@@ -18,6 +18,7 @@
 import threading
 from collections import deque
 from typing import Dict, Set
+import os
 
 
 def parse_devices(device_string):
@@ -81,9 +82,12 @@ def get_user_config(flags_d: str, flags_nstreams: str, flags_nthreads: int)-> Di
 
 
 class AsyncPipeline:
-    def __init__(self, ie, model, plugin_config, device='CPU', max_num_requests=1):
+    def __init__(self, ie, model, model_path, plugin_config, device='CPU', max_num_requests=1):
         self.model = model
         self.logger = logging.getLogger()
+        
+        ie.set_config(config={"CACHE_DIR": os.path.join(model_path, '..', 'cache')}, device_name=device)
+        self.logger.info('Model Cached')        
 
         self.logger.info('Loading network to {} plugin...'.format(device))
         self.exec_net = ie.load_network(network=self.model.net, device_name=device,

gimpopenvino/tools/openvino_common/semseg_run_ov.py

@@ -120,7 +120,7 @@ def run(frame, model_path, device):
     model, visualizer = get_model(ie, model_path)
     log.info('Loading network: %s',model_path )
     log.info('Device: %s',device)   
-    pipeline = AsyncPipeline(ie, model, plugin_config, device, 1)
+    pipeline = AsyncPipeline(ie, model, model_path, plugin_config, device, 1)
     log.info('Starting inference...')
 
     if pipeline.is_ready():

gimpopenvino/tools/openvino_common/styletransfer_run_ov.py

@@ -47,7 +47,7 @@ def run(frame, model_path, device):
      log.info('Device: %s',device)
 
      model = StyleTransfer(ie, model_path, frame.shape)
-     pipeline = AsyncPipeline(ie, model, plugin_config, device, 1)
+     pipeline = AsyncPipeline(ie, model, model_path, plugin_config, device, 1)
 
      log.info('Starting inference...')
 

gimpopenvino/tools/openvino_common/superes_run_ov.py

@@ -45,7 +45,7 @@ def run(frame, model_path, device, model_name):
      log.info('Device: %s',device)
 
      model = SuperResolution(ie, model_path, frame.shape, model_name)
-     pipeline = AsyncPipeline(ie, model, plugin_config, device, 1)
+     pipeline = AsyncPipeline(ie, model, model_path, plugin_config, device, 1)
 
      log.info('Starting inference...')