FaceEnhance / install.py
Rishi Desai
fixing lfs install issue
eb9fb03
import os
COMFYUI_PATH = "./ComfyUI"
HF_CACHE = os.getenv('HF_HOME')
os.makedirs(HF_CACHE, exist_ok=True)
def run_cmd(command):
"""Run a shell command"""
print(f"πŸ”„ Running: {command}")
exit_code = os.system(command)
if exit_code != 0:
print(f"❌ Command failed: {command} (Exit Code: {exit_code})")
exit(1)
def install_lfs_files():
"""Install dependencies, pull git lfs files."""
run_cmd("apt-get update && apt-get install -y git-lfs")
hook_path = os.path.join(".git", "hooks", "pre-push")
if os.path.exists(hook_path):
print("⚠️ Git LFS hook already exists. Skipping installation.")
else:
run_cmd("git lfs install")
run_cmd("git lfs pull")
print("βœ… Git LFS installed and files pulled.")
def install_git_repo(repo_url, install_path, requirements=False, submodules=False):
"""Clone or update a git repository and handle its dependencies"""
original_dir = os.getcwd()
if not os.path.exists(install_path) or not os.path.isdir(install_path) or not os.path.exists(
os.path.join(install_path, ".git")):
print(f"πŸ“‚ Cloning {os.path.basename(install_path)}...")
run_cmd(f"git clone {repo_url} {install_path}")
else:
print(f"πŸ”„ {os.path.basename(install_path)} exists. Checking for updates...")
# Change to repo directory and update
os.chdir(install_path)
run_cmd("git pull")
if submodules:
run_cmd("git submodule update --init --recursive")
if requirements:
run_cmd("python -m pip install -r requirements.txt")
print(f"βœ… {os.path.basename(install_path)} installed and updated.")
os.chdir(original_dir)
def install_comfyui():
"""Clone and set up ComfyUI if not already installed."""
install_git_repo(
"https://github.com/comfyanonymous/ComfyUI.git",
COMFYUI_PATH,
requirements=True
)
def download_huggingface_models():
"""Download required models from Hugging Face and symlink to ComfyUI models directory."""
from huggingface_hub import hf_hub_download
hf_models = [
{"repo_id": "black-forest-labs/FLUX.1-dev", "filename": "flux1-dev.safetensors", "folder": "unet"},
{"repo_id": "black-forest-labs/FLUX.1-dev", "filename": "ae.safetensors", "folder": "vae"},
{"repo_id": "Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro", "filename": "diffusion_pytorch_model.safetensors",
"folder": "controlnet"},
{"repo_id": "guozinan/PuLID", "filename": "pulid_flux_v0.9.1.safetensors", "folder": "pulid"},
{"repo_id": "comfyanonymous/flux_text_encoders", "filename": "t5xxl_fp16.safetensors",
"folder": "text_encoders"},
{"repo_id": "comfyanonymous/flux_text_encoders", "filename": "clip_l.safetensors", "folder": "text_encoders"},
]
# More specific filenames
filename_mappings = {
"Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro": "Flux_Dev_ControlNet_Union_Pro_ShakkerLabs.safetensors",
}
for model in hf_models:
try:
model_path = hf_hub_download(
repo_id=model["repo_id"],
filename=model["filename"],
cache_dir=HF_CACHE
,
repo_type=model.get("repo_type", "model"),
token=os.getenv('HUGGINGFACE_TOKEN')
)
target_dir = os.path.join(COMFYUI_PATH, "models", model["folder"])
os.makedirs(target_dir, exist_ok=True)
# Use mapping if it exists, otherwise use original filename
file_name_only = filename_mappings.get(model["repo_id"], os.path.basename(model["filename"]))
target_path = os.path.join(target_dir, file_name_only)
if not os.path.exists(target_path):
os.symlink(model_path, target_path)
print(f"βœ… Linked: {file_name_only}")
else:
print(f"βœ… Already exists: {file_name_only}")
except Exception as e:
print(f"❌ Failed to download {model['filename']}: {e}")
def download_and_extract_antelopev2():
"""Download and extract AntelopeV2 model for insightface."""
import zipfile, requests, shutil
base_path = os.path.join(COMFYUI_PATH, "models", "insightface/models")
model_target_path = os.path.join(base_path, "antelopev2")
download_url = "https://huggingface.co/MonsterMMORPG/tools/resolve/main/antelopev2.zip"
zip_path = os.path.join(base_path, "antelopev2.zip")
temp_extract_path = os.path.join(base_path, "temp_antelopev2")
os.makedirs(base_path, exist_ok=True)
if not os.path.exists(model_target_path) or not os.listdir(model_target_path):
# First, remove any existing problematic directories
if os.path.exists(model_target_path):
shutil.rmtree(model_target_path)
print(f"πŸ“₯ Downloading AntelopeV2 model...")
try:
response = requests.get(download_url, stream=True)
response.raise_for_status()
with open(zip_path, "wb") as file:
for chunk in response.iter_content(chunk_size=8192):
file.write(chunk)
print("βœ… Download complete.")
# Create a temporary extraction directory
os.makedirs(temp_extract_path, exist_ok=True)
print("πŸ“‚ Extracting AntelopeV2 model...")
with zipfile.ZipFile(zip_path, "r") as zip_ref:
zip_ref.extractall(temp_extract_path)
print("βœ… Extraction complete.")
# Create the target directory
os.makedirs(model_target_path, exist_ok=True)
# Move the model files to the correct location
# The ZIP contains a nested antelopev2 directory we need to move files from
nested_model_dir = os.path.join(temp_extract_path, "antelopev2")
if os.path.exists(nested_model_dir):
for item in os.listdir(nested_model_dir):
source = os.path.join(nested_model_dir, item)
target = os.path.join(model_target_path, item)
shutil.move(source, target)
# Clean up
if os.path.exists(temp_extract_path):
shutil.rmtree(temp_extract_path)
if os.path.exists(zip_path):
os.remove(zip_path)
print("πŸ—‘οΈ Cleaned up temporary files.")
print("βœ… AntelopeV2 model installed correctly.")
except Exception as e:
print(f"❌ Failed to download/extract AntelopeV2: {e}")
else:
print("βœ… AntelopeV2 model already exists")
def install_custom_nodes():
"""Install all custom nodes for ComfyUI."""
custom_nodes = [
{
"repo": "https://github.com/sipie800/ComfyUI-PuLID-Flux-Enhanced",
"name": "ComfyUI-PuLID-Flux-Enhanced",
"requirements": True
},
{
"repo": "https://github.com/rgthree/rgthree-comfy",
"name": "rgthree-comfy",
"requirements": True
},
{ # we already have insightface so don't need requirements (no dlib)
"repo": "https://github.com/cubiq/ComfyUI_FaceAnalysis",
"name": "ComfyUI_FaceAnalysis",
"requirements": False
},
]
for node in custom_nodes:
repo_name = node["name"]
repo_path = os.path.join(COMFYUI_PATH, "custom_nodes", repo_name)
install_git_repo(
node["repo"],
repo_path,
requirements=node.get("requirements", False),
submodules=node.get("submodules", False)
)
print(f"βœ… {repo_name} installed and updated.")
print("βœ… Installed and updated all ComfyUI nodes.")
def install():
install_lfs_files()
install_comfyui()
install_custom_nodes()
download_huggingface_models()
download_and_extract_antelopev2()
print("πŸŽ‰ Setup Complete!")
if __name__ == "__main__":
install()