Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
TheLastBen
GitHub Repository: TheLastBen/fast-stable-diffusion
 Path: blob/main/fast_stable_diffusion_ComfyUI.ipynb
539 views
Kernel: Python 3
#@markdown # Connect Google Drive
from google.colab import drive
from IPython.display import clear_output
import ipywidgets as widgets
import os

def inf(msg, style, wdth): inf = widgets.Button(description=msg, disabled=True, button_style=style, layout=widgets.Layout(min_width=wdth));display(inf)
Shared_Drive = "" #@param {type:"string"}
#@markdown - Leave empty if you're not using a shared drive

print("Connecting...")
drive.mount('/content/gdrive')

if Shared_Drive!="" and os.path.exists("/content/gdrive/Shareddrives"):
  mainpth="Shareddrives/"+Shared_Drive
else:
  mainpth="MyDrive"

clear_output()
inf('\u2714 Done','success', '50px')

#@markdown ---

#@markdown # Install/Update ComfyUI repo
from IPython.utils import capture
from IPython.display import clear_output
from subprocess import getoutput
import ipywidgets as widgets
import sys
import fileinput
import os
import time
import base64
import requests
from urllib.request import urlopen, Request
from urllib.parse import urlparse, parse_qs, unquote
from tqdm import tqdm
import six


if not os.path.exists("/content/gdrive"):
  print('Gdrive not connected, using temporary colab storage ...')
  time.sleep(4)
  mainpth="MyDrive"
  !mkdir -p /content/gdrive/$mainpth
  Shared_Drive=""

if Shared_Drive!="" and not os.path.exists("/content/gdrive/Shareddrives"):
  print('Shared drive not detected, using default MyDrive')
  mainpth="MyDrive"

with capture.capture_output() as cap:
  def inf(msg, style, wdth): inf = widgets.Button(description=msg, disabled=True, button_style=style, layout=widgets.Layout(min_width=wdth));display(inf)
  fgitclone = "git clone --depth 1"
  !git clone -q --depth 1 --branch main https://github.com/TheLastBen/diffusers
  %cd /content/gdrive/$mainpth/
  !git clone --depth 1 https://github.com/comfyanonymous/ComfyUI
  !mkdir -p /content/gdrive/$mainpth/sd/cache/
  os.environ['TRANSFORMERS_CACHE']=f"/content/gdrive/{mainpth}/sd/cache"
  os.environ['TORCH_HOME'] = f"/content/gdrive/{mainpth}/sd/cache"

with capture.capture_output() as cap:
  %cd /content/gdrive/$mainpth/ComfyUI
  !git reset --hard
  time.sleep(1)
  !git pull
clear_output()
inf('\u2714 Done','success', '50px')

#@markdown ---

#@markdown # Requirements

print('Installing requirements...')

with capture.capture_output() as cap:
  %cd /content/
  !wget -q -i https://raw.githubusercontent.com/TheLastBen/fast-stable-diffusion/main/Dependencies/A1111.txt
  !dpkg -i *.deb
  !pip uninstall wandb -y
  !tar -C / --zstd -xf gcolabdeps.tar.zst
  !rm *.deb | rm *.zst | rm *.txt

  if not os.path.exists('gdrive/'+mainpth+'/sd/libtcmalloc/libtcmalloc_minimal.so.4'):
    %env CXXFLAGS=-std=c++14
    !wget -q https://github.com/gperftools/gperftools/releases/download/gperftools-2.5/gperftools-2.5.tar.gz && tar zxf gperftools-2.5.tar.gz && mv gperftools-2.5 gperftools
    !wget -q https://github.com/TheLastBen/fast-stable-diffusion/raw/main/AUTOMATIC1111_files/Patch
    %cd /content/gperftools
    !patch -p1 < /content/Patch
    !./configure --enable-minimal --enable-libunwind --enable-frame-pointers --enable-dynamic-sized-delete-support --enable-sized-delete --enable-emergency-malloc; make -j4
    !mkdir -p /content/gdrive/$mainpth/sd/libtcmalloc && cp .libs/libtcmalloc*.so* /content/gdrive/$mainpth/sd/libtcmalloc
    %env LD_PRELOAD=/content/gdrive/$mainpth/sd/libtcmalloc/libtcmalloc_minimal.so.4
    %cd /content
    !rm *.tar.gz Patch && rm -r /content/gperftools
  else:
    %env LD_PRELOAD=/content/gdrive/$mainpth/sd/libtcmalloc/libtcmalloc_minimal.so.4

  !pip install wandb==0.15.12 pydantic==2.10.5 numpy==1.26 scipy==1.15.3 -qq
  !pip install diffusers accelerate transformers av comfyui-frontend-package comfyui-workflow-templates alembic -U -qq
  !rm -r /usr/local/lib/python3.12/dist-packages/tensorflow*
  os.environ['PYTHONWARNINGS'] = 'ignore'
  !sed -i 's@text = _formatwarnmsg(msg)@text =\"\"@g' /usr/lib/python3.12/warnings.py
  !sed -i 's@raise AttributeError(f"module {module!r} has no attribute {name!r}")@@g' /usr/local/lib/python3.12/dist-packages/jax/_src/deprecations.py
  !sed -i 's@globalns, localns, set()@globalns, localns, recursive_guard=set()@g' /usr/local/lib/python3.12/dist-packages/pydantic/typing.py

