from IPython.display import clear_output from subprocess import call, getoutput, Popen from IPython.display import display import ipywidgets as widgets import io from PIL import Image, ImageDraw, ImageOps import fileinput import time import os from os import listdir from os.path import isfile import random import sys from io import BytesIO import requests from collections import defaultdict from math import log, sqrt import numpy as np import sys import fileinput from subprocess import check_output import six import base64 import re from urllib.parse import urlparse, parse_qs, unquote import urllib.request from urllib.request import urlopen, Request import tempfile from tqdm import tqdm def Deps(force_reinstall): if not force_reinstall and os.path.exists('/usr/local/lib/python3.10/dist-packages/safetensors'): ntbks() call('pip install --root-user-action=ignore --disable-pip-version-check -qq diffusers==0.18.1', shell=True, stdout=open('/dev/null', 'w')) print('Modules and notebooks updated, dependencies already installed') os.environ['TORCH_HOME'] = '/workspace/cache/torch' os.environ['PYTHONWARNINGS'] = 'ignore' else: call('pip install --root-user-action=ignore --disable-pip-version-check --no-deps -qq gdown PyWavelets numpy==1.23.5 accelerate==0.12.0 --force-reinstall', shell=True, stdout=open('/dev/null', 'w')) ntbks() if os.path.exists('deps'): call("rm -r deps", shell=True) if os.path.exists('diffusers'): call("rm -r diffusers", shell=True) call('mkdir deps', shell=True) if not os.path.exists('cache'): call('mkdir cache', shell=True) os.chdir('deps') dwn("https://huggingface.co/TheLastBen/dependencies/resolve/main/rnpddeps-t2.tar.zst", "/workspace/deps/rnpddeps-t2.tar.zst", "Installing dependencies") call('tar -C / --zstd -xf rnpddeps-t2.tar.zst', shell=True, stdout=open('/dev/null', 'w')) call("sed -i 's@~/.cache@/workspace/cache@' /usr/local/lib/python3.10/dist-packages/transformers/utils/hub.py", shell=True) os.chdir('/workspace') call('pip install --root-user-action=ignore --disable-pip-version-check -qq diffusers==0.18.1', shell=True, stdout=open('/dev/null', 'w')) call("git clone --depth 1 -q --branch main https://github.com/TheLastBen/diffusers", shell=True, stdout=open('/dev/null', 'w')) #call('pip install --root-user-action=ignore --disable-pip-version-check -qq gradio==3.41.0', shell=True, stdout=open('/dev/null', 'w')) call("rm -r deps", shell=True) os.chdir('/workspace') os.environ['TORCH_HOME'] = '/workspace/cache/torch' os.environ['PYTHONWARNINGS'] = 'ignore' call("sed -i 's@text = _formatwarnmsg(msg)@text =\"\"@g' /usr/lib/python3.10/warnings.py", shell=True) clear_output() done() 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 ntbks(): os.chdir('/workspace') if not os.path.exists('Latest_Notebooks'): call('mkdir Latest_Notebooks', shell=True) else: call('rm -r Latest_Notebooks', shell=True) call('mkdir Latest_Notebooks', shell=True) os.chdir('/workspace/Latest_Notebooks') call('wget -q -i https://huggingface.co/datasets/TheLastBen/RNPD/raw/main/Notebooks.txt', shell=True) call('rm Notebooks.txt', shell=True) os.chdir('/workspace') def done(): done = widgets.Button( description='Done!', disabled=True, button_style='success', tooltip='', icon='check' ) display(done) def mdlvxl(): os.chdir('/workspace') if os.path.exists('stable-diffusion-XL') and not os.path.exists('/workspace/stable-diffusion-XL/unet/diffusion_pytorch_model.safetensors'): call('rm -r stable-diffusion-XL', shell=True) if not os.path.exists('stable-diffusion-XL'): print('Downloading SDXL model...') call('mkdir stable-diffusion-XL', shell=True) os.chdir('stable-diffusion-XL') call('git init', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w')) call('git lfs install --system --skip-repo', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w')) call('git remote add -f origin https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w')) call('git config core.sparsecheckout true', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w')) call('echo -e "\nscheduler\ntext_encoder\ntext_encoder_2\ntokenizer\ntokenizer_2\nunet\nvae\nfeature_extractor\nmodel_index.json\n!*.safetensors\n!*.bin\n!*.onnx*\n!*.xml" > .git/info/sparse-checkout', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w')) call('git pull origin main', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w')) dwn('https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/text_encoder/model.safetensors', 'text_encoder/model.safetensors', '1/4') dwn('https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/text_encoder_2/model.safetensors', 'text_encoder_2/model.safetensors', '2/4') dwn('https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/vae/diffusion_pytorch_model.safetensors', 'vae/diffusion_pytorch_model.safetensors', '3/4') dwn('https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/unet/diffusion_pytorch_model.safetensors', 'unet/diffusion_pytorch_model.safetensors', '4/4') call('rm -r .git', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w')) os.chdir('/workspace') clear_output() while not os.path.exists('/workspace/stable-diffusion-XL/unet/diffusion_pytorch_model.safetensors'): print('Invalid HF token, make sure you have access to the model') time.sleep(8) if os.path.exists('/workspace/stable-diffusion-XL/unet/diffusion_pytorch_model.safetensors'): print('Using SDXL model') else: print('Using SDXL model') call("sed -i 's@\"force_upcast.*@@' /workspace/stable-diffusion-XL/vae/config.json", shell=True) def downloadmodel_hfxl(Path_to_HuggingFace): os.chdir('/workspace') if os.path.exists('stable-diffusion-custom'): call("rm -r stable-diffusion-custom", shell=True) clear_output() if os.path.exists('Fast-Dreambooth/token.txt'): with open("Fast-Dreambooth/token.txt") as f: token = f.read() authe=f'https://USER:{token}@' else: authe="https://" clear_output() call("mkdir stable-diffusion-custom", shell=True) os.chdir("stable-diffusion-custom") call("git init", shell=True) call("git lfs install --system --skip-repo", shell=True) call('git remote add -f origin '+authe+'huggingface.co/'+Path_to_HuggingFace, shell=True) call("git config core.sparsecheckout true", shell=True) call('echo -e "\nscheduler\ntext_encoder\ntokenizer\nunet\nvae\nfeature_extractor\nmodel_index.json\n!*.fp16.safetensors" > .git/info/sparse-checkout', shell=True) call("git pull origin main", shell=True) if os.path.exists('unet/diffusion_pytorch_model.safetensors'): call("rm -r .git", shell=True) os.chdir('/workspace') clear_output() done() while not os.path.exists('/workspace/stable-diffusion-custom/unet/diffusion_pytorch_model.safetensors'): print('Check the link you provided') os.chdir('/workspace') time.sleep(5) def downloadmodel_link_xl(MODEL_LINK): import wget import gdown from gdown.download import get_url_from_gdrive_confirmation 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(r"filename\*=UTF-8''(.*)", content_disposition).groups()[0].replace(os.path.sep, "_") return filenm if src=='civitai': modelname=get_name(MODEL_LINK, False) elif src=='gdrive': modelname=get_name(MODEL_LINK, True) else: modelname=os.path.basename(MODEL_LINK) os.chdir('/workspace') if src=='huggingface': dwn(MODEL_LINK, modelname,'Downloading the Model') else: call("gdown --fuzzy " +MODEL_LINK+ " -O "+modelname, shell=True) if os.path.exists(modelname): if os.path.getsize(modelname) > 1810671599: print('Converting to diffusers...') call('python /workspace/diffusers/scripts/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path '+modelname+' --dump_path stable-diffusion-custom --from_safetensors', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w')) if os.path.exists('stable-diffusion-custom/unet/diffusion_pytorch_model.bin'): os.chdir('/workspace') clear_output() done() else: while not os.path.exists('stable-diffusion-custom/unet/diffusion_pytorch_model.bin'): print('Conversion error') os.chdir('/workspace') time.sleep(5) else: while os.path.getsize(modelname) < 1810671599: print('Wrong link, check that the link is valid') os.chdir('/workspace') time.sleep(5) def downloadmodel_path_xl(MODEL_PATH): import wget os.chdir('/workspace') clear_output() if os.path.exists(str(MODEL_PATH)): print('Converting to diffusers...') call('python /workspace/diffusers/scripts/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path '+MODEL_PATH+' --dump_path stable-diffusion-custom --from_safetensors', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w')) if os.path.exists('stable-diffusion-custom/unet/diffusion_pytorch_model.bin'): clear_output() done() while not os.path.exists('stable-diffusion-custom/unet/diffusion_pytorch_model.bin'): print('Conversion error') os.chdir('/workspace') time.sleep(5) else: while not os.path.exists(str(MODEL_PATH)): print('Wrong path, use the file explorer to copy the path') os.chdir('/workspace') time.sleep(5) def dls_xlf(Path_to_HuggingFace, MODEL_PATH, MODEL_LINK): os.chdir('/workspace') if Path_to_HuggingFace != "": downloadmodel_hfxl(Path_to_HuggingFace) MODEL_NAMExl="/workspace/stable-diffusion-custom" elif MODEL_PATH !="": downloadmodel_path_xl(MODEL_PATH) MODEL_NAMExl="/workspace/stable-diffusion-custom" elif MODEL_LINK !="": downloadmodel_link_xl(MODEL_LINK) MODEL_NAMExl="/workspace/stable-diffusion-custom" else: mdlvxl() MODEL_NAMExl="/workspace/stable-diffusion-XL" return MODEL_NAMExl def sess_xl(Session_Name, MODEL_NAMExl): import gdown import wget os.chdir('/workspace') PT="" while Session_Name=="": print('Input the Session Name:') Session_Name=input("") Session_Name=Session_Name.replace(" ","_") WORKSPACE='/workspace/Fast-Dreambooth' INSTANCE_NAME=Session_Name OUTPUT_DIR="/workspace/models/"+Session_Name SESSION_DIR=WORKSPACE+"/Sessions/"+Session_Name INSTANCE_DIR=SESSION_DIR+"/instance_images" CAPTIONS_DIR=SESSION_DIR+'/captions' MDLPTH=str(SESSION_DIR+"/"+Session_Name+'.safetensors') if os.path.exists(str(SESSION_DIR)) and not os.path.exists(MDLPTH): print('Loading session with no previous LoRa model') if MODEL_NAMExl=="": print('No model found, use the "Model Download" cell to download a model.') else: print('Session Loaded, proceed') elif not os.path.exists(str(SESSION_DIR)): call('mkdir -p '+INSTANCE_DIR, shell=True) print('Creating session...') if MODEL_NAMExl=="": print('No model found, use the "Model Download" cell to download a model.') else: print('Session created, proceed to uploading instance images') if MODEL_NAMExl=="": print('No model found, use the "Model Download" cell to download a model.') else: print('Session Loaded, proceed') return WORKSPACE, Session_Name, INSTANCE_NAME, OUTPUT_DIR, SESSION_DIR, INSTANCE_DIR, CAPTIONS_DIR, MDLPTH, MODEL_NAMExl def uplder(Remove_existing_instance_images, Crop_images, Crop_size, IMAGES_FOLDER_OPTIONAL, INSTANCE_DIR, CAPTIONS_DIR): if os.path.exists(INSTANCE_DIR+"/.ipynb_checkpoints"): call('rm -r '+INSTANCE_DIR+'/.ipynb_checkpoints', shell=True) uploader = widgets.FileUpload(description="Choose images",accept='image/*, .txt', multiple=True) Upload = widgets.Button( description='Upload', disabled=False, button_style='info', tooltip='Click to upload the chosen instance images', icon='' ) def up(Upload): with out: uploader.close() Upload.close() upld(Remove_existing_instance_images, Crop_images, Crop_size, IMAGES_FOLDER_OPTIONAL, INSTANCE_DIR, CAPTIONS_DIR, uploader) done() out=widgets.Output() if IMAGES_FOLDER_OPTIONAL=="": Upload.on_click(up) display(uploader, Upload, out) else: upld(Remove_existing_instance_images, Crop_images, Crop_size, IMAGES_FOLDER_OPTIONAL, INSTANCE_DIR, CAPTIONS_DIR, uploader) done() def upld(Remove_existing_instance_images, Crop_images, Crop_size, IMAGES_FOLDER_OPTIONAL, INSTANCE_DIR, CAPTIONS_DIR, uploader): from tqdm import tqdm if Remove_existing_instance_images: if os.path.exists(str(INSTANCE_DIR)): call("rm -r " +INSTANCE_DIR, shell=True) if os.path.exists(str(CAPTIONS_DIR)): call("rm -r " +CAPTIONS_DIR, shell=True) if not os.path.exists(str(INSTANCE_DIR)): call("mkdir -p " +INSTANCE_DIR, shell=True) if not os.path.exists(str(CAPTIONS_DIR)): call("mkdir -p " +CAPTIONS_DIR, shell=True) if IMAGES_FOLDER_OPTIONAL !="": if os.path.exists(IMAGES_FOLDER_OPTIONAL+"/.ipynb_checkpoints"): call('rm -r '+IMAGES_FOLDER_OPTIONAL+'/.ipynb_checkpoints', shell=True) if any(file.endswith('.{}'.format('txt')) for file in os.listdir(IMAGES_FOLDER_OPTIONAL)): call('mv '+IMAGES_FOLDER_OPTIONAL+'/*.txt '+CAPTIONS_DIR, shell=True) if Crop_images: os.chdir(str(IMAGES_FOLDER_OPTIONAL)) call('find . -name "* *" -type f | rename ' "'s/ /-/g'", shell=True) os.chdir('/workspace') for filename in tqdm(os.listdir(IMAGES_FOLDER_OPTIONAL), bar_format=' |{bar:15}| {n_fmt}/{total_fmt} Uploaded'): extension = filename.split(".")[-1] identifier=filename.split(".")[0] new_path_with_file = os.path.join(INSTANCE_DIR, filename) file = Image.open(IMAGES_FOLDER_OPTIONAL+"/"+filename) file=file.convert("RGB") file=ImageOps.exif_transpose(file) width, height = file.size if file.size !=(Crop_size, Crop_size): image=crop_image(file, Crop_size) if extension.upper()=="JPG" or extension.upper()=="jpg": image[0].save(new_path_with_file, format="JPEG", quality = 100) else: image[0].save(new_path_with_file, format=extension.upper()) else: call("cp \'"+IMAGES_FOLDER_OPTIONAL+"/"+filename+"\' "+INSTANCE_DIR, shell=True) else: for filename in tqdm(os.listdir(IMAGES_FOLDER_OPTIONAL), bar_format=' |{bar:15}| {n_fmt}/{total_fmt} Uploaded'): call("cp -r " +IMAGES_FOLDER_OPTIONAL+"/. " +INSTANCE_DIR, shell=True) elif IMAGES_FOLDER_OPTIONAL =="": up="" for file in uploader.value: filename = file['name'] if filename.split(".")[-1]=="txt": with open(CAPTIONS_DIR+'/'+filename, 'w') as f: f.write(bytes(file['content']).decode()) up=[file for file in uploader.value if not file['name'].endswith('.txt')] if Crop_images: for file in tqdm(up, bar_format=' |{bar:15}| {n_fmt}/{total_fmt} Uploaded'): filename = file['name'] img = Image.open(io.BytesIO(file['content'])) img=img.convert("RGB") img=ImageOps.exif_transpose(img) extension = filename.split(".")[-1] identifier=filename.split(".")[0] if extension.upper()=="JPG" or extension.upper()=="jpg": img.save(INSTANCE_DIR+"/"+filename, format="JPEG", quality = 100) else: img.save(INSTANCE_DIR+"/"+filename, format=extension.upper()) new_path_with_file = os.path.join(INSTANCE_DIR, filename) file = Image.open(new_path_with_file) width, height = file.size if file.size !=(Crop_size, Crop_size): image=crop_image(file, Crop_size) if extension.upper()=="JPG" or extension.upper()=="jpg": image[0].save(new_path_with_file, format="JPEG", quality = 100) else: image[0].save(new_path_with_file, format=extension.upper()) else: for file in tqdm(uploader.value, bar_format=' |{bar:15}| {n_fmt}/{total_fmt} Uploaded'): filename = file['name'] img = Image.open(io.BytesIO(file['content'])) img=img.convert("RGB") extension = filename.split(".")[-1] identifier=filename.split(".")[0] if extension.upper()=="JPG" or extension.upper()=="jpg": img.save(INSTANCE_DIR+"/"+filename, format="JPEG", quality = 100) else: img.save(INSTANCE_DIR+"/"+filename, format=extension.upper()) os.chdir(INSTANCE_DIR) call('find . -name "* *" -type f | rename ' "'s/ /-/g'", shell=True) os.chdir(CAPTIONS_DIR) call('find . -name "* *" -type f | rename ' "'s/ /-/g'", shell=True) os.chdir('/workspace') def caption(CAPTIONS_DIR, INSTANCE_DIR): paths="" out="" widgets_l="" clear_output() def Caption(path): if path!="Select an instance image to caption": name = os.path.splitext(os.path.basename(path))[0] ext=os.path.splitext(os.path.basename(path))[-1][1:] if ext=="jpg" or "JPG": ext="JPEG" if os.path.exists(CAPTIONS_DIR+"/"+name + '.txt'): with open(CAPTIONS_DIR+"/"+name + '.txt', 'r') as f: text = f.read() else: with open(CAPTIONS_DIR+"/"+name + '.txt', 'w') as f: f.write("") with open(CAPTIONS_DIR+"/"+name + '.txt', 'r') as f: text = f.read() img=Image.open(os.path.join(INSTANCE_DIR,path)) img=img.convert("RGB") img=img.resize((420, 420)) image_bytes = BytesIO() img.save(image_bytes, format=ext, qualiy=10) image_bytes.seek(0) image_data = image_bytes.read() img= image_data image = widgets.Image( value=img, width=420, height=420 ) text_area = widgets.Textarea(value=text, description='', disabled=False, layout={'width': '300px', 'height': '120px'}) def update_text(text): with open(CAPTIONS_DIR+"/"+name + '.txt', 'w') as f: f.write(text) button = widgets.Button(description='Save', button_style='success') button.on_click(lambda b: update_text(text_area.value)) return widgets.VBox([widgets.HBox([image, text_area, button])]) paths = os.listdir(INSTANCE_DIR) widgets_l = widgets.Select(options=["Select an instance image to caption"]+paths, rows=25) out = widgets.Output() def click(change): with out: out.clear_output() display(Caption(change.new)) widgets_l.observe(click, names='value') display(widgets.HBox([widgets_l, out])) def dbtrainxl(Unet_Training_Epochs, Text_Encoder_Training_Epochs, Unet_Learning_Rate, Text_Encoder_Learning_Rate, dim, Offset_Noise, Resolution, MODEL_NAME, SESSION_DIR, INSTANCE_DIR, CAPTIONS_DIR, External_Captions, INSTANCE_NAME, Session_Name, OUTPUT_DIR, ofstnselvl, Save_VRAM): if os.path.exists(INSTANCE_DIR+"/.ipynb_checkpoints"): call('rm -r '+INSTANCE_DIR+'/.ipynb_checkpoints', shell=True) if os.path.exists(CAPTIONS_DIR+"/.ipynb_checkpoints"): call('rm -r '+CAPTIONS_DIR+'/.ipynb_checkpoints', shell=True) Seed=random.randint(1, 999999) ofstnse="" if Offset_Noise: ofstnse="--offset_noise" GC='' if Save_VRAM: GC='--gradient_checkpointing' extrnlcptn="" if External_Captions: extrnlcptn="--external_captions" precision="fp16" def train_only_text(SESSION_DIR, MODEL_NAME, INSTANCE_DIR, OUTPUT_DIR, Seed, Resolution, ofstnse, extrnlcptn, precision, Training_Epochs): print('Training the Text Encoder...') call('accelerate launch /workspace/diffusers/examples/dreambooth/train_dreambooth_rnpd_sdxl_TI.py \ '+ofstnse+' \ '+extrnlcptn+' \ --dim='+str(dim)+' \ --ofstnselvl='+str(ofstnselvl)+' \ --image_captions_filename \ --Session_dir='+SESSION_DIR+' \ --pretrained_model_name_or_path='+MODEL_NAME+' \ --instance_data_dir='+INSTANCE_DIR+' \ --output_dir='+OUTPUT_DIR+' \ --captions_dir='+CAPTIONS_DIR+' \ --seed='+str(Seed)+' \ --resolution='+str(Resolution)+' \ --mixed_precision='+str(precision)+' \ --train_batch_size=1 \ --gradient_accumulation_steps=1 '+GC+ ' \ --use_8bit_adam \ --learning_rate='+str(Text_Encoder_Learning_Rate)+' \ --lr_scheduler="cosine" \ --lr_warmup_steps=0 \ --num_train_epochs='+str(Training_Epochs), shell=True) def train_only_unet(SESSION_DIR, MODEL_NAME, INSTANCE_DIR, OUTPUT_DIR, Seed, Resolution, ofstnse, extrnlcptn, precision, Training_Epochs): print('Training the UNet...') call('accelerate launch /workspace/diffusers/examples/dreambooth/train_dreambooth_rnpd_sdxl_lora.py \ '+ofstnse+' \ '+extrnlcptn+' \ --dim='+str(dim)+' \ --ofstnselvl='+str(ofstnselvl)+' \ --image_captions_filename \ --Session_dir='+SESSION_DIR+' \ --pretrained_model_name_or_path='+MODEL_NAME+' \ --instance_data_dir='+INSTANCE_DIR+' \ --output_dir='+OUTPUT_DIR+' \ --captions_dir='+CAPTIONS_DIR+' \ --seed='+str(Seed)+' \ --resolution='+str(Resolution)+' \ --mixed_precision='+str(precision)+' \ --train_batch_size=1 \ --gradient_accumulation_steps=1 '+GC+ ' \ --use_8bit_adam \ --learning_rate='+str(Unet_Learning_Rate)+' \ --lr_scheduler="cosine" \ --lr_warmup_steps=0 \ --num_train_epochs='+str(Training_Epochs), shell=True) if Unet_Training_Epochs!=0: if Text_Encoder_Training_Epochs!=0: train_only_text(SESSION_DIR, MODEL_NAME, INSTANCE_DIR, OUTPUT_DIR, Seed, Resolution, ofstnse, extrnlcptn, precision, Training_Epochs=Text_Encoder_Training_Epochs) clear_output() train_only_unet(SESSION_DIR, MODEL_NAME, INSTANCE_DIR, OUTPUT_DIR, Seed, Resolution, ofstnse, extrnlcptn, precision, Training_Epochs=Unet_Training_Epochs) else : print('Nothing to do') if os.path.exists(SESSION_DIR+'/'+Session_Name+'.safetensors'): clear_output() print("DONE, the LoRa model is in the session's folder") else: print("Something went wrong") def sdcmff(Huggingface_token_optional, MDLPTH, restored): from slugify import slugify from huggingface_hub import HfApi, CommitOperationAdd, create_repo os.chdir('/workspace') if restored: Huggingface_token_optional="" if Huggingface_token_optional!="": username = HfApi().whoami(Huggingface_token_optional)["name"] backup=f"https://huggingface.co/datasets/{username}/fast-stable-diffusion/resolve/main/sdcomfy_backup_rnpd.tar.zst" headers = {"Authorization": f"Bearer {Huggingface_token_optional}"} response = requests.head(backup, headers=headers) if response.status_code == 302: restored=True print('Restoring ComfyUI...') open('/workspace/sdcomfy_backup_rnpd.tar.zst', 'wb').write(requests.get(backup, headers=headers).content) call('tar --zstd -xf sdcomfy_backup_rnpd.tar.zst', shell=True) call('rm sdcomfy_backup_rnpd.tar.zst', shell=True) else: print('Backup not found, using a fresh/existing repo...') time.sleep(2) if not os.path.exists('ComfyUI'): call('git clone -q --depth 1 https://github.com/comfyanonymous/ComfyUI', shell=True) else: print('Installing/Updating the repo...') if not os.path.exists('ComfyUI'): call('git clone -q --depth 1 https://github.com/comfyanonymous/ComfyUI', shell=True) os.chdir('ComfyUI') call('git reset --hard', shell=True) print('') call('git pull', shell=True) if os.path.exists(MDLPTH): call('ln -s '+MDLPTH+' models/loras', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w')) clean_symlinks('models/loras') if not os.path.exists('models/checkpoints/sd_xl_base_1.0.safetensors'): call('ln -s /workspace/auto-models/* models/checkpoints', shell=True) podid=os.environ.get('RUNPOD_POD_ID') localurl=f"https://{podid}-3001.proxy.runpod.net" call("sed -i 's@print(\"To see the GUI go to: http://{}:{}\".format(address, port))@print(\"\u2714 Connected\")\\n print(\""+localurl+"\")@' /workspace/ComfyUI/server.py", shell=True) os.chdir('/workspace') return restored def test(MDLPTH, User, Password, Huggingface_token_optional, restoreda): from slugify import slugify from huggingface_hub import HfApi, CommitOperationAdd, create_repo import gradio gradio.close_all() auth=f"--gradio-auth {User}:{Password}" if User =="" or Password=="": auth="" if restoreda: Huggingface_token_optional="" if Huggingface_token_optional!="": username = HfApi().whoami(Huggingface_token_optional)["name"] backup=f"https://huggingface.co/datasets/{username}/fast-stable-diffusion/resolve/main/sd_backup_rnpd.tar.zst" headers = {"Authorization": f"Bearer {Huggingface_token_optional}"} response = requests.head(backup, headers=headers) if response.status_code == 302: restoreda=True print('Restoring the SD folder...') open('/workspace/sd_backup_rnpd.tar.zst', 'wb').write(requests.get(backup, headers=headers).content) call('tar --zstd -xf sd_backup_rnpd.tar.zst', shell=True) call('rm sd_backup_rnpd.tar.zst', shell=True) else: print('Backup not found, using a fresh/existing repo...') time.sleep(2) if not os.path.exists('/workspace/sd/stablediffusiond'): #reset later call('wget -q -O sd_mrep.tar.zst https://huggingface.co/TheLastBen/dependencies/resolve/main/sd_mrep.tar.zst', shell=True) call('tar --zstd -xf sd_mrep.tar.zst', shell=True) call('rm sd_mrep.tar.zst', shell=True) os.chdir('/workspace/sd') if not os.path.exists('stable-diffusion-webui'): call('git clone -q --depth 1 --branch master https://github.com/AUTOMATIC1111/stable-diffusion-webui', shell=True) else: print('Installing/Updating the repo...') os.chdir('/workspace') if not os.path.exists('/workspace/sd/stablediffusiond'): #reset later call('wget -q -O sd_mrep.tar.zst https://huggingface.co/TheLastBen/dependencies/resolve/main/sd_mrep.tar.zst', shell=True) call('tar --zstd -xf sd_mrep.tar.zst', shell=True) call('rm sd_mrep.tar.zst', shell=True) os.chdir('/workspace/sd') if not os.path.exists('stable-diffusion-webui'): call('git clone -q --depth 1 --branch master https://github.com/AUTOMATIC1111/stable-diffusion-webui', shell=True) os.chdir('/workspace/sd/stable-diffusion-webui/') call('git reset --hard', shell=True) print('') call('git pull', shell=True) if os.path.exists(MDLPTH): call('mkdir models/Lora', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w')) call('ln -s '+MDLPTH+' models/Lora', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w')) if not os.path.exists('models/Stable-diffusion/sd_xl_base_1.0.safetensors'): call('ln -s /workspace/auto-models/* models/Stable-diffusion', shell=True) clean_symlinks('models/Lora') os.chdir('/workspace') call('wget -q -O /usr/local/lib/python3.10/dist-packages/gradio/blocks.py https://raw.githubusercontent.com/TheLastBen/fast-stable-diffusion/main/AUTOMATIC1111_files/blocks.py', shell=True) os.chdir('/workspace/sd/stable-diffusion-webui/modules') call("sed -i 's@possible_sd_paths =.*@possible_sd_paths = [\"/workspace/sd/stablediffusion\"]@' /workspace/sd/stable-diffusion-webui/modules/paths.py", shell=True) call("sed -i 's@\.\.\/@src/@g' /workspace/sd/stable-diffusion-webui/modules/paths.py", shell=True) call("sed -i 's@src\/generative-models@generative-models@g' /workspace/sd/stable-diffusion-webui/modules/paths.py", shell=True) call("sed -i 's@\[\"sd_model_checkpoint\"\]@\[\"sd_model_checkpoint\", \"sd_vae\", \"CLIP_stop_at_last_layers\", \"inpainting_mask_weight\", \"initial_noise_multiplier\"\]@g' /workspace/sd/stable-diffusion-webui/modules/shared.py", shell=True) call("sed -i 's@print(\"No module.*@@' /workspace/sd/stablediffusion/ldm/modules/diffusionmodules/model.py", shell=True) os.chdir('/workspace/sd/stable-diffusion-webui') clear_output() podid=os.environ.get('RUNPOD_POD_ID') localurl=f"{podid}-3001.proxy.runpod.net" for line in fileinput.input('/usr/local/lib/python3.10/dist-packages/gradio/blocks.py', inplace=True): if line.strip().startswith('self.server_name ='): line = f' self.server_name = "{localurl}"\n' if line.strip().startswith('self.protocol = "https"'): line = ' self.protocol = "https"\n' if line.strip().startswith('if self.local_url.startswith("https") or self.is_colab'): line = '' if line.strip().startswith('else "http"'): line = '' sys.stdout.write(line) configf="--disable-console-progressbars --upcast-sampling --no-half-vae --disable-safe-unpickle --api --opt-sdp-attention --enable-insecure-extension-access --no-download-sd-model --skip-version-check --listen --port 3000 --ckpt /workspace/sd/stable-diffusion-webui/models/Stable-diffusion/sd_xl_base_1.0.safetensors "+auth return configf, restoreda def clean(): Sessions=os.listdir("/workspace/Fast-Dreambooth/Sessions") s = widgets.Select( options=Sessions, rows=5, description='', disabled=False ) out=widgets.Output() d = widgets.Button( description='Remove', disabled=False, button_style='warning', tooltip='Removet the selected session', icon='warning' ) def rem(d): with out: if s.value is not None: clear_output() print("THE SESSION "+s.value+" HAS BEEN REMOVED FROM THE STORAGE") call('rm -r /workspace/Fast-Dreambooth/Sessions/'+s.value, shell=True) if os.path.exists('/workspace/models/'+s.value): call('rm -r /workspace/models/'+s.value, shell=True) s.options=os.listdir("/workspace/Fast-Dreambooth/Sessions") else: d.close() s.close() clear_output() print("NOTHING TO REMOVE") d.on_click(rem) if s.value is not None: display(s,d,out) else: print("NOTHING TO REMOVE") def crop_image(im, size): import cv2 GREEN = "#0F0" BLUE = "#00F" RED = "#F00" def focal_point(im, settings): corner_points = image_corner_points(im, settings) if settings.corner_points_weight > 0 else [] entropy_points = image_entropy_points(im, settings) if settings.entropy_points_weight > 0 else [] face_points = image_face_points(im, settings) if settings.face_points_weight > 0 else [] pois = [] weight_pref_total = 0 if len(corner_points) > 0: weight_pref_total += settings.corner_points_weight if len(entropy_points) > 0: weight_pref_total += settings.entropy_points_weight if len(face_points) > 0: weight_pref_total += settings.face_points_weight corner_centroid = None if len(corner_points) > 0: corner_centroid = centroid(corner_points) corner_centroid.weight = settings.corner_points_weight / weight_pref_total pois.append(corner_centroid) entropy_centroid = None if len(entropy_points) > 0: entropy_centroid = centroid(entropy_points) entropy_centroid.weight = settings.entropy_points_weight / weight_pref_total pois.append(entropy_centroid) face_centroid = None if len(face_points) > 0: face_centroid = centroid(face_points) face_centroid.weight = settings.face_points_weight / weight_pref_total pois.append(face_centroid) average_point = poi_average(pois, settings) return average_point def image_face_points(im, settings): np_im = np.array(im) gray = cv2.cvtColor(np_im, cv2.COLOR_BGR2GRAY) tries = [ [ f'{cv2.data.haarcascades}haarcascade_eye.xml', 0.01 ], [ f'{cv2.data.haarcascades}haarcascade_frontalface_default.xml', 0.05 ], [ f'{cv2.data.haarcascades}haarcascade_profileface.xml', 0.05 ], [ f'{cv2.data.haarcascades}haarcascade_frontalface_alt.xml', 0.05 ], [ f'{cv2.data.haarcascades}haarcascade_frontalface_alt2.xml', 0.05 ], [ f'{cv2.data.haarcascades}haarcascade_frontalface_alt_tree.xml', 0.05 ], [ f'{cv2.data.haarcascades}haarcascade_eye_tree_eyeglasses.xml', 0.05 ], [ f'{cv2.data.haarcascades}haarcascade_upperbody.xml', 0.05 ] ] for t in tries: classifier = cv2.CascadeClassifier(t[0]) minsize = int(min(im.width, im.height) * t[1]) # at least N percent of the smallest side try: faces = classifier.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=7, minSize=(minsize, minsize), flags=cv2.CASCADE_SCALE_IMAGE) except: continue if len(faces) > 0: rects = [[f[0], f[1], f[0] + f[2], f[1] + f[3]] for f in faces] return [PointOfInterest((r[0] +r[2]) // 2, (r[1] + r[3]) // 2, size=abs(r[0]-r[2]), weight=1/len(rects)) for r in rects] return [] def image_corner_points(im, settings): grayscale = im.convert("L") # naive attempt at preventing focal points from collecting at watermarks near the bottom gd = ImageDraw.Draw(grayscale) gd.rectangle([0, im.height*.9, im.width, im.height], fill="#999") np_im = np.array(grayscale) points = cv2.goodFeaturesToTrack( np_im, maxCorners=100, qualityLevel=0.04, minDistance=min(grayscale.width, grayscale.height)*0.06, useHarrisDetector=False, ) if points is None: return [] focal_points = [] for point in points: x, y = point.ravel() focal_points.append(PointOfInterest(x, y, size=4, weight=1/len(points))) return focal_points def image_entropy_points(im, settings): landscape = im.height < im.width portrait = im.height > im.width if landscape: move_idx = [0, 2] move_max = im.size[0] elif portrait: move_idx = [1, 3] move_max = im.size[1] else: return [] e_max = 0 crop_current = [0, 0, settings.crop_width, settings.crop_height] crop_best = crop_current while crop_current[move_idx[1]] < move_max: crop = im.crop(tuple(crop_current)) e = image_entropy(crop) if (e > e_max): e_max = e crop_best = list(crop_current) crop_current[move_idx[0]] += 4 crop_current[move_idx[1]] += 4 x_mid = int(crop_best[0] + settings.crop_width/2) y_mid = int(crop_best[1] + settings.crop_height/2) return [PointOfInterest(x_mid, y_mid, size=25, weight=1.0)] def image_entropy(im): # greyscale image entropy # band = np.asarray(im.convert("L")) band = np.asarray(im.convert("1"), dtype=np.uint8) hist, _ = np.histogram(band, bins=range(0, 256)) hist = hist[hist > 0] return -np.log2(hist / hist.sum()).sum() def centroid(pois): x = [poi.x for poi in pois] y = [poi.y for poi in pois] return PointOfInterest(sum(x)/len(pois), sum(y)/len(pois)) def poi_average(pois, settings): weight = 0.0 x = 0.0 y = 0.0 for poi in pois: weight += poi.weight x += poi.x * poi.weight y += poi.y * poi.weight avg_x = round(weight and x / weight) avg_y = round(weight and y / weight) return PointOfInterest(avg_x, avg_y) def is_landscape(w, h): return w > h def is_portrait(w, h): return h > w def is_square(w, h): return w == h class PointOfInterest: def __init__(self, x, y, weight=1.0, size=10): self.x = x self.y = y self.weight = weight self.size = size def bounding(self, size): return [ self.x - size//2, self.y - size//2, self.x + size//2, self.y + size//2 ] class Settings: def __init__(self, crop_width=512, crop_height=512, corner_points_weight=0.5, entropy_points_weight=0.5, face_points_weight=0.5): self.crop_width = crop_width self.crop_height = crop_height self.corner_points_weight = corner_points_weight self.entropy_points_weight = entropy_points_weight self.face_points_weight = face_points_weight settings = Settings( crop_width = size, crop_height = size, face_points_weight = 0.9, entropy_points_weight = 0.15, corner_points_weight = 0.5, ) scale_by = 1 if is_landscape(im.width, im.height): scale_by = settings.crop_height / im.height elif is_portrait(im.width, im.height): scale_by = settings.crop_width / im.width elif is_square(im.width, im.height): if is_square(settings.crop_width, settings.crop_height): scale_by = settings.crop_width / im.width elif is_landscape(settings.crop_width, settings.crop_height): scale_by = settings.crop_width / im.width elif is_portrait(settings.crop_width, settings.crop_height): scale_by = settings.crop_height / im.height im = im.resize((int(im.width * scale_by), int(im.height * scale_by))) im_debug = im.copy() focus = focal_point(im_debug, settings) # take the focal point and turn it into crop coordinates that try to center over the focal # point but then get adjusted back into the frame y_half = int(settings.crop_height / 2) x_half = int(settings.crop_width / 2) x1 = focus.x - x_half if x1 < 0: x1 = 0 elif x1 + settings.crop_width > im.width: x1 = im.width - settings.crop_width y1 = focus.y - y_half if y1 < 0: y1 = 0 elif y1 + settings.crop_height > im.height: y1 = im.height - settings.crop_height x2 = x1 + settings.crop_width y2 = y1 + settings.crop_height crop = [x1, y1, x2, y2] results = [] results.append(im.crop(tuple(crop))) return results def clean_symlinks(path): for item in os.listdir(path): lnk = os.path.join(path, item) if os.path.islink(lnk) and not os.path.exists(os.readlink(lnk)): os.remove(lnk)