main2.py

import random
import string
import subprocess
import sys
import os
import concurrent.futures
import time
import shutil

params = {}
wi_list = ["webui.bat", "webui.sh", "webui-macos-env.sh", "webui-user.bat", "webui-user.sh"]
for arg in sys.argv[1:]:
    if arg.startswith('--'):
        key_value = arg[len('--'):].split('=')
        if len(key_value) == 2:
            key, value = key_value
            params[key] = value

subprocess.run(f'cd {params["dir"]}/', shell=True)
if params["ui"] == "anapnoe手机端完美适配":
    subprocess.run(f'git clone https://github.com/anapnoe/stable-diffusion-webui {params["dir"]}', shell=True)
elif params["ui"] == "AUTOMATIC1111原版v1.6.0[sdxl]":
    subprocess.run(
        f'git clone  -b v1.6.0 https://github.com/s4afa451dgf415f/sdwi {params["dir"]}', shell=True)
else:
    subprocess.run(
        f'git clone -b v1.5.1 https://github.com/s4afa451dgf415f/sdwi {params["dir"]}', shell=True)
for wi_file in wi_list:
    try:
        os.remove(f"{params['dir']}/{wi_file}")
    except FileNotFoundError:
        print(f"文件 '{wi_file}' 不存在")
try:
    os.rename(f"{params['dir']}/webui.py", f"{params['dir']}/wi.py")
except Exception as e:
    print(f"文件命名失败: {str(e)}")

file_path = f'{params["dir"]}/modules/launch_utils.py'

try:
    with open(file_path, 'r') as file:
        content = file.read()
        content = content.replace('import webui', 'import wi')
        content = content.replace('webui.webui', 'wi.webui')
        modified_content = content.replace('webui.api_only','wi.api_only')
        file.close()

    # 将修改后的内容写回到原文件
    with open(file_path, 'w') as file:
        file.write(modified_content)
        file.close()

except FileNotFoundError:
    print(f"文件 '{file_path}' 不存在")
except Exception as e:
    print(f"修改文件 '{file_path}' 时出错: {str(e)}")


checkpoint_url = {
    "Dark_sushi_mix.safetensors": "https://huggingface.co/mdl-mirror/dark-sushi-mix/resolve/main/darkSushiMixMix_brighter.safetensors",
    "AnythingV5V3_v5PrtRE.safetensors": "https://huggingface.co/ckpt/anything-v5.0/resolve/main/AnythingV5V3_v5PrtRE.safetensors",
    "chilloutmix_NiPrunedFp16Fix.safetensors": "https://huggingface.co/naonovn/chilloutmix_NiPrunedFp32Fix/resolve/main/chilloutmix_NiPrunedFp32Fix.safetensors",
    "rpg_V4.safetensors": "https://huggingface.co/Anashel/rpg/resolve/main/RPG-V4-Model-Download/RPG-v4.safetensors",
    "ProtoGen_X5.8-pruned-fp16.safetensors": "https://huggingface.co/darkstorm2150/Protogen_x5.8_Official_Release/resolve/main/ProtoGen_X5.8-pruned-fp16.safetensors",
    "none": "",
}


