58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
import os
|
|
import subprocess
|
|
import requests
|
|
import tempfile
|
|
import shutil
|
|
from tqdm.auto import tqdm
|
|
|
|
|
|
def download_installer(url, dir):
|
|
response = requests.get(url, stream=True, allow_redirects=True)
|
|
response.raise_for_status()
|
|
file_size = int(response.headers.get('Content-Length', 0))
|
|
|
|
with tqdm.wrapattr(response.raw, "read", total=file_size, desc=url) as read:
|
|
dest = os.path.join(dir, "NIMSetup.exe")
|
|
with open(dest, "wb") as f:
|
|
shutil.copyfileobj(read, f)
|
|
|
|
return dest
|
|
|
|
|
|
def run_installer(installer_path):
|
|
# Run the installer via powershell with elevated privileges and wait for it to finish
|
|
print("Launching installer...")
|
|
powershell_command = [
|
|
"powershell.exe",
|
|
"-Command",
|
|
f"""
|
|
try {{
|
|
$process = Start-Process '{installer_path}' -Verb RunAs -Wait -PassThru
|
|
exit $process.ExitCode
|
|
}} catch {{
|
|
exit 403 # UAC cancellation
|
|
}}
|
|
"""
|
|
]
|
|
|
|
print("Waiting for completion...")
|
|
subprocess.run(powershell_command, check=True)
|
|
|
|
|
|
if os.name == 'nt':
|
|
import ctypes
|
|
res = ctypes.windll.user32.MessageBoxW(None, "Do you want to automatically download and run the NVIDIA NIM installer?\n\n" +
|
|
"If you choose not to please ensure you manually setup NVIDIA NIM before attempting to use the node.",
|
|
"NIM Installer", 4)
|
|
|
|
if res == 6:
|
|
url = "https://storage.googleapis.com/comfy-assets/NimSetup.exe"
|
|
|
|
# Download to temp directory so files get deleted
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
installer_path = download_installer(url, tmpdir)
|
|
run_installer(installer_path)
|
|
|
|
else:
|
|
print("NIM node setup is only supported for Windows")
|