Spaces:
Running on Zero
Running on Zero
File size: 5,290 Bytes
bbf161a | 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | import yaml
import os
from collections import OrderedDict
CHECKPOINT_DIR = "models/checkpoints"
LORA_DIR = "models/loras"
EMBEDDING_DIR = "models/embeddings"
CONTROLNET_DIR = "models/controlnet"
MODEL_PATCHES_DIR = "models/model_patches"
DIFFUSION_MODELS_DIR = "models/diffusion_models"
VAE_DIR = "models/vae"
TEXT_ENCODERS_DIR = "models/text_encoders"
STYLE_MODELS_DIR = "models/style_models"
CLIP_VISION_DIR = "models/clip_vision"
IPADAPTER_DIR = "models/ipadapter"
IPADAPTER_FLUX_DIR = "models/ipadapter-flux"
INPUT_DIR = "input"
OUTPUT_DIR = "output"
CATEGORY_TO_DIR_MAP = {
"diffusion_models": DIFFUSION_MODELS_DIR,
"text_encoders": TEXT_ENCODERS_DIR,
"vae": VAE_DIR,
"checkpoints": CHECKPOINT_DIR,
"loras": LORA_DIR,
"controlnet": CONTROLNET_DIR,
"model_patches": MODEL_PATCHES_DIR,
"embeddings": EMBEDDING_DIR,
"style_models": STYLE_MODELS_DIR,
"clip_vision": CLIP_VISION_DIR,
"ipadapter": IPADAPTER_DIR,
"ipadapter-flux": IPADAPTER_FLUX_DIR
}
_PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_FILE_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'file_list.yaml')
_CONSTANTS_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'constants.yaml')
def load_constants_from_yaml(filepath=_CONSTANTS_PATH):
if not os.path.exists(filepath):
print(f"Warning: Constants file not found at {filepath}. Using fallback values.")
return {}
with open(filepath, 'r', encoding='utf-8') as f:
return yaml.safe_load(f)
def load_file_download_map(filepath=_FILE_LIST_PATH):
if not os.path.exists(filepath):
raise FileNotFoundError(f"The file list (for downloads) was not found at: {filepath}")
with open(filepath, 'r', encoding='utf-8') as f:
file_list_data = yaml.safe_load(f)
download_info_map = {}
for category, files in file_list_data.get('file', {}).items():
if isinstance(files, list):
for file_info in files:
if 'filename' in file_info:
file_info['category'] = category
download_info_map[file_info['filename']] = file_info
return download_info_map
try:
ALL_FILE_DOWNLOAD_MAP = load_file_download_map()
category_to_model_type = {
"diffusion_models": "UNET",
"text_encoders": "TEXT_ENCODER",
"vae": "VAE",
"checkpoints": "SDXL",
"loras": "LORA",
"controlnet": "CONTROLNET",
"model_patches": "MODEL_PATCH",
"style_models": "STYLE",
"clip_vision": "CLIP_VISION",
"ipadapter": "IPADAPTER",
"ipadapter-flux": "IPADAPTER_FLUX"
}
MODEL_MAP_CHECKPOINT = OrderedDict()
ALL_MODEL_MAP = OrderedDict()
for filename, file_info in ALL_FILE_DOWNLOAD_MAP.items():
category = file_info.get('category')
repo_id = file_info.get('repo_id', '')
model_type = category_to_model_type.get(category, 'UNKNOWN')
model_tuple = (repo_id, filename, model_type, "latent", category)
ALL_MODEL_MAP[filename] = model_tuple
if category == 'checkpoints':
MODEL_MAP_CHECKPOINT[filename] = model_tuple
MODEL_TYPE_MAP = {k: v[2] for k, v in ALL_MODEL_MAP.items()}
ARCH_CATEGORIES_MAP = {}
for display_name, info in MODEL_MAP_CHECKPOINT.items():
arch = info[2]
cat = info[4] if len(info) > 4 else None
if arch not in ARCH_CATEGORIES_MAP:
ARCH_CATEGORIES_MAP[arch] = []
if cat and cat not in ARCH_CATEGORIES_MAP[arch]:
ARCH_CATEGORIES_MAP[arch].append(cat)
except Exception as e:
print(f"FATAL: Could not load file download map from YAML. Error: {e}")
ALL_FILE_DOWNLOAD_MAP = {}
MODEL_MAP_CHECKPOINT, ALL_MODEL_MAP = OrderedDict(), OrderedDict()
MODEL_TYPE_MAP = {}
ARCH_CATEGORIES_MAP = {}
try:
_constants = load_constants_from_yaml()
MAX_LORAS = _constants.get('MAX_LORAS', 10)
MAX_EMBEDDINGS = _constants.get('MAX_EMBEDDINGS', 5)
MAX_CONDITIONINGS = _constants.get('MAX_CONDITIONINGS', 10)
MAX_CONTROLNETS = _constants.get('MAX_CONTROLNETS', 5)
MAX_H3_GUIDES = _constants.get('MAX_H3_GUIDES', 5)
MAX_IPADAPTERS = _constants.get('MAX_IPADAPTERS', 5)
LORA_SOURCE_CHOICES = _constants.get('LORA_SOURCE_CHOICES', ["Hugging Face", "Civitai", "File"])
RESOLUTION_MAP = _constants.get('RESOLUTION_MAP', {})
MULTIPLIERS_MAP = _constants.get('MULTIPLIERS_MAP', {})
ARCHITECTURES_CONFIG = {"architectures": {}, "architecture_order": []}
FEATURES_CONFIG = {"default": {"enabled_chains": ["lora"]}}
MODEL_DEFAULTS_CONFIG = {"Default": {}}
except Exception as e:
print(f"FATAL: Could not load constants from YAML. Error: {e}")
MAX_LORAS, MAX_EMBEDDINGS, MAX_CONDITIONINGS, MAX_CONTROLNETS, MAX_IPADAPTERS = 10, 5, 10, 5, 5
MAX_H3_GUIDES = 5
LORA_SOURCE_CHOICES = ["Hugging Face", "Civitai", "File"]
RESOLUTION_MAP = {}
MULTIPLIERS_MAP = {}
ARCHITECTURES_CONFIG = {"architectures": {}, "architecture_order": []}
FEATURES_CONFIG = {"default": {"enabled_chains": ["lora"]}}
MODEL_DEFAULTS_CONFIG = {"Default": {}} |