generated from fofr/cog-comfyui
-
Notifications
You must be signed in to change notification settings - Fork 7
/
weights_manifest.py
44 lines (36 loc) · 1.36 KB
/
weights_manifest.py
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
import json
import os
BASE_URL = "https://weights.replicate.delivery/default/comfy-ui"
BASE_PATH = "ComfyUI/models"
from helpers.ComfyUI_Controlnet_Aux import ComfyUI_Controlnet_Aux
class WeightsManifest:
def __init__(self):
self.weights_manifest = self._load_weights_manifest()
self.weights_map = self._initialize_weights_map()
def _load_weights_manifest(self):
return self._load_local_manifest()
def _load_local_manifest(self):
WEIGHTS_MANIFEST_PATH = "weights.json"
if os.path.exists(WEIGHTS_MANIFEST_PATH):
with open(WEIGHTS_MANIFEST_PATH, "r") as f:
return json.load(f)
else:
print("Local weights manifest file does not exist.")
return {}
def _generate_weights_map(self, keys, dest):
return {
key: {
"url": f"{BASE_URL}/{dest}/{key}.tar",
"dest": f"{BASE_PATH}/{dest}",
}
for key in keys
}
def _initialize_weights_map(self):
weights_map = {}
for key in self.weights_manifest.keys():
if key.isupper():
weights_map.update(
self._generate_weights_map(self.weights_manifest[key], key.lower())
)
weights_map.update(ComfyUI_Controlnet_Aux.weights_map(BASE_URL))
return weights_map