123.py

import argparse
import binascii
import sys
import os
import json
import base64
import importlib.util
from IPython import get_ipython
from IPython.display import clear_output
from google.colab import drive
import tensorflow as tf
import pickle
import re
import subprocess
import shlex
import shutil

# 参数解析
def str_to_bool(value):
    if value.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif value.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Invalid boolean value: {}'.format(value))

parser = argparse.ArgumentParser()
parser.add_argument('--UI', type=str)
parser.add_argument('--Version', type=str)
parser.add_argument('--ControlNet', type=str)
parser.add_argument('--Drive_Map', type=str_to_bool)
parser.add_argument('--Key_words', type=str_to_bool)
parser.add_argument('--colabpath', type=str)  # 新增用于传递第一段代码的colabpath的参数
args = parser.parse_args()

UI = args.UI
Version = args.Version  
ControlNet = args.ControlNet
Drive_Map = args.Drive_Map
Key_words = args.Key_words
colabpath = args.colabpath  # 提取colabpath参数

# 初始化环境
sd0 = "s" + "t" + "a" + "b" + "l" + "e" + "-" + "d" + "i" + "f" + "f" + "u" + "s" + "i" + "o" + "n" + "-" + "w" + "e" + "b" + "u"
sd = sd0 + 'i'
wi = 'w' + 'e' + 'b' + 'u' + 'i'
wb = f'/content/{sd}'  
gwb = f'/content/drive/MyDrive/{sd}'

get_ipython().run_line_magic('cd', '/content')
get_ipython().run_line_magic('env', 'TF_CPP_MIN_LOG_LEVEL=1')

print("TensorFlow version:", tf.__version__)
if tf.test.gpu_device_name():
    drive.mount('/content/drive')
else:
    raise Exception("\n请在《代码执行程序》-《更改运行时类型》-设置为GPU~")

# 云盘操作    
def cloudDriveSync(cloudPath, localPath='', sync=False):
    if not os.path.exists(cloudPath):
        get_ipython().system(f'mkdir {cloudPath}')
    if not sync:
        return        
    get_ipython().system(f'rm -rf {localPath}')
    get_ipython().system(f'ln -s {cloudPath} {localPath}')

def initCloudDrive():
    cloudDriveSync(f'{gwb}')
    cloudDriveSync(f'{gwb}/Config')
    cloudDriveSync(f'{gwb}/Models', f'{wb}/models/Stable-diffusion', Drive_Map)
    cloudDriveSync(f'{gwb}/Lora', f'{wb}/models/Lora', Drive_Map)
    cloudDriveSync(f'{gwb}/LyCORIS', f'{wb}/models/LyCORIS', Drive_Map)
    cloudDriveSync(f'{gwb}/hypernetworks', f'{wb}/models/hypernetworks', Drive_Map)
    cloudDriveSync(f'{gwb}/Vae', f'{wb}/models/VAE', Drive_Map)
    cloudDriveSync(f'{gwb}/Outputs', f'{wb}/outputs', Drive_Map)

    if not os.path.exists(f'{gwb}/Config/config.json'):
        get_ipython().system(f'wget -O {gwb}/Config/config.json "https://huggingface.co/wageguagua/main/raw/main/config.json"')
        
    if not os.path.exists(f'{gwb}/styles.csv'):
        get_ipython().system(f'wget -O {gwb}/styles.csv "https://huggingface.co/wageguagua/main/raw/main/styles.csv"')

# 处理colab文件        
def process_colab(colabpath):

    extensionlines = []

    with open(colabpath, 'r', encoding='utf-8') as f:
        pattern = r"(?<!\S)https://github.com/camenduru/stable-diffusion-webui(?!\S)"
        for line in f:
            stripped_line = line.strip()
            if stripped_line.startswith('"'):
                stripped_line = stripped_line[1:]
            if stripped_line.startswith('!'):
                stripped_line = stripped_line[1:]
            if stripped_line.endswith('\\n",'):
                stripped_line = stripped_line[:-4]
            if not startcapture:
                if re.search(pattern, stripped_line):
                    startcapture = True
            else:
                camendururepo = 'camenduru/stable-diffusion-webui'
                if camendururepo in stripped_line and not '/content/volatile-concentration-localux' in stripped_line:
                    if camendururepo in stripped_line and (stripped_line.find(camendururepo) + len(camendururepo) == len(stripped_line) or stripped_line[stripped_line.find(camendururepo) + len(camendururepo)] in [' ', '\n']):
                        stripped_line += ' /content/volatile-concentration-localux'
                if stripped_line.startswith("git clone") and "https://github.com" in stripped_line:
                    commandtoappend = stripped_line.replace('/content/stable-diffusion-webui', '/content/volatile-concentration-localux')
                    extensionlines.append(commandtoappend)

    return extensionlines

