import os # Define paths BASE_PATH = "./" COMFYUI_PATH = os.path.join(BASE_PATH, "ComfyUI") MODEL_PATH = os.path.join(COMFYUI_PATH, "models") CACHE_PATH = "/workspace/huggingface_cache" os.environ["HF_HOME"] = CACHE_PATH os.makedirs(CACHE_PATH, exist_ok=True) def run_command(command): """Run a shell command using os.system() (simpler than subprocess).""" print(f"🔄 Running: {command}") exit_code = os.system(command) if exit_code != 0: print(f"❌ Command failed: {command} (Exit Code: {exit_code})") exit(1) # Exit on failure # def install_dependencies(): # """Install system dependencies and Python packages.""" # print("📦 Installing dependencies...") # # run_command("apt-get update && apt-get install -y git wget curl libgl1-mesa-glx libglib2.0-0 tmux emacs git-lfs") # run_command("pip install --upgrade pip") # run_command("pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124") # run_command("pip install xformers --index-url https://download.pytorch.org/whl/cu124") # # run_command("pip install -r requirements.txt") # print("✅ Dependencies installed.") def manage_git_repo(repo_url, install_path, requirements=False, submodules=False): """Clone or update a git repository and handle its dependencies. Args: repo_url: URL of the git repository install_path: Where to install/update the repository requirements: Whether to install requirements.txt submodules: Whether to update git submodules """ # Save the original directory 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_command(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_command("git pull") if submodules: run_command("git submodule update --init --recursive") if requirements: run_command("python -m pip install -r requirements.txt") print(f"✅ {os.path.basename(install_path)} installed and updated.") # Change back to the original directory os.chdir(original_dir) def install_comfyui(): """Clone and set up ComfyUI if not already installed.""" manage_git_repo( "https://github.com/comfyanonymous/ComfyUI.git", COMFYUI_PATH, requirements=True ) def download_huggingface_models(): """Download required models from Hugging Face.""" 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"}, ] # Dictionary mapping repo_ids to 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=CACHE_PATH, repo_type=model.get("repo_type", "model") ) target_dir = os.path.join(MODEL_PATH, 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 base_path = os.path.join(BASE_PATH, "models/insightface/models") download_url = "https://huggingface.co/MonsterMMORPG/tools/resolve/main/antelopev2.zip" zip_path = os.path.join(base_path, "antelopev2.zip") extract_path = os.path.join(base_path, "antelopev2") os.makedirs(base_path, exist_ok=True) if not os.path.exists(extract_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.") print("📂 Extracting AntelopeV2 model...") with zipfile.ZipFile(zip_path, "r") as zip_ref: zip_ref.extractall(base_path) print("✅ Extraction complete.") os.remove(zip_path) print("🗑️ Cleaned up ZIP file.") 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_git = [ { "repo": "https://github.com/ltdrdata/ComfyUI-Manager", "name": "ComfyUI-Manager", "requirements": True }, { "repo": "https://github.com/balazik/ComfyUI-PuLID-Flux", "name": "ComfyUI-PuLID-Flux", "requirements": True }, { "repo": "https://github.com/rgthree/rgthree-comfy", "name": "rgthree-comfy", "requirements": True }, { # we already have insightface so don't need requirements (for dlib) "repo": "https://github.com/cubiq/ComfyUI_FaceAnalysis", "name": "ComfyUI_FaceAnalysis", "requirements": False } ] for node in custom_nodes_git: repo_name = node["name"] repo_path = os.path.join(COMFYUI_PATH, "custom_nodes", repo_name) manage_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.") if __name__ == "__main__": # install_dependencies() # install_comfyui() # download_huggingface_models() install_custom_nodes() print("🎉 Setup Complete! Run `run.py` to start ComfyUI.")