clear_output()
inf('\u2714 Done','success', '50px')

#@markdown ---

#@markdown # Model Download/Load

import gdown
from gdown.download import get_url_from_gdrive_confirmation
import re

Use_Temp_Storage = False #@param {type:"boolean"}
#@markdown - If not, make sure you have enough space on your gdrive

#@markdown ---

Model_Version = "SDXL" #@param ["SDXL", "1.5", "v1.5 Inpainting", "flux"]

#@markdown Or
MODEL_LINK = "" #@param {type:"string"}


def getsrc(url):
    parsed_url = urlparse(url)
    if parsed_url.netloc == 'civitai.com':
        src='civitai'
    elif parsed_url.netloc == 'drive.google.com':
        src='gdrive'
    elif parsed_url.netloc == 'huggingface.co':
        src='huggingface'
    else:
        src='others'
    return src

src=getsrc(MODEL_LINK)

def get_name(url, gdrive):
    if not gdrive:
        response = requests.get(url, allow_redirects=False)
        if "Location" in response.headers:
            redirected_url = response.headers["Location"]
            quer = parse_qs(urlparse(redirected_url).query)
            if "response-content-disposition" in quer:
                disp_val = quer["response-content-disposition"][0].split(";")
                for vals in disp_val:
                    if vals.strip().startswith("filename="):
                        filenm=unquote(vals.split("=", 1)[1].strip())
                        return filenm.replace("\"","")
    else:
        headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_10_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36"}
        lnk="https://drive.google.com/uc?id={id}&export=download".format(id=url[url.find("/d/")+3:url.find("/view")])
        res = requests.session().get(lnk, headers=headers, stream=True, verify=True)
        res = requests.session().get(get_url_from_gdrive_confirmation(res.text), headers=headers, stream=True, verify=True)
        content_disposition = six.moves.urllib_parse.unquote(res.headers["Content-Disposition"])
        filenm = re.search('attachment; filename="(.*?)"', content_disposition).groups()[0]
        return filenm


def dwn(url, dst, msg):
    file_size = None
    req = Request(url, headers={"User-Agent": "torch.hub"})
    u = urlopen(req)
    meta = u.info()
    if hasattr(meta, 'getheaders'):
        content_length = meta.getheaders("Content-Length")
    else:
        content_length = meta.get_all("Content-Length")
    if content_length is not None and len(content_length) > 0:
        file_size = int(content_length[0])

    with tqdm(total=file_size, disable=False, mininterval=0.5,
              bar_format=msg+' |{bar:20}| {percentage:3.0f}%') as pbar:
        with open(dst, "wb") as f:
            while True:
                buffer = u.read(8192)
                if len(buffer) == 0:
                    break
                f.write(buffer)
                pbar.update(len(buffer))
            f.close()


def sdmdls(ver, Use_Temp_Storage):

  if ver=='1.5':
    if Use_Temp_Storage:
      os.makedirs('/content/temp_models', exist_ok=True)
      model='/content/temp_models/v1-5-pruned-emaonly.safetensors'
    else:
      model='/content/gdrive/'+mainpth+'/ComfyUI/models/checkpoints/v1-5-pruned-emaonly.safetensors'
    link='https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors'
  elif ver=='flux':
    if Use_Temp_Storage:
      os.makedirs('/content/temp_models', exist_ok=True)
      model='/content/temp_models/flux1-dev-fp8.safetensors'
    else:
      model='/content/gdrive/'+mainpth+'/ComfyUI/models/checkpoints/flux1-dev-fp8.safetensors'
    link='https://huggingface.co/lllyasviel/flux1_dev/resolve/main/flux1-dev-fp8.safetensors'
  elif ver=='v1.5 Inpainting':
    if Use_Temp_Storage:
      os.makedirs('/content/temp_models', exist_ok=True)
      model='/content/temp_models/sd-v1-5-inpainting.ckpt'
    else:
      model='/content/gdrive/'+mainpth+'/ComfyUI/models/checkpoints/sd-v1-5-inpainting.ckpt'
    link='https://huggingface.co/runwayml/stable-diffusion-inpainting/resolve/main/sd-v1-5-inpainting.ckpt'
  elif ver=='SDXL':
    if Use_Temp_Storage:
      os.makedirs('/content/temp_models', exist_ok=True)
      model='/content/temp_models/sd_xl_base_1.0.safetensors'
    else:
      model='/content/gdrive/'+mainpth+'/ComfyUI/models/checkpoints/sd_xl_base_1.0.safetensors'
    link='https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors'

  if not os.path.exists(model):
    !gdown --fuzzy -O $model $link
    if os.path.exists(model):
      clear_output()
      inf('\u2714 Done','success', '50px')
    else:
      inf('\u2718 Something went wrong, try again','danger', "250px")
  else:
      clear_output()
      inf('\u2714 Model already exists','primary', '300px')

  return model