def run_git_download():
    start_time = time.time()
    subprocess.run(
         f'git clone https://github.com/s4afa451dgf415f/sd-wi-negpip {params["dir"]}/extensions/sd-wi-negpip',
         shell=True)
    subprocess.run(
         f'git clone https://github.com/Bing-su/adetailer {params["dir"]}/extensions/adetailer',
         shell=True)
    subprocess.run(
         f'git clone https://github.com/s4afa451dgf415f/sd-wi-aio {params["dir"]}/extensions/prompt-aio',
         shell=True)
    subprocess.run(
        f'git clone https://github.com/s4afa451dgf415f/sd-webui-controlnet {params["dir"]}/extensions/sd-wi-controlnet',
        shell=True)
    subprocess.run(
        f'git clone https://github.com/dtlnor/stable-diffusion-webui-localization-zh_CN {params["dir"]}/extensions/stable-diffusion-wi-localization-zh_CN',
        shell=True)
    subprocess.run(
        f'git clone https://github.com/fkunn1326/openpose-editor {params["dir"]}/extensions/openpose-editor',
         shell=True)
    subprocess.run(
        f'git clone https://github.com/s4afa451dgf415f/sd-wi-tagcomplet {params["dir"]}/extensions/a1111-sd-wi-tagcomplete',
        shell=True)
    subprocess.run(
        f'git clone https://github.com/picobyte/stable-diffusion-webui-wd14-tagger {params["dir"]}/extensions/stable-diffusion-wi-wd14-tagger',
        shell=True)
    subprocess.run(
        f'git clone https://github.com/Coyote-A/ultimate-upscale-for-automatic1111 {params["dir"]}/extensions/ultimate-upscale',
        shell=True)
    subprocess.run(
        f'git clone https://github.com/pkuliyi2015/multidiffusion-upscaler-for-automatic1111 {params["dir"]}/extensions/multidiffusion-upscaler',
        shell=True)
    # subprocess.run(
    #     f'git clone https://github.com/nonnonstop/sd-webui-3d-open-pose-editor {params["dir"]}/extensions/sd-wi-3d-open-pose-editor',
    #     shell=True)
    subprocess.run(
        f'git clone https://github.com/hako-mikan/sd-webui-lora-block-weight {params["dir"]}/extensions/sd-wi-lora-block-weight',
        shell=True)
    subprocess.run(
        f'git clone https://github.com/joyjoy0213/tt {params["dir"]}/extensions/tt',
        shell=True)
    subprocess.run(
        f'git clone  https://github.com/joyjoy0213/LCH {params["dir"]}/extensions/LCH',
        shell=True)
    subprocess.run(
        f'git clone  https://github.com/joyjoy0213/sd-com-la {params["dir"]}/extensions/sd-com-la',
        shell=True)
    subprocess.run(
        f'git clone  https://github.com/joyjoy0213/sd-lr {params["dir"]}/extensions/sd-lr',
        shell=True)
    subprocess.run(
        f'git clone  https://github.com/joyjoy0213/iCH {params["dir"]}/extensions/iCH',
        shell=True)
    subprocess.run(
        f'git clone  https://github.com/joyjoy0213/civ {params["dir"]}/extensions/civ',
        shell=True)
    subprocess.run(
        f'git clone  https://github.com/joyjoy0213/-ima {params["dir"]}/extensions/-ima',
        shell=True)
    subprocess.run(
        f'git clone  https://github.com/joyjoy0213/promapt-wb {params["dir"]}/extensions/promapt-wb',
        shell=True)
    subprocess.run(
        f'git clone  https://github.com/joyjoy0213/netdisk {params["dir"]}/extensions/netdisk',
        shell=True)
    subprocess.run(
        f'git clone  https://github.com/joyjoy0213/mn {params["dir"]}/extensions/mn',
        shell=True)
    end_time = time.time()
    print("已克隆git耗时：", end_time - start_time, "秒")
def run_aria2c_download():
    start_time = time.time()
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M {checkpoint_url[params["model"]]} -d {params["dir"]}/models/Stable-diffusion -o {params["model"]}',
        shell=True)
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors/resolve/main/control_v11e_sd15_ip2p_fp16.safetensors -d {params["dir"]}/models/ControlNet -o control_v11e_sd15_ip2p.safetensors',
        shell=True)
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors/resolve/main/control_v11f1p_sd15_depth_fp16.safetensors -d {params["dir"]}/models/ControlNet -o control_v11f1p_sd15_depth.safetensors',
        shell=True)
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors/resolve/main/control_v11p_sd15_canny_fp16.safetensors -d {params["dir"]}/models/ControlNet -o control_v11p_sd15_canny.safetensors',
        shell=True)
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors/resolve/main/control_v11p_sd15_inpaint_fp16.safetensors -d {params["dir"]}/models/ControlNet -o control_v11p_sd15_inpaint.safetensors',
        shell=True)
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors/resolve/main/control_v11p_sd15_lineart_fp16.safetensors -d {params["dir"]}/models/ControlNet -o control_v11f1p_sd15_depth.safetensors',
        shell=True)
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors/resolve/main/control_v11p_sd15_mlsd_fp16.safetensors -d {params["dir"]}/models/ControlNet -o control_v11p_sd15_mlsd.safetensors',
        shell=True)
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors/resolve/main/control_v11p_sd15_normalbae_fp16.safetensors -d {params["dir"]}/models/ControlNet -o control_v11p_sd15_normalbae.safetensors',
        shell=True)
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors/resolve/main/control_v11p_sd15_openpose_fp16.safetensors -d {params["dir"]}/models/ControlNet -o control_v11p_sd15_openpose.safetensors',
        shell=True)
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors/resolve/main/control_v11f1e_sd15_tile_fp16.safetensors -d {params["dir"]}/models/ControlNet -o control_v11f1e_sd15_tile_fp16.safetensors',
        shell=True)
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/lokCX/4x-Ultrasharp/resolve/main/4x-UltraSharp.pth -d {params["dir"]}/models/ESRGAN/ -o 4x-UltraSharp.pth',
        shell=True)
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/konohashinobi4/4xAnimesharp/tree/main -d {params["dir"]}/models/ESRGAN/ -o 4x-AnimeSharp.pth',
        shell=True)
    # 删除文件
    subprocess.run(f'rm {params["dir"]}/extensions/a1111-sd-wi-tagcomplete/tags/danbooru.csv', shell=True)
    
    subprocess.run(
        f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/caocaocoa/joyLora/resolve/main/danbooru.csv -d {params["dir"]}/extensions/a1111-sd-wi-tagcomplete/tags -o danbooru.csv',
        shell=True)
    end_time = time.time()
    print("aria2c完成下载耗时：", end_time - start_time, "秒")


