# -*- coding: utf-8 -*-
"""stable_diffusion_1_5_webui.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/github/PinkQween/Discord-AI-bot-advanced/blob/main/stable_diffusion_1_5_webui.ipynb

# AUTOMATIC1111 Stable Diffusion WebUI 1.5 + ChilloutMix Checkpoint

For generating AI

> [@nang](https://github.com/nathan-149)
> References: [@wibus-wee](https://github.com/wibus-wee) [camenduru/stable-diffusion-webui-colab](https://github.com/camenduru/stable-diffusion-webui-colab)
"""

# Commented out IPython magic to ensure Python compatibility.
from IPython.display import display
import ipywidgets as widgets
import requests

endpoint = 'https://civitai.com/api/v1/models'
checkpointName = 'Chilloutmix'
checkpointID = '' #@param {type:"string"}
checkpointURL = '' #@param {type:"string"}

#@markdown ---

#@markdown ## **LoRAs**

#@markdown Korean Doll Likeness `v1.0`
koreanDollLikenessv10 = True #@param {type:"boolean"}
#@markdown Korean Doll Likeness `v1.5`
koreanDollLikenessv15 = True #@param {type:"boolean"}
#@markdown Korean Doll Likeness `v2.0`
koreanDollLikenessv20 = True #@param {type:"boolean"}
#@markdown St. Louis (Luxurious Wheels) (Azur Lane)
stLouis = True #@param {type:"boolean"}
#@markdown Girls' Frontline-OTs-14"lightning"
girlsFrontlineOTs14Lightning = True #@param {type:"boolean"}
#@markdown Yae Miko | Realistic Genshin (Mixed)
yaeMikoRealisticGenshinMixed = True #@param {type:"boolean"}
#@markdown Fashion Girl
fashionGirl = True #@param {type:"boolean"}
#@markdown Cute Girl mix4
cuteGirlMix4 = True #@param {type:"boolean"}
#@markdown More LoRA? (Enter the URL of the lora file, use comma to separate multiple URLs)
moreLoRA = "" #@param {type:"string"}

#@markdown ---

#@markdown ## **Textual Inversion**

#@markdown Ulzzang-6500 (Korean doll aesthetic)
ulzzang6500 = True #@param {type:"boolean"}
#@markdown Pure Eros Face
pureErosFace = True #@param {type:"boolean"}


defaultCheckpoint = {
    'Chilloutmix': '6424',
    'Sunshinemix': '9291',
    'grapefruit_hentai': '2583'
}

loraDownloadIDs = {
  'koreanDollLikenessv10': 'https://huggingface.co/Kanbara/doll-likeness-series/resolve/main/koreanDollLikeness_v10.safetensors',
  'koreanDollLikenessv15': 'https://huggingface.co/Kanbara/doll-likeness-series/resolve/main/koreanDollLikeness_v15.safetensors',
  'koreanDollLikenessv20': 'https://huggingface.co/Kanbara/doll-likeness-series/resolve/main/koreanDollLikeness_v20.safetensors',
  'stLouis': 6669,
  'girlsFrontlineOTs14Lightning': 6525,
  'yaeMikoRealisticGenshinMixed': 8484,
  # 'chilloutMixss': 10850,
  'fashionGirl': {
    "id": 8217,
    "version": "v4.5"
  },
  'cuteGirlMix4': 14171,
}

textualInversionDownloadIDs = {
  'ulzzang6500': 8109,
  'pureErosFace': 4514,
}

downloadIds = []

if checkpointID != '':
    downloadIds = checkpointID.split(',')
if checkpointName != 'others':
        downloadIds.append(defaultCheckpoint[checkpointName])

globalDropdowns = []
globalVerions = []
globalNames = []
globalTexts = []
checkpoints = []
downloadLinks = []
customLoras = []
loraDownloadLinks = []

def text_on_submit(change):
    checkpoints[checkpoints.index(change['old'])] = change['new']