# 下载git库
def gitDownload(url, localPath):
    if os.path.exists(localPath):
        return
    get_ipython().system(f'git clone {url} {localPath}')

# 安装插件  
def installAdditional():

    # 处理colab文件
    extensionlines = process_colab(colabpath)
    
    for line in extensionlines:
        if 'sd-webui' in line:
            line = line.replace('sd-webui', sd)
        get_ipython().system(line)
      
    # 下载git库
    gitDownload(f'https://huggingface.co/embed/negative',f'{wb}/embeddings/negative')
    gitDownload(f'https://huggingface.co/embed/lora',f'{wb}/models/Lora/positive')
    
    # 云盘模型
    if len(os.listdir(f"{gwb}/Models")) == 0:
        get_ipython().system(f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/coco233/colab/resolve/main/ZR888.safetensors -d {wb}/models/Stable-diffusion -o zr888.safetensors')
        get_ipython().system(f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/wageguagua/sdmodels/resolve/main/long9.safetensors -d {wb}/models/Stable-diffusion -o long9.safetensors')

    if len(os.listdir(f"{gwb}/Vae")) == 0:
        get_ipython().system(f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors -d {wb}/models/VAE -o vae-ft-mse-840000-ema-pruned.safetensors')
        
# 初始化本地环境
def initLocal():

    get_ipython().system(f'apt -y update -qq')
    get_ipython().system(f'wget https://github.com/camenduru/gperftools/releases/download/v1.0/libtcmalloc_minimal.so.4 -O /content/libtcmalloc_minimal.so.4')
    get_ipython().run_line_magic('env', 'LD_PRELOAD=/content/libtcmalloc_minimal.so.4')

    get_ipython().system(f'apt -y install -qq aria2 libcairo2-dev pkg-config python3-dev')
    get_ipython().system(f'pip install -q torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2+cu118 torchtext==0.15.2 torchdata==0.6.1 --extra-index-url https://download.pytorch.org/whl/cu118 -U')
    get_ipython().system(f'pip install -q xformers==0.0.20 triton==2.0.0 gradio_client==0.2.7 -U')

    if Version == "A1111":
        get_ipython().system(f'git clone -b master https://github.com/AUTOMATIC1111/{sd} {wb}')
    elif Version == "V2.5":
        get_ipython().system(f'git clone -b v2.5 https://github.com/camenduru/{sd} {wb}') 
    elif Version == "V2.4":
        get_ipython().system(f'git clone -b v2.4 https://github.com/camenduru/{sd} {wb}')

    get_ipython().system(f'git -C {wb}/repositories/stable-diffusion-stability-ai reset --hard')
    
    initCloudDrive()

    installAdditional()

    get_ipython().system(f'rm -f {wb}/config.json')
    get_ipython().system(f'ln -s {gwb}/Config/config.json {wb}/config.json')

    if os.path.exists(f"{gwb}/styles.csv"):
        get_ipython().system(f'ln -s {gwb}/styles.csv {wb}/styles.csv')

# 运行
def run(script):
    if script == 'webui.py':
        script = f'{wi}.py'
        
    clear_output()
    get_ipython().run_line_magic('cd', f'{wb}')
    get_ipython().system(f'python {script} --listen --xformers --enable-insecure-extension-access --theme dark --gradio-queue --disable-console-progressbars --multiple --api --cors-allow-origins=*')
    
# 主入口
if __name__ == '__main__':
    if os.path.exists(f'{wb}'):
        run('webui.py')
    else:
        initLocal()
        execute_first_segment_code(colabpath)  # 调用执行第一段代码的函数
        run('launch.py')