Spaces:
Running
Running
| from jaa import JaaCore | |
| from roop.utilities import get_device | |
| from typing import Any | |
| version = "4.0.0" | |
| class ChainImgProcessor(JaaCore): | |
| def __init__(self): | |
| JaaCore.__init__(self) | |
| self.processors:dict = { | |
| } | |
| self.processors_objects:dict[str,list[ChainImgPlugin]] = {} | |
| self.default_chain = "" | |
| self.init_on_start = "" | |
| self.inited_processors = [] | |
| self.is_demo_row_render = False | |
| def process_plugin_manifest(self, modname, manifest): | |
| # adding processors from plugin manifest | |
| if "img_processor" in manifest: # process commands | |
| for cmd in manifest["img_processor"].keys(): | |
| self.processors[cmd] = manifest["img_processor"][cmd] | |
| return manifest | |
| def init_with_plugins(self): | |
| self.init_plugins(["core"]) | |
| self.display_init_info() | |
| #self.init_translator_engine(self.default_translator) | |
| init_on_start_arr = self.init_on_start.split(",") | |
| for proc_id in init_on_start_arr: | |
| self.init_processor(proc_id) | |
| def run_chain(self, img, params:dict[str,Any] = None, chain:str = None, thread_index:int = 0): | |
| if chain is None: | |
| chain = self.default_chain | |
| if params is None: | |
| params = {} | |
| params["_thread_index"] = thread_index | |
| chain_ar = chain.split(",") | |
| # init all not inited processors first | |
| for proc_id in chain_ar: | |
| if proc_id != "": | |
| if not proc_id in self.inited_processors: | |
| self.init_processor(proc_id) | |
| # run processing | |
| if self.is_demo_row_render: | |
| import cv2 | |
| import numpy as np | |
| height, width, channels = img.shape | |
| img_blank = np.zeros((height+30, width*(1+len(chain_ar)), 3), dtype=np.uint8) | |
| img_blank.fill(255) | |
| y = 30 | |
| x = 0 | |
| img_blank[y:y + height, x:x + width] = img | |
| # Set the font scale and thickness | |
| font_scale = 1 | |
| thickness = 2 | |
| # Set the font face to a monospace font | |
| font_face = cv2.FONT_HERSHEY_SIMPLEX | |
| cv2.putText(img_blank, "original", (x+4, y-7), font_face, font_scale, (0, 0, 0), thickness) | |
| i = 0 | |
| for proc_id in chain_ar: | |
| i += 1 | |
| if proc_id != "": | |
| #img = self.processors[proc_id][1](self, img, params) # params can be modified inside | |
| y = 30 | |
| img = self.processors_objects[proc_id][thread_index].process(img,params) | |
| if self.is_demo_row_render: | |
| x = width*i | |
| img_blank[y:y + height, x:x + width] = img | |
| cv2.putText(img_blank, proc_id, (x + 4, y - 7), font_face, font_scale, (0, 0, 0), thickness) | |
| if self.is_demo_row_render: | |
| return img_blank, params | |
| return img, params | |
| # ---------------- init translation stuff ---------------- | |
| def fill_processors_for_thread_chains(self, threads:int = 1, chain:str = None): | |
| if chain is None: | |
| chain = self.default_chain | |
| chain_ar = chain.split(",") | |
| # init all not initialized processors first | |
| for processor_id in chain_ar: | |
| if processor_id != "": | |
| if self.processors_objects.get(processor_id) is None: | |
| self.processors_objects[processor_id] = [] | |
| while len(self.processors_objects[processor_id]) < threads: | |
| self.add_processor_to_list(processor_id) | |
| def add_processor_to_list(self, processor_id: str): | |
| obj = self.processors[processor_id](self) | |
| obj.init_plugin() | |
| if self.processors_objects.get(processor_id) is None: | |
| self.processors_objects[processor_id] = [] | |
| self.processors_objects[processor_id].append(obj) | |
| def init_processor(self, processor_id: str): | |
| if processor_id == "": # blank line case | |
| return | |
| if processor_id in self.inited_processors: | |
| return | |
| try: | |
| if self.verbose: | |
| self.print_blue("TRY: init processor plugin '{0}'...".format(processor_id)) | |
| self.add_processor_to_list(processor_id) | |
| self.inited_processors.append(processor_id) | |
| if self.verbose: | |
| self.print_blue("SUCCESS: '{0}' initialized!".format(processor_id)) | |
| except Exception as e: | |
| self.print_error("Error init processor plugin {0}...".format(processor_id), e) | |
| # ------------ formatting stuff ------------------- | |
| def display_init_info(self): | |
| if self.verbose: | |
| print("ChainImgProcessor v{0}:".format(version)) | |
| self.format_print_key_list("processors:", self.processors.keys()) | |
| def format_print_key_list(self, key:str, value:list): | |
| print(key+": ".join(value)) | |
| def print_error(self,err_txt,e:Exception = None): | |
| print(err_txt,"red") | |
| # if e != None: | |
| # cprint(e,"red") | |
| import traceback | |
| traceback.print_exc() | |
| def print_red(self,txt): | |
| print(txt) | |
| def print_blue(self, txt): | |
| print(txt) | |
| class ChainImgPlugin: | |
| device = 'cpu' | |
| def __init__(self, core: ChainImgProcessor): | |
| self.core = core | |
| self.device = get_device() | |
| def init_plugin(self): # here you can init something. Called once | |
| pass | |
| def process(self, img, params:dict): # process img. Called multiple | |
| return img | |
| _img_processor:ChainImgProcessor = None | |
| def get_single_image_processor() -> ChainImgProcessor: | |
| global _img_processor | |
| if _img_processor is None: | |
| _img_processor = ChainImgProcessor() | |
| _img_processor.init_with_plugins() | |
| return _img_processor |