if checkpointURL != '':
    _downloadLinks = checkpointURL.split(',')
    for _downloadLink in _downloadLinks:
        checkpoints.append(_downloadLink.split('/')[-1])
        downloadLinks.append(_downloadLink)
        text = widgets.Text(value=_downloadLink.split('/')[-1], description=_downloadLink.split('/')[-1], disabled=False)
        text.observe(text_on_submit, names='value')
        form = widgets.VBox([text])
        display(form)
if moreLoRA != '':
    _downloadLinks = moreLoRA.split(',')
    for _downloadLink in _downloadLinks:
        customLoras.append(_downloadLink)
        loraDownloadLinks.append(_downloadLink)
        text = widgets.Text(value=_downloadLink.split('/')[-1], description=_downloadLink.split('/')[-1], disabled=False)
        text.observe(text_on_submit, names='value')
        form = widgets.VBox([text])
        display(form)


def showVerionOptions(downloadId):
    res = requests.get(endpoint + '/' + downloadId).json()
    globalNames.append(res['name'])
    versions = res['modelVersions']
    globalVerions.append(versions)
    options = []
    for version in versions:
        options.append(version['files'][0]['name'])
    dropdown = widgets.Dropdown(options=options, description=res['name'])
    globalDropdowns.append(dropdown)
    form = widgets.VBox([dropdown])
    display(form)

for downloadId in downloadIds:
    showVerionOptions(downloadId)

def on_button_clicked(b):
    downloadLink = None
    for dropdown in globalDropdowns:
        checkpoint = dropdown.value
        versions = globalVerions[globalDropdowns.index(dropdown)]
        for version in versions:
            if version['files'][0]['name'] == checkpoint:
                downloadLink = version['files'][0]['downloadUrl']
                break
        if downloadLink is None:
            print('Error: downloadLink not assigned')
            return
        checkpoints.append(checkpoint)
        downloadLinks.append(downloadLink)

    print("已选择模型: " + str(checkpoints) + " <===> " + str(downloadLinks))
    print("自定义 LoRA (仅显示填入 moreLoRA 中的 LoRA): " + str(customLoras) + " <===> " + str(loraDownloadLinks))
#     %store checkpoints
#     %store downloadLinks


button = widgets.Button(description='Use it!')
button.on_click(on_button_clicked)
display(button)

# Commented out IPython magic to ensure Python compatibility.
#@title 2. Check GPU & Dev Environment

import os, subprocess
paperspace_m4000 = False
#@markdown Paperspace platform?
isPaperspace = False #@param {type:"boolean"}
appPrefix = 'stable' + '-' + 'diffusion' + '-' + 'webui' # app prefix
sdModelPrefix = 'Stable' + '-' + 'diffusion'
ecosystemUIPrefix = "sd-webui"

try:
    subprocess.run(['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'], stdout=subprocess.PIPE)
    if 'M4000' in subprocess.run(['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'], stdout=subprocess.PIPE).stdout.decode('utf-8'):
        print("WARNING: You're using Quadro M4000 GPU，xformers won't work。")
        paperspace_m4000 = True
        isPaperspace = True
    else:
        print("Your GPU is suitable - " + subprocess.run(['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'], stdout=subprocess.PIPE).stdout.decode('utf-8') + "。")
        print("Platform: Paperspace" if isPaperspace else "Platform: Colab")
except:
    print("No GPU appears to be available. Please check your runtime type")
    exit()

rootDir = isPaperspace and '/tmp' or '/content'
diffusionWebUIInstalled = os.path.exists(rootDir + '/' + appPrefix)
# %store rootDir
# %store paperspace_m4000
# %store isPaperspace
# %store diffusionWebUIInstalled

# Commented out IPython magic to ensure Python compatibility.
import requests
#@title 3. Install dependencies and extensions

#@markdown ## **Extensions**

