add model manager

inference_manager.py

@@ -84,12 +84,15 @@ class AuthHelper:
         raise Exception("wrong token")
 
 class InferenceManager:
-    def __init__(self, model_version="xl", config_path="config.json", lora_options_path="loras.json"):
-        self.model_version = model_version
+    def __init__(self, config_path="config.json"):
+        self.cfg = {}
+        with open(config_path, "r", encoding="utf-8") as f:
+            self.cfg = json.load(f)
+        lora_options_path = cfg.get("loras", "")
         self.lora_load_options = self.load_json(lora_options_path)  # Load LoRA load options
         self.lora_models = self.load_index_file("index.json")  # Load index.json
         self.preloaded_loras = []  # Array to store preloaded LoRAs with name and weights
-        self.base_model_pipeline = self.load_base_model(config_path)  # Load the base model
+        self.base_model_pipeline = self.load_base_model()  # Load the base model
         self.preload_loras()  # Preload LoRAs based on options
 
     def load_json(self, filepath):
@@ -132,13 +135,12 @@ class InferenceManager:
         self.base_model_pipeline = pipe
         print(f"OneDiff compile in {compile_time}s")
     
-    def load_base_model(self, config_path):
+    def load_base_model(self):
         """Load base model and return the pipeline."""
         start = time.time()
-        with open(config_path, "r", encoding="utf-8") as f:
-            cfg = json.load(f)
+        cfg = self.cfg
 
-        model_version = cfg.get("model_version", self.model_version)
+        self.model_version = cfg.get("model_version", self.model_version)
         ckpt_dir = snapshot_download(repo_id=cfg["model_id"], local_files_only=False)
 
         if model_version == "1.5":
@@ -257,6 +259,81 @@ class InferenceManager:
         torch.cuda.empty_cache()
         print("Memory released and cache cleared.")
 
+class ModelManager:
+    def __init__(self, config_urls):
+        """
+        Initialize the ModelManager by loading all models specified by the URLs.
+
+        :param config_urls: List of URLs pointing to model config files (e.g., ["model1/config.json", "model2/config.json"]).
+        """
+        self.models = {}
+        self.load_models(config_urls)
+
+    def load_models(self, config_urls):
+        """
+        Load and initialize InferenceManager instances for each config URL.
+
+        :param config_urls: List of config file URLs.
+        """
+        for url in config_urls:
+            model_name = self.get_model_name_from_url(url)
+            print(f"Initializing model: {model_name} from {url}")
+            try:
+                # Initialize InferenceManager for each model
+                self.models[model_name] = InferenceManager(config_path=url)
+            except Exception as e:
+                print(f"Failed to initialize model {model_name} from {url}: {e}")
+
+    def get_model_name_from_url(self, url):
+        """
+        Extract the model name from the config URL (filename without extension).
+
+        :param url: The URL of the configuration file.
+        :return: The model name (file name without extension).
+        """
+        filename = os.path.basename(url)
+        model_name, _ = os.path.splitext(filename)
+        return model_name
+
+    def get_model_pipeline(self, model_id, lora_list, sampler="DPM2 a", new_pipeline=False):
+        """
+        Build the pipeline with specific LoRAs for a model.
+
+        :param model_id: The model ID (the model name extracted from the config URL).
+        :param lora_list: List of LoRAs to be applied to the model pipeline.
+        :param sampler: The sampler to be used for the pipeline.
+        :param new_pipeline: Flag to indicate whether to create a new pipeline or reuse the existing one.
+        :return: The built pipeline with LoRAs applied.
+        """
+        model = self.models.get(model_id)
+        if not model:
+            print(f"Model {model_id} not found.")
+            return None
+        try:
+            print(f"Building pipeline with LoRAs for model {model_id}...")
+            return model.build_pipeline_with_lora(lora_list, sampler, new_pipeline)
+        except Exception as e:
+            print(f"Failed to build pipeline for model {model_id}: {e}")
+            return None
+
+    def release_model(self, model_id):
+        """
+        Release resources and clear memory for a specific model.
+
+        :param model_id: The model ID (the model name extracted from the config URL).
+        """
+        model = self.models.get(model_id)
+        if not model:
+            print(f"Model {model_id} not found.")
+            return
+        try:
+            print(f"Releasing model {model_id}...")
+            model.release(model.base_model_pipeline)
+        except Exception as e:
+            print(f"Failed to release model {model_id}: {e}")
+
+
+
 
 # Hugging Face file download function - returns only file path
 def download_from_hf(filename, local_dir=None):