File size: 3,955 Bytes
0163a2c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
import os
from datetime import datetime
from pathlib import Path
from huggingface_hub import snapshot_download
from modules import shared
from ia_config import get_webui_setting
from ia_logging import ia_logging
class IAFileManager:
DOWNLOAD_COMPLETE = "Download complete"
def __init__(self) -> None:
self.update_ia_outputs_dir()
self.update_ia_sam_models_dir()
def update_ia_outputs_dir(self) -> None:
"""Update inpaint-anything outputs directory.
Returns:
None
"""
config_save_folder = get_webui_setting("inpaint_anything_save_folder", "inpaint-anything")
self.folder_is_webui = False if config_save_folder == "inpaint-anything" else True
if config_save_folder == "inpaint-anything":
self._ia_outputs_dir = os.path.join(shared.data_path, "outputs", config_save_folder, datetime.now().strftime("%Y-%m-%d"))
else:
try:
webui_save_folder = Path(get_webui_setting("outdir_img2img_samples", os.path.join("outputs", "img2img-images")))
if webui_save_folder.is_absolute():
self._ia_outputs_dir = os.path.join(str(webui_save_folder), datetime.now().strftime("%Y-%m-%d"))
else:
self._ia_outputs_dir = os.path.join(shared.data_path, str(webui_save_folder), datetime.now().strftime("%Y-%m-%d"))
except Exception:
self._ia_outputs_dir = os.path.join(shared.data_path, "outputs", "img2img-images", datetime.now().strftime("%Y-%m-%d"))
def update_ia_sam_models_dir(self) -> None:
"""Update inpaint-anything sam models directory.
Returns:
None
"""
sam_models_dir = get_webui_setting("inpain_anything_sam_models_dir", "")
if (len(sam_models_dir) == 0) or (sam_models_dir.isspace()):
sam_models_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models")
self._ia_models_dir = sam_models_dir
@property
def outputs_dir(self) -> str:
"""Get inpaint-anything outputs directory.
Returns:
str: inpaint-anything outputs directory
"""
self.update_ia_outputs_dir()
if not os.path.isdir(self._ia_outputs_dir):
os.makedirs(self._ia_outputs_dir, exist_ok=True)
return self._ia_outputs_dir
@property
def models_dir(self) -> str:
"""Get inpaint-anything models directory.
Returns:
str: inpaint-anything models directory
"""
self.update_ia_sam_models_dir()
if not os.path.isdir(self._ia_models_dir):
os.makedirs(self._ia_models_dir, exist_ok=True)
return self._ia_models_dir
@property
def savename_prefix(self) -> str:
"""Get inpaint-anything savename prefix.
Returns:
str: inpaint-anything savename prefix
"""
config_save_folder = get_webui_setting("inpaint_anything_save_folder", "inpaint-anything")
self.folder_is_webui = False if config_save_folder == "inpaint-anything" else True
basename = "inpainta-" if self.folder_is_webui else ""
return basename + datetime.now().strftime("%Y%m%d-%H%M%S")
ia_file_manager = IAFileManager()
def download_model_from_hf(hf_model_id, local_files_only=False):
"""Download model from HuggingFace Hub.
Args:
sam_model_id (str): HuggingFace model id
local_files_only (bool, optional): If True, use only local files. Defaults to False.
Returns:
str: download status
"""
if not local_files_only:
ia_logging.info(f"Downloading {hf_model_id}")
try:
snapshot_download(repo_id=hf_model_id, local_files_only=local_files_only)
except FileNotFoundError:
return f"{hf_model_id} not found, please download"
except Exception as e:
return str(e)
return IAFileManager.DOWNLOAD_COMPLETE
|