def curl_download():
    start_time = time.time()
    subprocess.run(
        f"curl -Lo '{params['dir']}/models/VAE/vae-ft-mse-840000-ema-pruned.safetensors' https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors",
        shell=True)
    subprocess.run(
        f"curl -Lo '{params['dir']}/models/VAE/kl-f8-anime2.ckpt' https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/4c4f05104055c029ad577c18ac176462f0d1d7c1/vae/kl-f8-anime2.ckpt",
        shell=True)
    subprocess.run(
        f"curl -Lo '{params['dir']}/models/VAE/animevae.pt' https://huggingface.co/swl-models/animvae/resolve/main/animevae.pt",
        shell=True)
    end_time = time.time()
    print("curl完成下载耗时：", end_time - start_time, "秒")


def wget_download():
    start_time = time.time()
    subprocess.run("apt install libunwind8-dev -yqq", shell=True)
    os.environ["LD_PRELOAD"] = "libtcmalloc.so.4"
    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
    subprocess.run("sudo apt-get install sox ffmpeg libcairo2 libcairo2-dev", shell=True)
    end_time = time.time()
    print("wget完成下载耗时：", end_time - start_time, "秒")


def pip_download():
    start_time = time.time()
    # subprocess.run("pip install xformers xformers==0.0.20", shell=True)
    subprocess.run("pip install xformers==0.0.22.post4 --index-url https://download.pytorch.org/whl/cu121", shell=True)
    # all in one和controlnet所需依赖
    subprocess.run("pip install mediapipe svglib fvcore opencv-python>=4.8.0 PyExecJS pathos openai boto3 aliyun-python-sdk-core aliyun-python-sdk-alimt", shell=True)
    for root, dirs, files in os.walk(f"{params['dir']}/extensions"):
        for file in files:
            if file == 'install.py':
                file_path = os.path.join(root, file)
                os.remove(file_path)
                print(f"Deleted {file_path}")
    # subprocess.run("pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2 torchtext==0.15.2+cpu torchdata==0.6.1 --index-url https://download.pytorch.org/whl/cu118", shell=True)
    end_time = time.time()
    print("pip完成下载耗时：", end_time - start_time, "秒")

executor = concurrent.futures.ThreadPoolExecutor(max_workers=5)
task1 = executor.submit(run_git_download)
task2 = executor.submit(run_aria2c_download)
task3 = executor.submit(curl_download)
task4 = executor.submit(wget_download)
task5 = executor.submit(pip_download)
concurrent.futures.wait([ task1,task3, task4, task5])

if os.path.exists(f'{params["dir"]}/embeddings'):
    shutil.rmtree(f'{params["dir"]}/embeddings')
subprocess.run(f'git clone https://huggingface.co/nolanaatama/embeddings {params["dir"]}/embeddings', shell=True)

# 个人插件从云盘的extensions文件夹与VAE文件夹加载
if params["extensions"] == "True":
    if os.path.exists("/content/drive/MyDrive/extensions"):
        subprocess.run(f'rsync -a /content/drive/MyDrive/extensions/* {params["dir"]}/extensions', shell=True)
        print('已加载云盘里的插件')
    if os.path.exists("/content/drive/MyDrive/VAE"):
        subprocess.run(f'rsync -a /content/drive/MyDrive/VAE/* {params["dir"]}/models/VAE', shell=True)
        print('已加载云盘里的VAE')
    if os.path.exists("/content/drive/MyDrive/embeddings"):
        subprocess.run(f'rsync -a /content/drive/MyDrive/embeddings/* {params["dir"]}/embeddings', shell=True)
        print('已加载云盘里的embeddings')
    if os.path.exists("/content/drive/MyDrive/lora"):
        subprocess.run(f'mkdir -p {params["dir"]}/models/Lora', shell=True)
        subprocess.run(f'rsync -a /content/drive/MyDrive/lora/* {params["dir"]}/models/Lora', shell=True)
        print('已加载云盘里的lora')
    if os.path.exists("/content/drive/MyDrive/checkpoint"):
        subprocess.run(f'rsync -a /content/drive/MyDrive/checkpoint/* {params["dir"]}/models/Stable-diffusion',
                       shell=True)
        print('已加载云盘里的Stable-diffusion')