#@markdown xformer
xformersInstall = True #@param {type:"boolean"}
#@markdown ControlNet
controlNetExtension = False #@param {type:"boolean"}
#@markdown OpenPose Editor
openPoseExtension = False #@param {type:"boolean"}
#@markdown Civitai Browser
civitaiBrowserExtension = False #@param {type:"boolean"}
#@markdown HuggingFace
huggingFaceExtension = False #@param {type:"boolean"}
#@markdown Images Browser
imagesBrowserExtension = False #@param {type:"boolean"}
#@markdown Additional Networks
additionalNetworksExtension = True #@param {type:"boolean"}
#@markdown Deforum
deforumExtension = False #@param {type:"boolean"}
#@markdown Kohya sd-scripts
kohyaExtension = False #@param {type:"boolean"}
#@markdown DreamBooth
dreamBoothExtension = False #@param {type:"boolean"}

#@markdown ---

#@markdown ## **Others**

#@markdown Styles CSV File URL
stylesCSVFileURL = "" #@param {type:"string"}
#@markdown Use [Cloudflared](https://github.com/cloudflare/cloudflared) instead of Gradio share
useCloudflared = False #@param {type:"boolean"}

#@markdown ---

#@markdown ## **Startup Options**

#@markdown API Support
apiSupport = True #@param {type:"boolean"}
#@markdown Cors Allow Origins
corsAllowOrigins = "*" #@param {type:"string"}

def getLatestModelDownloadURL(id):
  try:
    if type(id) == int:
      res = requests.get(endpoint + '/' + str(id)).json()
      latest = res['modelVersions'][0]
      downloadLink = latest['files'][0]['downloadUrl']
      name = latest['files'][0]['name']
      return {
        'url': downloadLink,
        'name': name
      }
    else:
      return {
        'url': id,
        'name': id.split('/')[-1]
      }
  except:
    print("Lora model " + str(id) + " not found. Skip.")
    return None

def getSpecificModelDownloadURL(id, version):
  try:
    if type(id) == int:
      res = requests.get(endpoint + '/' + str(id)).json()
      for modelVersion in res['modelVersions']:
        if modelVersion['name'] == version:
          # if modelVersion["baseModel"] != "SD 1.5":
          #   print("Lora model " + str(id) + " is not SD 1.5, may not work. Skip.")
          #   return None
          downloadLink = modelVersion['files'][0]['downloadUrl']
          name = modelVersion['files'][0]['name']
          return {
            'url': downloadLink,
            'name': name
          }
    else:
      return {
        'url': id,
        'name': id.split('/')[-1]
      }
  except:
    print("Lora model " + str(id) + " version " + version + " not found. Skip.")
    return None

def getLoraDownloadURLs():
  downloadURLs = []
  for key in loraDownloadIDs:
    if not eval(key): # skip if not selected
      continue
    if type(loraDownloadIDs[key]) is int:
      downloadURLs.append(getLatestModelDownloadURL(loraDownloadIDs[key]))
    elif type(loraDownloadIDs[key]) is dict: # {'id': 123, 'version': 'v1.0'}
      downloadURLs.append(getSpecificModelDownloadURL(loraDownloadIDs[key]['id'], loraDownloadIDs[key]['version']))
    elif type(loraDownloadIDs[key]) is str: # url
      downloadURLs.append({ 'url': loraDownloadIDs[key], 'name': loraDownloadIDs[key].split('/')[-1] })
  downloadURLs = [x for x in downloadURLs if x is not None] # remove None
  for custom in customLoras:
    downloadURLs.append({ 'url': loraDownloadLinks[customLoras.index(custom)], 'name': custom })
  return downloadURLs