if MODEL_LINK != "":

      if src=='civitai':
         modelname=get_name(MODEL_LINK, False)
         if Use_Temp_Storage:
            os.makedirs('/content/temp_models', exist_ok=True)
            model=f'/content/temp_models/{modelname}'
         else:
            model=f'/content/gdrive/{mainpth}/ComfyUI/models/checkpoints/{modelname}'
         if not os.path.exists(model):
            dwn(MODEL_LINK, model, 'Downloading the custom model')
            clear_output()
         else:
            inf('\u2714 Model already exists','primary', '300px')
      elif src=='gdrive':
         modelname=get_name(MODEL_LINK, True)
         if Use_Temp_Storage:
            os.makedirs('/content/temp_models', exist_ok=True)
            model=f'/content/temp_models/{modelname}'
         else:
            model=f'/content/gdrive/{mainpth}/ComfyUI/models/checkpoints/{modelname}'
         if not os.path.exists(model):
            gdown.download(url=MODEL_LINK, output=model, quiet=False, fuzzy=True)
            clear_output()
         else:
            inf('\u2714 Model already exists','primary', '300px')
      else:
         modelname=os.path.basename(MODEL_LINK)
         if Use_Temp_Storage:
            os.makedirs('/content/temp_models', exist_ok=True)
            model=f'/content/temp_models/{modelname}'
         else:
            model=f'/content/gdrive/{mainpth}/ComfyUI/models/checkpoints/{modelname}'
         if not os.path.exists(model):
            gdown.download(url=MODEL_LINK, output=model, quiet=False, fuzzy=True)
            clear_output()
         else:
            inf('\u2714 Model already exists','primary', '700px')

      if os.path.exists(model) and os.path.getsize(model) > 1810671599:
        inf('\u2714 Model downloaded, using the custom model.','success', '300px')
      else:
        !rm model
        inf('\u2718 Wrong link, check that the link is valid','danger', "300px")

else:
  model=sdmdls(Model_Version, Use_Temp_Storage)

#@markdown ---

#@markdown # Download LoRA

LoRA_LINK = "" #@param {type:"string"}

if LoRA_LINK == "":
  inf('\u2714 Nothing to do','primary', '200px')
else:
  os.makedirs('/content/gdrive/'+mainpth+'/ComfyUI/models/loras', exist_ok=True)

  src=getsrc(LoRA_LINK)

  if src=='civitai':
      modelname=get_name(LoRA_LINK, False)
      loramodel=f'/content/gdrive/{mainpth}/ComfyUI/models/loras/{modelname}'
      if not os.path.exists(loramodel):
        dwn(LoRA_LINK, loramodel, 'Downloading the LoRA model '+modelname)
        clear_output()
      else:
        inf('\u2714 Model already exists','primary', '200px')
  elif src=='gdrive':
      modelname=get_name(LoRA_LINK, True)
      loramodel=f'/content/gdrive/{mainpth}/ComfyUI/models/loras/{modelname}'
      if not os.path.exists(loramodel):
        gdown.download(url=LoRA_LINK, output=loramodel, quiet=False, fuzzy=True)
        clear_output()
      else:
        inf('\u2714 Model already exists','primary', '200px')
  else:
      modelname=os.path.basename(LoRA_LINK)
      loramodel=f'/content/gdrive/{mainpth}/ComfyUI/models/loras/{modelname}'
      if not os.path.exists(loramodel):
        gdown.download(url=LoRA_LINK, output=loramodel, quiet=False, fuzzy=True)
        clear_output()
      else:
        inf('\u2714 Model already exists','primary', '200px')

  if os.path.exists(loramodel) :
    inf('\u2714 LoRA downloaded','success', '200px')
  else:
    inf('\u2718 Wrong link, check that the link is valid','danger', "300px")

#@markdown ---

#@markdown # Start ComfyUI
from IPython.utils import capture
import time
import sys
import fileinput
from pyngrok import ngrok, conf
import re
from subprocess import call


Ngrok_Token = "" #@param {type:"string"}

#@markdown - Ngrok token must be entered to be able to run ComfyUI

#ngrok.kill()
#time.sleep(2)
with capture.capture_output() as cap:
  localurl=ngrok.connect(666, pyngrok_config=conf.PyngrokConfig(auth_token=Ngrok_Token) , bind_tls=True).public_url
  !rm -r /content/gdrive/MyDrive/ComfyUI/custom_nodes/.ipynb_checkpoints
call("sed -i 's@^            if verbose:@@' /content/gdrive/MyDrive/ComfyUI/server.py", shell=True)
call("sed -i 's@^                logging.info(\"To see the GUI go to: {}://{}:{}\".format(scheme, address_print, port))@        logging.info(\"\u2714 Connected\");logging.info(\""+localurl+"\")@' /content/gdrive/MyDrive/ComfyUI/server.py", shell=True)
!sed -i 's@https:.*@{localurl}\")@' /content/gdrive/MyDrive/ComfyUI/server.py
!python /content/gdrive/MyDrive/ComfyUI/main.py --listen --port 666 #--preview-method auto