Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Allow passing of a civit api key

Browse Source
This commit is contained in:
Brandon Rising 2024-01-30 20:35:01 -05:00 committed by Kent Keirsey
parent f48a2c5fd2
commit 14efc95707
2 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
invokeai/backend/install/model_install_backend.py
View File

@ -104,12 +104,14 @@ class ModelInstall(object):
prediction_type_helper: Optional[Callable[[Path], SchedulerPredictionType]] = None, prediction_type_helper: Optional[Callable[[Path], SchedulerPredictionType]] = None,
model_manager: Optional[ModelManager] = None, model_manager: Optional[ModelManager] = None,
access_token: Optional[str] = None, access_token: Optional[str] = None,
civit_api_key: Optional[str] = None,
): ):
self.config = config self.config = config
self.mgr = model_manager or ModelManager(config.model_conf_path) self.mgr = model_manager or ModelManager(config.model_conf_path)
self.datasets = OmegaConf.load(Dataset_path) self.datasets = OmegaConf.load(Dataset_path)
self.prediction_helper = prediction_type_helper self.prediction_helper = prediction_type_helper
self.access_token = access_token or HfFolder.get_token() 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) self.reverse_paths = self._reverse_paths(self.datasets)
def all_models(self) -> Dict[str, ModelLoadInfo]: def all_models(self) -> Dict[str, ModelLoadInfo]:
@ -326,7 +328,10 @@ class ModelInstall(object):
def _install_url(self, url: str) -> AddModelResult: def _install_url(self, url: str) -> AddModelResult:
with TemporaryDirectory(dir=self.config.models_path) as staging: 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: if not location:
logger.error(f"Unable to download {url}. Skipping.") logger.error(f"Unable to download {url}. Skipping.")
info = ModelProbe().heuristic_probe(location, self.prediction_helper) info = ModelProbe().heuristic_probe(location, self.prediction_helper)

3
invokeai/backend/util/util.py
View File

@ -286,9 +286,8 @@ def download_with_resume(url: str, dest: Path, access_token: str = None) -> Path
open_mode = "wb" open_mode = "wb"
exist_size = 0 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)) content_length = int(resp.headers.get("content-length", 0))
if dest.is_dir(): if dest.is_dir():
try: try:
file_name = re.search('filename="(.+)"', resp.headers.get("Content-Disposition")).group(1) file_name = re.search('filename="(.+)"', resp.headers.get("Content-Disposition")).group(1)