def getTextualInversionDownloadURLs():
  downloadURLs = []
  for key in textualInversionDownloadIDs:
    if not eval(key): # skip if not selected
      continue
    if type(textualInversionDownloadIDs[key]) is int:
      downloadURLs.append(getLatestModelDownloadURL(textualInversionDownloadIDs[key]))
    elif type(textualInversionDownloadIDs[key]) is dict: # {'id': 123, 'version': 'v1.0'}
      downloadURLs.append(getSpecificModelDownloadURL(textualInversionDownloadIDs[key]['id'], textualInversionDownloadIDs[key]['version']))
    elif type(textualInversionDownloadIDs[key]) is str: # url
      downloadURLs.append({ 'url': textualInversionDownloadIDs[key], 'name': textualInversionDownloadIDs[key].split('/')[-1] })
  downloadURLs = [x for x in downloadURLs if x is not None]
  return downloadURLs


loraDownloadURLs = getLoraDownloadURLs()
textualInversionDownloadURLs = getTextualInversionDownloadURLs()


# %store -r paperspace_m4000
# %store -r isPaperspace
# %store -r rootDir
# %store -r checkpoints
# %store -r downloadLinks
# %store -r diffusionWebUIInstalled

import subprocess

!apt-get -y install -qq aria2
ariaInstalled = False

try:
    subprocess.run(['aria2c', '--version'], stdout=subprocess.PIPE)
    ariaInstalled = True
except:
    pass

!pip install -q torch==2.0.0+cu118 torchvision==0.15.1+cu118 torchaudio==2.0.1+cu118 torchtext==0.15.1 torchdata==0.6.0 --extra-index-url https://download.pytorch.org/whl/cu118 -U
!pip install -q xformers==0.0.19 triton==2.0.0 -U

!git clone https://github.com/AUTOMATIC1111/{appPrefix} {rootDir}/{appPrefix}
!git clone https://huggingface.co/embed/negative {rootDir}/{appPrefix}/embeddings/negative
!git clone https://huggingface.co/embed/lora {rootDir}/{appPrefix}/models/Lora/positive
!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/embed/upscale/resolve/main/4x-UltraSharp.pth -d {rootDir}/{appPrefix}/models/ESRGAN -o 4x-UltraSharp.pth
!git clone https://github.com/Bing-su/{ecosystemUIPrefix}-tunnels {rootDir}/{appPrefix}/extensions/{ecosystemUIPrefix}-tunnels
!git clone https://github.com/thomasasfk/{ecosystemUIPrefix}-aspect-ratio-helper {rootDir}/{appPrefix}/extensions/{ecosystemUIPrefix}-aspect-ratio-helper
!wget https://raw.githubusercontent.com/camenduru/{appPrefix}-scripts/main/run_n_times.py -O {rootDir}/{appPrefix}/scripts/run_n_times.py
if deforumExtension:
  !git clone https://github.com/deforum-art/deforum-for-automatic1111-webui {rootDir}/{appPrefix}/extensions/deforum-for-automatic1111-webui
if imagesBrowserExtension:
  !git clone https://github.com/AlUlkesh/{appPrefix}-images-browser {rootDir}/{appPrefix}/extensions/{appPrefix}-images-browser
if huggingFaceExtension:
  !git clone https://github.com/camenduru/{appPrefix}-huggingface {rootDir}/{appPrefix}/extensions/{appPrefix}-huggingface
if civitaiBrowserExtension:
  !git clone https://github.com/Vetchems/sd-civitai-browser {rootDir}/{appPrefix}/extensions/sd-civitai-browser
if openPoseExtension:
  !git clone https://github.com/fkunn1326/openpose-editor {rootDir}/{appPrefix}/extensions/openpose-editor
if controlNetExtension:
  !git clone https://github.com/Mikubill/{ecosystemUIPrefix}-controlnet {rootDir}/{appPrefix}/extensions/{ecosystemUIPrefix}-controlnet
if additionalNetworksExtension:
  !git clone https://github.com/kohya-ss/{ecosystemUIPrefix}-additional-networks {rootDir}/{appPrefix}/extensions/{ecosystemUIPrefix}-additional-networks
if kohyaExtension:
  !git clone https://github.com/ddpn08/kohya-sd-scripts-webui.git {rootDir}/{appPrefix}/extensions/kohya-sd-scripts-webui
if dreamBoothExtension:
  !git clone https://github.com/d8ahazard/sd_dreambooth_extension {rootDir}/{appPrefix}/extensions/sd_dreambooth_extension

if isPaperspace:
#   %cd /{appPrefix}
else:
#   %cd {rootDir}/{appPrefix}


webuiControlNetModels = [
  "https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_canny-fp16.safetensors",
  "https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_depth-fp16.safetensors",
  "https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_hed-fp16.safetensors",
  "https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_mlsd-fp16.safetensors",
  "https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_normal-fp16.safetensors",
  "https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_openpose-fp16.safetensors",
  "https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_scribble-fp16.safetensors",
  "https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/control_seg-fp16.safetensors",
]
annotatorLink = [
  "https://huggingface.co/ckpt/ControlNet/resolve/main/hand_pose_model.pth",
  "https://huggingface.co/ckpt/ControlNet/resolve/main/body_pose_model.pth",
  "https://huggingface.co/ckpt/ControlNet/resolve/main/dpt_hybrid-midas-501f0c75.pt",
  "https://huggingface.co/ckpt/ControlNet/resolve/main/mlsd_large_512_fp32.pth",
  "https://huggingface.co/ckpt/ControlNet/resolve/main/mlsd_tiny_512_fp32.pth",
  "https://huggingface.co/ckpt/ControlNet/resolve/main/network-bsds500.pth",
  "https://huggingface.co/ckpt/ControlNet/resolve/main/upernet_global_small.pth",
]

def ariaDownload(downloadLink, checkpoint, path):
  if (type(downloadLink) == list and type(checkpoint) == list):
    for i in downloadLink:
      !aria2c --console-log-level=error -c -x 16 -s 16 -k 1M {i} -d {path} -o {checkpoint[downloadLink.index(i)]}
  else:
    !aria2c --console-log-level=error -c -x 16 -s 16 -k 1M {downloadLink} -d {path} -o {checkpoint}
def wgetDownload(downloadLink, checkpoint, path):
  if (type(downloadLink) == list and type(checkpoint) == list):
    for i in downloadLink:
      !wget -c {i} -P {path} -O {checkpoint[downloadLink.index(i)]}
  else:
    !wget -c {downloadLink} -P {path} -O {checkpoint}
def autoDetectDownload(downloadLink, checkpoint, path):
  if ariaInstalled:
    ariaDownload(downloadLink, checkpoint, path)
  else:
    wgetDownload(downloadLink, checkpoint, path)

autoDetectDownload(downloadLinks, checkpoints, rootDir + "/" + appPrefix + "/models/" + sdModelPrefix)
if controlNetExtension:
  for model in webuiControlNetModels:
    autoDetectDownload(model, model.split('/')[-1], rootDir + "/" + appPrefix + "/extensions/" + ecosystemUIPrefix + "-controlnet/models")
  for model in annotatorLink:
    autoDetectDownload(model, model.split('/')[-1], rootDir + "/" + appPrefix + "/extensions/" + ecosystemUIPrefix + "-controlnet/annotator")
for model in loraDownloadURLs:
  autoDetectDownload(model["url"], model["name"], rootDir + "/" + appPrefix + "/models/Lora")
for model in textualInversionDownloadURLs:
  autoDetectDownload(model["url"], model["name"], rootDir + "/" + appPrefix + "/embeddings")
if stylesCSVFileURL:
  # autoDetectDownload(stylesCSVFileURL, "styles.csv", rootDir + "/{appPrefix}")
  autoDetectDownload(stylesCSVFileURL, "styles.csv", rootDir + "/" + appPrefix)

