Mirrors
/
EveryDream-trainer
mirror of https://github.com/victorchall/EveryDream-trainer.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Porting changes over from my branch.

This commit is contained in:
David Bielejeski 2022-09-26 08:01:18 -05:00
parent 3d97f68047
commit fd298537c2
2 changed files with 537 additions and 414 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

924
dreambooth_runpod_joepenna.ipynb
View File

@ -1,412 +1,518 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "aa2c1ada",
"metadata": {
"id": "aa2c1ada"
},
"source": [
"# Dreambooth\n",
"### Notebook implementation by Joe Penna (@MysteryGuitarM on Twitter) - Improvements by David Bielejeski\n",
"https://github.com/JoePenna/Dreambooth-Stable-Diffusion\n",
"\n",
"### If on runpod / vast.ai / etc, spin up an A6000 or A100 pod using a Stable Diffusion template with Jupyter pre-installed."
]
},
{
"cell_type": "markdown",
"id": "7b971cc0",
"metadata": {
"id": "7b971cc0"
},
"source": [
"## Build Environment"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2AsGA1xpNQnb",
"metadata": {
"id": "2AsGA1xpNQnb"
},
"outputs": [],
"source": [
"# If running on Vast.AI, copy the code in this cell into a new notebook. Run it, then launch the `dreambooth_runpod_joepenna.ipynb` notebook from the jupyter interface.\n",
"!git clone https://github.com/JoePenna/Dreambooth-Stable-Diffusion"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9e1bc458-091b-42f4-a125-c3f0df20f29d",
"metadata": {
"id": "9e1bc458-091b-42f4-a125-c3f0df20f29d",
"scrolled": true
},
"outputs": [],
"source": [
"%cd Dreambooth-Stable-Diffusion\n",
"\n",
"#HOUSEKEEPING\n",
"!ln -s /usr/bin/python3.6 /usr/bin/python\n",
"!ln -s /usr/bin/pip3 /usr/bin/pip\n",
"!mkdir training_samples\n",
"\n",
"#BUILD ENV\n",
"%pip install omegaconf\n",
"%pip install einops\n",
"%pip install pytorch-lightning==1.6.5\n",
"%pip install protobuf==3.20.1\n",
"%pip install test-tube\n",
"%pip install transformers\n",
"%pip install kornia\n",
"%pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers\n",
"%pip install -e git+https://github.com/openai/CLIP.git@main#egg=clip\n",
"%pip install setuptools==59.5.0\n",
"%pip install pillow==9.0.1\n",
"%pip install torchmetrics==0.6.0\n",
"%pip install torchvision\n",
"%pip install -e .\n",
"%pip install gdown\n",
"%pip install pydrive"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "dae11c10",
"metadata": {
"id": "dae11c10"
},
"outputs": [],
"source": [
"## DOWNLOAD SD\n",
"%pip install gdown\n",
"\n",
"#Link to a Google Drive ckpt of Stable Diffusion\n",
"!gdown https://drive.google.com/uc?id={{FILE ID}}"
]
},
{
"cell_type": "markdown",
"id": "17d1d11a",
"metadata": {
"id": "17d1d11a"
},
"source": [
"# Regularization Images"
]
},
{
"cell_type": "markdown",
"id": "ed07a5df",
"metadata": {
"id": "ed07a5df"
},
"source": [
"Training teaches your new model both your token **but** re-trains your class simultaneously.\n",
"\n",
"From cursory testing, it does not seem like reg images affect the model too much. However, they do affect your class greatly, which will in turn affect your generations.\n",
"\n",
"You can either generate your images here, or use the repos below to quickly download 1500 images."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "67f9ff0c-b529-4c7c-8e26-8388d70a5d91",
"metadata": {
"id": "67f9ff0c-b529-4c7c-8e26-8388d70a5d91"
},
"outputs": [],
"source": [
"# GENERATE 400 images\n",
"!python scripts/stable_txt2img.py --seed 10 --ddim_eta 0.0 --n_samples 1 --n_iter 400 --scale 10.0 --ddim_steps 50 --ckpt model.ckpt \\\n",
"--prompt \"person\"\n",
"\n",
"# If you don't want to train against \"person\", change it to whatever you want above, and on some of the cells below:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3d1c7e1c",
"metadata": {
"id": "3d1c7e1c"
},
"outputs": [],
"source": [
"# zip up the files for downloading and reuse.\n",
"!apt-get install -y zip\n",
"!zip -r all_images.zip outputs/"
]
},
{
"cell_type": "markdown",
"id": "mxPL2O0OLvBW",
"metadata": {
"id": "mxPL2O0OLvBW"
},
"source": [
"## Download pre-generated regularization images"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e7EydXCjOV1v",
"metadata": {
"id": "e7EydXCjOV1v"
},
"outputs": [],
"source": [
"# grab the files here: https://github.com/djbielejeski/Stable-Diffusion-Regularization-Images\n",
"!git clone https://github.com/JoePenna/Stable-Diffusion-Regularization-Images\n",
"\n",
"!mkdir outputs\n",
"!mkdir outputs/txt2img-samples\n",
"!mkdir outputs/txt2img-samples/samples"
]
},
{
"cell_type": "markdown",
"id": "aeca9198",
"metadata": {},
"source": [
"## RECOMMENDED train with \"person\":"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2402864a",
"metadata": {},
"outputs": [],
"source": [
"# 1500 images - ddim @ 50 steps, CFG 7.0\n",
"!mv -v Stable-Diffusion-Regularization-Images/person_ddim outputs/txt2img-samples outputs/txt2img-samples/samples"
]
},
{
"cell_type": "markdown",
"id": "e8eb85b5",
"metadata": {},
"source": [
"### If you'd rather try `man` or `woman`, comment the line above, and uncomment one of the ones below:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8891aa90",
"metadata": {},
"outputs": [],
"source": [
"# Use man_euler images - provided by Niko Pueringer (Corridor Digital) - euler @ 40 steps, CFG 7.5\n",
"# !mv -v Stable-Diffusion-Regularization-Images/man_euler outputs/txt2img-samples\n",
"\n",
"# Use man_unsplash images - pictures from various photographers\n",
"# !mv -v Stable-Diffusion-Regularization-Images/man_unsplash outputs/txt2img-samples\n",
"\n",
"# Use woman_dimm images - provided by David Bielejeski - ddim\n",
"# !mv -v Stable-Diffusion-Regularization-Images/woman_dimm outputs/txt2img-samples"
]
},
{
"cell_type": "markdown",
"id": "zshrC_JuMXmM",
"metadata": {
"id": "zshrC_JuMXmM"
},
"source": [
"# Upload your training images\n",
"Upload 10-20 images of someone to /workspace/Dreambooth-Stable-Diffusion/training_samples\n",
"\n",
"WARNING: Be sure to upload an *even* amount of images, otherwise the training inexplicably stops at 1500 steps.\n",
"\n",
"* 2-3 full body\n",
"* 3-5 upper body \n",
"* 5-12 close-up on face"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "bff9314e",
"metadata": {
"id": "bff9314e"
},
"outputs": [],
"source": [
"# or upload from your Google Drive\n",
"%pip install gdown\n",
"!gdown https://drive.google.com/uc?id={fileId}\n",
"!mv {name of file} training_samples/{name of file}"
]
},
{
"cell_type": "markdown",
"id": "ad4e50df",
"metadata": {
"id": "ad4e50df"
},
"source": [
"## Prep Training Variables"
]
},
{
"cell_type": "markdown",
"id": "a26964af",
"metadata": {
"id": "a26964af"
},
"source": [
"Navigate to:\n",
"/workspace/Dreambooth-Stable-Diffusion/ldm/data/personalized.py\n",
"\n",
"Change `demoura` in line 11 to whatever you want your token to be.\n",
"\n",
"(Note: The original repo uses `sks` as a token -- but that has a relatively high correlation with the SKS semi-automatic rifle, which affects low-step generations)\n",
"\n",
"e.g. I changed mine to\n",
"```python\n",
"training_templates_smallest = [\n",
" 'joepenna {}',\n",
"]\n",
"```"
]
},
{
"cell_type": "markdown",
"id": "37612c32",
"metadata": {
"id": "37612c32"
},
"source": [
"The last line of this file trains for `3000` steps.\n",
"/workspace/Dreambooth-Stable-Diffusion/configs/stable-diffusion/v1-finetune_unfrozen.yaml\n",
"\n",
"If training a person or subject, keep an eye on your project's `/images/train/samples_scaled_gs-00xxxx` generations.\n",
"\n",
"If training a style, keep an eye on your project's `/images/train/samples_gs-00xxxx` generations."
]
},
{
"cell_type": "markdown",
"id": "12f19de7",
"metadata": {
"id": "12f19de7"
},
"source": [
"## Start Training (you should also change the code below)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6fa5dd66-2ca0-4819-907e-802e25583ae6",
"metadata": {
"id": "6fa5dd66-2ca0-4819-907e-802e25583ae6",
"tags": []
},
"outputs": [],
"source": [
"# START THE TRAINING\n",
"!python \"main.py\" \\\n",
" --base configs/stable-diffusion/v1-finetune_unfrozen.yaml \\\n",
" -t \\\n",
" --actual_resume \"model.ckpt\" \\\n",
" --reg_data_root \"/workspace/Dreambooth-Stable-Diffusion/outputs/txt2img-samples/samples\" \\\n",
" -n \"project_name\" \\\n",
" --gpus 0, \\\n",
" --data_root \"/workspace/Dreambooth-Stable-Diffusion/training_samples\" \\\n",
" --class_word \"person\" # << match this word to the class word from regularization images above"
]
},
{
"cell_type": "markdown",
"id": "dc49d0bd",
"metadata": {},
"source": [
"## Pruning (12GB to 2GB)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# This will prune around 10GB from the ckpt file, then move it the \"trained_models\" folder\n",
"!python \"scripts/prune-ckpt.py\" --ckpt logs/{training_name}/checkpoints/last.ckpt\n",
"!mkdir trained_models\n",
"!mv logs/{training_name}/checkpoints/last-pruned.ckpt workspace/Dreambooth-Stable-Diffusion/trained_models/{training_name}.ckpt"
]
},
{
"cell_type": "markdown",
"id": "d28d0139",
"metadata": {},
"source": [
"## Generate Images With Your Trained Model!"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "80ddb03b",
"metadata": {},
"outputs": [],
"source": [
"!python scripts/stable_txt2img.py \\\n",
"--ddim_eta 0.0 --n_samples 1 --n_iter 4 \\\n",
"--scale 7.0 --ddim_steps 50 --ckpt \"/workspace/Dreambooth-Stable-Diffusion/logs/{training_name}/last.ckpt\" \\\n",
"--prompt \"demoura person as a masterpiece portrait painting by John Singer Sargent in the style of Rembrandt\""
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "cfm1iW_9OAw1",
"metadata": {
"id": "cfm1iW_9OAw1"
},
"outputs": [],
"source": [
"#Download the .ckpt file and use in your favorite Stable Diffusion repo!"
]
}
],
"metadata": {
"colab": {
"collapsed_sections": [],
"provenance": []
},
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.6"
},
"vscode": {
"interpreter": {
"hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
}
}
"cells": [
{
"cell_type": "markdown",
"id": "aa2c1ada",
"metadata": {
"id": "aa2c1ada"
},
"source": [
"# Dreambooth\n",
"### Notebook implementation by Joe Penna (@MysteryGuitarM on Twitter) - Improvements by David Bielejeski\n",
"https://github.com/JoePenna/Dreambooth-Stable-Diffusion\n",
"\n",
"### If on runpod / vast.ai / etc, spin up an A6000 or A100 pod using a Stable Diffusion template with Jupyter pre-installed."
]
},
"nbformat": 4,
"nbformat_minor": 5
{
"cell_type": "markdown",
"id": "7b971cc0",
"metadata": {
"id": "7b971cc0"
},
"source": [
"## Build Environment"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2AsGA1xpNQnb",
"metadata": {
"id": "2AsGA1xpNQnb"
},
"outputs": [],
"source": [
"# If running on Vast.AI, copy the code in this cell into a new notebook. Run it, then launch the `dreambooth_runpod_joepenna.ipynb` notebook from the jupyter interface.\n",
"!git clone https://github.com/JoePenna/Dreambooth-Stable-Diffusion"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9e1bc458-091b-42f4-a125-c3f0df20f29d",
"metadata": {
"id": "9e1bc458-091b-42f4-a125-c3f0df20f29d",
"scrolled": true
},
"outputs": [],
"source": [
"#BUILD ENV\n",
"!pip install omegaconf\n",
"!pip install einops\n",
"!pip install pytorch-lightning==1.6.5\n",
"!pip install test-tube\n",
"!pip install transformers\n",
"!pip install kornia\n",
"!pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers\n",
"!pip install -e git+https://github.com/openai/CLIP.git@main#egg=clip\n",
"!pip install setuptools==59.5.0\n",
"!pip install pillow==9.0.1\n",
"!pip install torchmetrics==0.6.0\n",
"!pip install -e .\n",
"!pip install protobuf==3.20.1\n",
"!pip install gdown\n",
"!pip install pydrive\n",
"!pip install -qq diffusers[\"training\"]==0.3.0 transformers ftfy\n",
"!pip install -qq \"ipywidgets>=7,<8\"\n",
"!pip install huggingface_hub\n",
"!pip install ipywidgets"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "dae11c10",
"metadata": {
"id": "dae11c10"
},
"outputs": [],
"source": [
"## Login to stable diffusion\n",
"from huggingface_hub import notebook_login\n",
"\n",
"notebook_login()"
]
},
{
"cell_type": "code",
"execution_count": null,
"outputs": [],
"source": [
"## Download the 1.4 sd model\n",
"from huggingface_hub import hf_hub_download\n",
"downloaded_model_path = hf_hub_download(\n",
" repo_id=\"CompVis/stable-diffusion-v-1-4-original\",\n",
" filename=\"sd-v1-4.ckpt\",\n",
" use_auth_token=True\n",
")"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "code",
"execution_count": null,
"outputs": [],
"source": [
"## Move the sd-v1-4.ckpt to the root of this directory as \"model.ckpt\"\n",
"actual_locations_of_model_blob = !readlink -f {downloaded_model_path}\n",
"!mv {actual_locations_of_model_blob[-1]} model.ckpt"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "markdown",
"id": "17d1d11a",
"metadata": {
"id": "17d1d11a"
},
"source": [
"# Regularization Images (Skip this section if you are uploading your own or using the provided images)"
]
},
{
"cell_type": "markdown",
"id": "ed07a5df",
"metadata": {
"id": "ed07a5df"
},
"source": [
"Training teaches your new model both your token **but** re-trains your class simultaneously.\n",
"\n",
"From cursory testing, it does not seem like reg images affect the model too much. However, they do affect your class greatly, which will in turn affect your generations.\n",
"\n",
"You can either generate your images here, or use the repos below to quickly download 1500 images."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "67f9ff0c-b529-4c7c-8e26-8388d70a5d91",
"metadata": {
"id": "67f9ff0c-b529-4c7c-8e26-8388d70a5d91"
},
"outputs": [],
"source": [
"# GENERATE 400 images\n",
"!python scripts/stable_txt2img.py \\\n",
" --seed 10 \\\n",
" --ddim_eta 0.0 \\\n",
" --n_samples 1 \\\n",
" --n_iter 400 \\\n",
" --scale 10.0 \\\n",
" --ddim_steps 50 \\\n",
" --ckpt model.ckpt \\\n",
" --prompt \"person\"\n",
"\n",
"# If you don't want to train against \"person\", change it to whatever you want above, and on some of the cells below:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3d1c7e1c",
"metadata": {
"id": "3d1c7e1c"
},
"outputs": [],
"source": [
"# zip up the files for downloading and reuse.\n",
"!apt-get install -y zip\n",
"!zip -r all_images.zip outputs/\n",
"\n",
"# Download this file locally so you can reuse during another training on this dataset"
]
},
{
"cell_type": "markdown",
"id": "mxPL2O0OLvBW",
"metadata": {
"id": "mxPL2O0OLvBW"
},
"source": [
"## Download pre-generated regularization images\n",
"\n",
"We've created the following image sets\n",
"\n",
"* man_euler - provided by Niko Pueringer (Corridor Digital) - euler @ 40 steps, CFG 7.5\n",
"* man_unsplash - pictures from various photographers\n",
"* person_ddim\n",
"* woman_ddim - provided by David Bielejeski - ddim @ 50 steps, CFG 10.0\n",
"\n",
"`person_ddim` is recommended"
]
},
{
"cell_type": "code",
"execution_count": null,
"outputs": [],
"source": [
"import ipywidgets as widgets\n",
"from IPython.display import display\n",
"\n",
"dataset='person_ddim'\n",
"w = widgets.Dropdown(\n",
" options=['man_euler', 'man_unsplash', 'person_ddim', 'woman_ddim'],\n",
" value='person_ddim',\n",
" description='Dataset:',\n",
")\n",
"\n",
"def on_change(change):\n",
" if change['type'] == 'change' and change['name'] == 'value':\n",
" dataset=change['new']\n",
" print(dataset + \" selected\")\n",
"\n",
"w.observe(on_change)\n",
"\n",
"display(w)"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "code",
"execution_count": null,
"id": "e7EydXCjOV1v",
"metadata": {
"id": "e7EydXCjOV1v"
},
"outputs": [],
"source": [
"# Grab the existing regularization images\n",
"# Choose the dataset that best represents what you are trying to do and matches what you used for your token\n",
"!git clone https://github.com/djbielejeski/Stable-Diffusion-Regularization-Images-{dataset}.git\n",
"\n",
"!mkdir -p outputs/txt2img-samples/samples/{dataset}\n",
"!mv -v Stable-Diffusion-Regularization-Images-{dataset}/{dataset}/*.* outputs/txt2img-samples/samples/{dataset}"
]
},
{
"cell_type": "markdown",
"id": "zshrC_JuMXmM",
"metadata": {
"id": "zshrC_JuMXmM"
},
"source": [
"# Upload your training images\n",
"Upload 10-20 images of someone to\n",
"\n",
"```\n",
"/workspace/Dreambooth-Stable-Diffusion/training_samples\n",
"```\n",
"\n",
"WARNING: Be sure to upload an *even* amount of images, otherwise the training inexplicably stops at 1500 steps.\n",
"\n",
"* 2-3 full body\n",
"* 3-5 upper body \n",
"* 5-12 close-up on face"
]
},
{
"cell_type": "code",
"execution_count": null,
"outputs": [],
"source": [
"#@markdown Add here the URLs to the images of the concept you are adding\n",
"urls = [\n",
" \"https://i.imgur.com/test1.png\",\n",
" \"https://i.imgur.com/test2.png\",\n",
" \"https://i.imgur.com/test3.png\",\n",
" \"https://i.imgur.com/test4.png\",\n",
" \"https://i.imgur.com/test5.png\",\n",
" ## You can add additional images here\n",
"]"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "code",
"execution_count": null,
"outputs": [],
"source": [
"#@title Download and check the images you have just added\n",
"import os\n",
"import requests\n",
"from io import BytesIO\n",
"from PIL import Image\n",
"\n",
"\n",
"def image_grid(imgs, rows, cols):\n",
" assert len(imgs) == rows*cols\n",
"\n",
" w, h = imgs[0].size\n",
" grid = Image.new('RGB', size=(cols*w, rows*h))\n",
" grid_w, grid_h = grid.size\n",
"\n",
" for i, img in enumerate(imgs):\n",
" grid.paste(img, box=(i%cols*w, i//cols*h))\n",
" return grid\n",
"\n",
"def download_image(url):\n",
" try:\n",
" response = requests.get(url)\n",
" except:\n",
" return None\n",
" return Image.open(BytesIO(response.content)).convert(\"RGB\")\n",
"\n",
"images = list(filter(None,[download_image(url) for url in urls]))\n",
"save_path = \"./training_samples\"\n",
"if not os.path.exists(save_path):\n",
" os.mkdir(save_path)\n",
"[image.save(f\"{save_path}/{i}.png\", format=\"png\") for i, image in enumerate(images)]\n",
"image_grid(images, 1, len(images))"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "markdown",
"id": "ad4e50df",
"metadata": {
"id": "ad4e50df"
},
"source": [
"## Training\n",
"\n",
"If training a person or subject, keep an eye on your project's `logs/{folder}/images/train/samples_scaled_gs-00xxxx` generations.\n",
"\n",
"If training a style, keep an eye on your project's `logs/{folder}/images/train/samples_gs-00xxxx` generations."
]
},
{
"cell_type": "markdown",
"source": [
"## Edit the personalized.py file\n",
"Execute this cell `%load ldm/data/personalized.py`\n",
"\n",
"Change `joepenna` to whatever you want it to be (but keep the {})\n",
"\n",
"```\n",
"training_templates_smallest = [\n",
" 'joepenna {}',\n",
"]\n",
"```\n",
"\n",
"Then paste this at the very top of the cell\n",
"\n",
"```\n",
"%%writefile ldm/data/personalized.py\n",
"```\n",
"\n",
"Then run the cell again. This will save your changes.\n"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "code",
"execution_count": null,
"outputs": [],
"source": [
"%load ldm/data/personalized.py"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "code",
"execution_count": null,
"id": "6fa5dd66-2ca0-4819-907e-802e25583ae6",
"metadata": {
"id": "6fa5dd66-2ca0-4819-907e-802e25583ae6",
"tags": []
},
"outputs": [],
"source": [
"# START THE TRAINING\n",
"project_name = \"project_name\"\n",
"batch_size = 3000\n",
"class_word = \"person\" # << match this word to the class word from regularization images above\n",
"\n",
"!rm -rf training_samples/.ipynb_checkpoints\n",
"!python \"main.py\" \\\n",
" --base configs/stable-diffusion/v1-finetune_unfrozen.yaml \\\n",
" -t \\\n",
" --actual_resume \"model.ckpt\" \\\n",
" --reg_data_root \"/workspace/Dreambooth-Stable-Diffusion/outputs/txt2img-samples/samples/\" + {dataset} \\\n",
" -n {project_name} \\\n",
" --gpus 0, \\\n",
" --data_root \"/workspace/Dreambooth-Stable-Diffusion/training_samples\" \\\n",
" --batch_size {batch_size} \\\n",
" --class_word {class_word}"
]
},
{
"cell_type": "markdown",
"id": "dc49d0bd",
"metadata": {},
"source": [
"## Pruning (12GB to 2GB)\n",
"We are working on having this happen automatically"
]
},
{
"cell_type": "code",
"execution_count": null,
"outputs": [],
"source": [
"directory_paths = !ls -d logs/*"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "code",
"execution_count": null,
"outputs": [],
"source": [
"# This version should automatically prune around 10GB from the ckpt file\n",
"last_checkpoint_file = directory_paths[-1] + \"/checkpoints/last.ckpt\"\n",
"!python \"prune-ckpt.py\" --ckpt {last_checkpoint_file}"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "code",
"execution_count": null,
"outputs": [],
"source": [
"last_checkpoint_file_pruned = directory_paths[-1] + \"/checkpoints/last-pruned.ckpt\"\n",
"training_samples = !ls training_samples\n",
"date_string = !date +\"%Y-%m-%dT%H-%M-%S\"\n",
"file_name = date_string[-1] + \"_\" + {project_name} + \"_\" + len(training_samples) + \"_training_images_\" + {batch_size} + \"_batch_size_\" + {class_word} + \"_class_word.ckpt\"\n",
"!mkdir trained_models\n",
"!mv {last_checkpoint_file_pruned} trained_models/{file_name}"
],
"metadata": {
"collapsed": false
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Download your trained model file from `trained_models` and use in your favorite Stable Diffusion repo!"
]
},
{
"cell_type": "markdown",
"id": "d28d0139",
"metadata": {},
"source": [
"## Generate Images With Your Trained Model!"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "80ddb03b",
"metadata": {},
"outputs": [],
"source": [
"!python scripts/stable_txt2img.py \\\n",
" --ddim_eta 0.0 \\\n",
" --n_samples 1 \\\n",
" --n_iter 4 \\\n",
" --scale 7.0 \\\n",
" --ddim_steps 50 \\\n",
" --ckpt \"/workspace/Dreambooth-Stable-Diffusion/trained_models/\" + {file_name} \\\n",
" --prompt \"joepenna person as a masterpiece portrait painting by John Singer Sargent in the style of Rembrandt\""
]
}
],
"metadata": {
"colab": {
"collapsed_sections": [],
"provenance": []
},
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.6"
},
"vscode": {
"interpreter": {
"hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
}
}
},
"nbformat": 4,
"nbformat_minor": 5
}

27
main.py
View File

@ -21,6 +21,9 @@ from pytorch_lightning.utilities import rank_zero_info
from ldm.data.base import Txt2ImgIterableBaseDataset
from ldm.util import instantiate_from_config
## Un-comment this for windows
## os.environ["PL_TORCH_DISTRIBUTED_BACKEND"] = "gloo"
def load_model_from_config(config, ckpt, verbose=False):
print(f"Loading model from {ckpt}")
pl_sd = torch.load(ckpt, map_location="cpu")
@ -139,13 +142,20 @@ def get_parser(**parser_kwargs):
)
parser.add_argument(
"--datadir_in_name",
type=str2bool,
nargs="?",
const=True,
default=True,
"--datadir_in_name",
type=str2bool,
nargs="?",
const=True,
default=True,
help="Prepend the final directory in the data_root to the output directory name")
parser.add_argument(
"--batch_size",
type=int,
required=True,
default=1000,
help="Number of iterations to run")
parser.add_argument("--actual_resume",
type=str,
required=True,
@ -605,8 +615,15 @@ if __name__ == "__main__":
cli = OmegaConf.from_dotlist(unknown)
config = OmegaConf.merge(*configs, cli)
lightning_config = config.pop("lightning", OmegaConf.create())
# merge trainer cli with config
trainer_config = lightning_config.get("trainer", OmegaConf.create())
# Set the steps
trainer_config.max_steps = opt.batch_size
# default to ddp
trainer_config["accelerator"] = "ddp"
for k in nondefault_trainer_args(opt):
trainer_config[k] = getattr(opt, k)
if not "gpus" in trainer_config: