From 14efc95707937264eadf052a5827aba06e0181e3 Mon Sep 17 00:00:00 2001 From: Brandon Rising Date: Tue, 30 Jan 2024 20:35:01 -0500 Subject: [PATCH 1/7] Allow passing of a civit api key --- invokeai/backend/install/model_install_backend.py | 7 ++++++- invokeai/backend/util/util.py | 3 +-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/invokeai/backend/install/model_install_backend.py b/invokeai/backend/install/model_install_backend.py index 2ffc1e6ff4..d412d0226f 100644 --- a/invokeai/backend/install/model_install_backend.py +++ b/invokeai/backend/install/model_install_backend.py @@ -104,12 +104,14 @@ class ModelInstall(object): prediction_type_helper: Optional[Callable[[Path], SchedulerPredictionType]] = None, model_manager: Optional[ModelManager] = None, access_token: Optional[str] = None, + civit_api_key: Optional[str] = None, ): self.config = config self.mgr = model_manager or ModelManager(config.model_conf_path) self.datasets = OmegaConf.load(Dataset_path) self.prediction_helper = prediction_type_helper self.access_token = access_token or HfFolder.get_token() + self.civit_api_key = civit_api_key or os.environ.get("CIVIT_API_KEY") self.reverse_paths = self._reverse_paths(self.datasets) def all_models(self) -> Dict[str, ModelLoadInfo]: @@ -326,7 +328,10 @@ class ModelInstall(object): def _install_url(self, url: str) -> AddModelResult: with TemporaryDirectory(dir=self.config.models_path) as staging: - location = download_with_resume(url, Path(staging)) + CIVITAI_RE = r".*civitai.com.*" + civit_url = re.match(CIVITAI_RE, url, re.IGNORECASE) + print(civit_url) + location = download_with_resume(url, Path(staging), access_token=self.civit_api_key if civit_url else None) if not location: logger.error(f"Unable to download {url}. Skipping.") info = ModelProbe().heuristic_probe(location, self.prediction_helper) diff --git a/invokeai/backend/util/util.py b/invokeai/backend/util/util.py index 4612b42cb9..85a4488915 100644 --- a/invokeai/backend/util/util.py +++ b/invokeai/backend/util/util.py @@ -286,9 +286,8 @@ def download_with_resume(url: str, dest: Path, access_token: str = None) -> Path open_mode = "wb" exist_size = 0 - resp = requests.get(url, header, stream=True) + resp = requests.get(url, headers=header, stream=True, allow_redirects=True) content_length = int(resp.headers.get("content-length", 0)) - if dest.is_dir(): try: file_name = re.search('filename="(.+)"', resp.headers.get("Content-Disposition")).group(1) From 088e3420e6883398f9bafac5fa1a0a1e687d7598 Mon Sep 17 00:00:00 2001 From: Brandon Rising Date: Tue, 30 Jan 2024 20:46:42 -0500 Subject: [PATCH 2/7] Allow passing of civit api key via config --- invokeai/app/services/config/config_default.py | 2 ++ invokeai/backend/install/model_install_backend.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index 3af906bd04..c13900cfee 100644 --- a/invokeai/app/services/config/config_default.py +++ b/invokeai/app/services/config/config_default.py @@ -293,6 +293,8 @@ class InvokeAIAppConfig(InvokeAISettings): lora_dir : Optional[Path] = Field(default=None, description='Path to a directory of LoRA/LyCORIS models to be imported on startup.', json_schema_extra=Categories.Paths) embedding_dir : Optional[Path] = Field(default=None, description='Path to a directory of Textual Inversion embeddings to be imported on startup.', json_schema_extra=Categories.Paths) controlnet_dir : Optional[Path] = Field(default=None, description='Path to a directory of ControlNet embeddings to be imported on startup.', json_schema_extra=Categories.Paths) + + civit_api_key : Optional[str] = Field(default=os.environ.get("CIVIT_API_KEY"), description="API key for Civit", json_schema_extra=Categories.Other) # this is not referred to in the source code and can be removed entirely #free_gpu_mem : Optional[bool] = Field(default=None, description="If true, purge model from GPU after each generation.", json_schema_extra=Categories.MemoryPerformance) diff --git a/invokeai/backend/install/model_install_backend.py b/invokeai/backend/install/model_install_backend.py index d412d0226f..8f4fdece31 100644 --- a/invokeai/backend/install/model_install_backend.py +++ b/invokeai/backend/install/model_install_backend.py @@ -111,7 +111,7 @@ class ModelInstall(object): self.datasets = OmegaConf.load(Dataset_path) self.prediction_helper = prediction_type_helper self.access_token = access_token or HfFolder.get_token() - self.civit_api_key = civit_api_key or os.environ.get("CIVIT_API_KEY") + self.civit_api_key = civit_api_key or config.civit_api_key self.reverse_paths = self._reverse_paths(self.datasets) def all_models(self) -> Dict[str, ModelLoadInfo]: From 5d773dc94c9ed7acc21cca3d7ee5b4d6490e8950 Mon Sep 17 00:00:00 2001 From: Brandon Rising Date: Tue, 30 Jan 2024 20:47:36 -0500 Subject: [PATCH 3/7] Remove debug line --- invokeai/backend/install/model_install_backend.py | 1 - 1 file changed, 1 deletion(-) diff --git a/invokeai/backend/install/model_install_backend.py b/invokeai/backend/install/model_install_backend.py index 8f4fdece31..5d7c39beb9 100644 --- a/invokeai/backend/install/model_install_backend.py +++ b/invokeai/backend/install/model_install_backend.py @@ -330,7 +330,6 @@ class ModelInstall(object): with TemporaryDirectory(dir=self.config.models_path) as staging: CIVITAI_RE = r".*civitai.com.*" civit_url = re.match(CIVITAI_RE, url, re.IGNORECASE) - print(civit_url) location = download_with_resume(url, Path(staging), access_token=self.civit_api_key if civit_url else None) if not location: logger.error(f"Unable to download {url}. Skipping.") From 2c5ef92979f92717b860f1e9156d95e97a8aa0c3 Mon Sep 17 00:00:00 2001 From: Brandon Rising Date: Tue, 30 Jan 2024 21:05:15 -0500 Subject: [PATCH 4/7] Move location of config property, comment for explanation of use --- invokeai/app/services/config/config_default.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index c13900cfee..510e60d7bc 100644 --- a/invokeai/app/services/config/config_default.py +++ b/invokeai/app/services/config/config_default.py @@ -284,6 +284,9 @@ class InvokeAIAppConfig(InvokeAISettings): deny_nodes : Optional[List[str]] = Field(default=None, description="List of nodes to deny. Omit to deny none.", json_schema_extra=Categories.Nodes) node_cache_size : int = Field(default=512, description="How many cached nodes to keep in memory", json_schema_extra=Categories.Nodes) + # MODEL IMPORT + civit_api_key : Optional[str] = Field(default=os.environ.get("CIVIT_API_KEY"), description="API key for Civit", json_schema_extra=Categories.Other) + # DEPRECATED FIELDS - STILL HERE IN ORDER TO OBTAN VALUES FROM PRE-3.1 CONFIG FILES always_use_cpu : bool = Field(default=False, description="If true, use the CPU for rendering even if a GPU is available.", json_schema_extra=Categories.MemoryPerformance) max_cache_size : Optional[float] = Field(default=None, gt=0, description="Maximum memory amount used by model cache for rapid switching", json_schema_extra=Categories.MemoryPerformance) @@ -294,7 +297,6 @@ class InvokeAIAppConfig(InvokeAISettings): embedding_dir : Optional[Path] = Field(default=None, description='Path to a directory of Textual Inversion embeddings to be imported on startup.', json_schema_extra=Categories.Paths) controlnet_dir : Optional[Path] = Field(default=None, description='Path to a directory of ControlNet embeddings to be imported on startup.', json_schema_extra=Categories.Paths) - civit_api_key : Optional[str] = Field(default=os.environ.get("CIVIT_API_KEY"), description="API key for Civit", json_schema_extra=Categories.Other) # this is not referred to in the source code and can be removed entirely #free_gpu_mem : Optional[bool] = Field(default=None, description="If true, purge model from GPU after each generation.", json_schema_extra=Categories.MemoryPerformance) From a769f93be0218c1c4030c9c2a21748acf864c01b Mon Sep 17 00:00:00 2001 From: Brandon Rising Date: Tue, 30 Jan 2024 21:06:17 -0500 Subject: [PATCH 5/7] Remove unnecessary change --- invokeai/backend/util/util.py | 1 + 1 file changed, 1 insertion(+) diff --git a/invokeai/backend/util/util.py b/invokeai/backend/util/util.py index 85a4488915..13751e2770 100644 --- a/invokeai/backend/util/util.py +++ b/invokeai/backend/util/util.py @@ -288,6 +288,7 @@ def download_with_resume(url: str, dest: Path, access_token: str = None) -> Path resp = requests.get(url, headers=header, stream=True, allow_redirects=True) content_length = int(resp.headers.get("content-length", 0)) + if dest.is_dir(): try: file_name = re.search('filename="(.+)"', resp.headers.get("Content-Disposition")).group(1) From 522ff4a042007e1710645d3ad2cc6010f7704104 Mon Sep 17 00:00:00 2001 From: Brandon Rising Date: Tue, 30 Jan 2024 22:25:23 -0500 Subject: [PATCH 6/7] civit -> civitai --- invokeai/app/services/config/config_default.py | 2 +- invokeai/backend/install/model_install_backend.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index 510e60d7bc..65af5bc1c4 100644 --- a/invokeai/app/services/config/config_default.py +++ b/invokeai/app/services/config/config_default.py @@ -285,7 +285,7 @@ class InvokeAIAppConfig(InvokeAISettings): node_cache_size : int = Field(default=512, description="How many cached nodes to keep in memory", json_schema_extra=Categories.Nodes) # MODEL IMPORT - civit_api_key : Optional[str] = Field(default=os.environ.get("CIVIT_API_KEY"), description="API key for Civit", json_schema_extra=Categories.Other) + civitai_api_key : Optional[str] = Field(default=os.environ.get("CIVITAI_API_KEY"), description="API key for CivitAI", json_schema_extra=Categories.Other) # DEPRECATED FIELDS - STILL HERE IN ORDER TO OBTAN VALUES FROM PRE-3.1 CONFIG FILES always_use_cpu : bool = Field(default=False, description="If true, use the CPU for rendering even if a GPU is available.", json_schema_extra=Categories.MemoryPerformance) diff --git a/invokeai/backend/install/model_install_backend.py b/invokeai/backend/install/model_install_backend.py index 5d7c39beb9..601f83fc6f 100644 --- a/invokeai/backend/install/model_install_backend.py +++ b/invokeai/backend/install/model_install_backend.py @@ -104,14 +104,14 @@ class ModelInstall(object): prediction_type_helper: Optional[Callable[[Path], SchedulerPredictionType]] = None, model_manager: Optional[ModelManager] = None, access_token: Optional[str] = None, - civit_api_key: Optional[str] = None, + civitai_api_key: Optional[str] = None, ): self.config = config self.mgr = model_manager or ModelManager(config.model_conf_path) self.datasets = OmegaConf.load(Dataset_path) self.prediction_helper = prediction_type_helper self.access_token = access_token or HfFolder.get_token() - self.civit_api_key = civit_api_key or config.civit_api_key + self.civitai_api_key = civitai_api_key or config.civitai_api_key self.reverse_paths = self._reverse_paths(self.datasets) def all_models(self) -> Dict[str, ModelLoadInfo]: @@ -330,7 +330,7 @@ class ModelInstall(object): with TemporaryDirectory(dir=self.config.models_path) as staging: CIVITAI_RE = r".*civitai.com.*" civit_url = re.match(CIVITAI_RE, url, re.IGNORECASE) - location = download_with_resume(url, Path(staging), access_token=self.civit_api_key if civit_url else None) + location = download_with_resume(url, Path(staging), access_token=self.civitai_api_key if civit_url else None) if not location: logger.error(f"Unable to download {url}. Skipping.") info = ModelProbe().heuristic_probe(location, self.prediction_helper) From a0996b1c0aaf27d39dfc60b1bd5910e2ff3c6159 Mon Sep 17 00:00:00 2001 From: Brandon Rising Date: Tue, 30 Jan 2024 22:32:40 -0500 Subject: [PATCH 7/7] Fix ruff styling --- invokeai/backend/install/model_install_backend.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/invokeai/backend/install/model_install_backend.py b/invokeai/backend/install/model_install_backend.py index 601f83fc6f..fdbe714f62 100644 --- a/invokeai/backend/install/model_install_backend.py +++ b/invokeai/backend/install/model_install_backend.py @@ -330,7 +330,9 @@ class ModelInstall(object): with TemporaryDirectory(dir=self.config.models_path) as staging: CIVITAI_RE = r".*civitai.com.*" civit_url = re.match(CIVITAI_RE, url, re.IGNORECASE) - location = download_with_resume(url, Path(staging), access_token=self.civitai_api_key if civit_url else None) + location = download_with_resume( + url, Path(staging), access_token=self.civitai_api_key if civit_url else None + ) if not location: logger.error(f"Unable to download {url}. Skipping.") info = ModelProbe().heuristic_probe(location, self.prediction_helper)