if additionalNetworksExtension:
  !ln -s {rootDir}/{appPrefix}/models/Lora {rootDir}/{appPrefix}/extensions/{ecosystemUIPrefix}-additional-networks/models/lora


diffusionWebUIInstalled = True
# %store diffusionWebUIInstalled

# %cd {rootDir}/{appPrefix}
!sed -i -e 's/\"sd_model_checkpoint\"\,/\"sd_model_checkpoint\,sd_vae\,CLIP_stop_at_last_layers\"\,/g' {rootDir}/{appPrefix}/modules/shared.py

if dreamBoothExtension:
  !export REQS_FILE="./extensions/sd_dreambooth_extension/requirements.txt"

backgroundProcess = ''

if useCloudflared:
   backgroundProcess = 'cloudflared tunnel --url 127.0.0.1:7860'
  # !cloudflared tunnel --url 127.0.0.1:7860
if backgroundProcess:
   backgroundProcess = backgroundProcess + ' &'

if useCloudflared:
  !echo "Installing cloudflared"
  !curl -Lo /usr/bin/cloudflared https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64 && chmod +x /usr/bin/cloudflared

#@title Download Chilloutmix Checkpoint

checkpoint = 'chilloutmix.safetensors' #@param ["chilloutmix.safetensors"]

downloadLink = 'https://huggingface.co/naonovn/chilloutmix_NiPrunedFp32Fix/resolve/main/chilloutmix_NiPrunedFp32Fix.safetensors' #@param


!wget -c {downloadLink} -O /content/stable-diffusion-webui/models/Stable-diffusion/{checkpoint}
!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/ckpt/sd14/resolve/main/sd-v1-4.ckpt -d /content/stable-diffusion-webui/models/Stable-diffusion -o sd-v1-4.ckpt

#@title Download More Specific Loras

loraLinks = dict((
    ('koreanDollLikeness_v15.safetensors', 'https://civitai.com/api/download/models/31284'),
    ('koreanDollLikeness_v20.safetensors', 'https://civitai.com/api/download/models/31284'),
    ('xswltry1.safetensors', 'https://civitai.com/api/download/models/29131'),
    ('liyuuLora_liyuuV1.safetensors', 'https://civitai.com/api/download/models/11885'),
    ('aiBeautyIthlinni_ithlinniV1.safetensors', 'https://civitai.com/api/download/models/19671'),
    ('Cute_girl_mix4.safetensors', 'https://civitai.com/api/download/models/16677'),
    ('breastinclassBetter_v141.safetensors', 'https://civitai.com/api/download/models/23250'),
    ('chilloutmixss_xss10.safetensors', 'https://huggingface.co/HankChang/chilloutmixss_xss10/resolve/main/chilloutmixss_xss10.safetensors'),
    ('legspread10.safetensors', 'https://civitai.com/api/download/models/29760'),
))


for lora, link in loraLinks.items():
    print('\nKey: %s' % lora)
    print('Value: %s' % link)
    !wget -c {link} -O /content/stable-diffusion-webui/models/Lora/{lora}

# Commented out IPython magic to ensure Python compatibility.
# Restore variables if needed
# %store -r paperspace_m4000
# %store -r rootDir
# %store -r diffusionWebUIInstalled

# Navigate to application directory
# %cd {rootDir}/{appPrefix}

# Modify shared.py if necessary
!sed -i -e 's/\"sd_model_checkpoint\"\,/\"sd_model_checkpoint\,sd_vae\,CLIP_stop_at_last_layers\"\,/g' {rootDir}/{appPrefix}/modules/shared.py

# Launch Python application
!{backgroundProcess} python launch.py --enable-insecure-extension-access {'' if paperspace_m4000 and not xformersInstall else '--xformers'} {'--api --cors-allow-origins "*" --listen' if corsAllowOrigins else ''} {'--share' if not backgroundProcess else ''} --gradio-queue --api