Skip to content

Instantly share code, notes, and snippets.

@pszemraj
Created February 12, 2022 01:53
Show Gist options
  • Save pszemraj/dc411efb2ea28385bc0cc28c43fe4c1e to your computer and use it in GitHub Desktop.
Save pszemraj/dc411efb2ea28385bc0cc28c43fe4c1e to your computer and use it in GitHub Desktop.
v-diffusion-art
Display the source blob
Display the rendered blob
Raw
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "v-diffusion-art",
"provenance": [],
"collapsed_sections": [],
"toc_visible": true,
"authorship_tag": "ABX9TyNAZ+A6oIcY2TxHQ7fG3LMe",
"include_colab_link": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"language_info": {
"name": "python"
},
"accelerator": "GPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/pszemraj/dc411efb2ea28385bc0cc28c43fe4c1e/v-diffusion-art.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"source": [
"# <center> Art with Diffusion </center>\n",
"\n",
"- [repo](https://github.com/crowsonkb/v-diffusion-pytorch) by katherine crowson"
],
"metadata": {
"id": "2c7EwjbLj2hZ"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "LrDWdEzv3LaX"
},
"outputs": [],
"source": [
"#@markdown add auto-Colab formatting with `IPython.display`\n",
"from IPython.display import HTML, display\n",
"# colab formatting\n",
"def set_css():\n",
" display(\n",
" HTML(\n",
" \"\"\"\n",
" <style>\n",
" pre {\n",
" white-space: pre-wrap;\n",
" }\n",
" </style>\n",
" \"\"\"\n",
" )\n",
" )\n",
"\n",
"get_ipython().events.register(\"pre_run_cell\", set_css)"
]
},
{
"cell_type": "code",
"source": [
"#@markdown check gpu\n",
"!nvidia-smi"
],
"metadata": {
"cellView": "form",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 364
},
"id": "yFnaOsAjaYsS",
"outputId": "18e2210f-9a7a-411b-e5c0-66f0856fee8c"
},
"execution_count": null,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/html": [
"\n",
" <style>\n",
" pre {\n",
" white-space: pre-wrap;\n",
" }\n",
" </style>\n",
" "
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {}
},
{
"output_type": "stream",
"name": "stdout",
"text": [
"Sat Feb 12 00:28:31 2022 \n",
"+-----------------------------------------------------------------------------+\n",
"| NVIDIA-SMI 460.32.03 Driver Version: 460.32.03 CUDA Version: 11.2 |\n",
"|-------------------------------+----------------------+----------------------+\n",
"| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |\n",
"| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |\n",
"| | | MIG M. |\n",
"|===============================+======================+======================|\n",
"| 0 Tesla V100-SXM2... Off | 00000000:00:04.0 Off | 0 |\n",
"| N/A 33C P0 24W / 300W | 0MiB / 16160MiB | 0% Default |\n",
"| | | N/A |\n",
"+-------------------------------+----------------------+----------------------+\n",
" \n",
"+-----------------------------------------------------------------------------+\n",
"| Processes: |\n",
"| GPU GI CI PID Type Process name GPU Memory |\n",
"| ID ID Usage |\n",
"|=============================================================================|\n",
"| No running processes found |\n",
"+-----------------------------------------------------------------------------+\n"
]
}
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 34
},
"id": "B0Z__yyRWBcQ",
"outputId": "7c35579e-2c48-4ff2-c930-30517834e70e"
},
"outputs": [
{
"output_type": "display_data",
"data": {
"text/html": [
"\n",
" <style>\n",
" pre {\n",
" white-space: pre-wrap;\n",
" }\n",
" </style>\n",
" "
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {}
},
{
"output_type": "stream",
"name": "stdout",
"text": [
"fatal: destination path 'diffusion_repo' already exists and is not an empty directory.\n"
]
}
],
"source": [
"!git clone https://github.com/crowsonkb/v-diffusion-pytorch.git --recursive diffusion_repo\n",
"\n",
"import os\n",
"os.chdir('/content/diffusion_repo')"
]
},
{
"cell_type": "code",
"source": [
"!pip install -r requirements.txt -U -q\n",
"from tqdm.auto import tqdm"
],
"metadata": {
"id": "7tBLNOxUWt5f",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"outputId": "5a82918d-5133-4d8f-d594-75dbe81a6ae6"
},
"execution_count": null,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/html": [
"\n",
" <style>\n",
" pre {\n",
" white-space: pre-wrap;\n",
" }\n",
" </style>\n",
" "
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {}
}
]
},
{
"cell_type": "code",
"source": [
"from pathlib import Path \n",
"_chk_dir = Path('/content/diffusion_repo/checkpoints/')\n",
"_chk_dir.mkdir(exist_ok=True)"
],
"metadata": {
"id": "c0rm1Z5WZoFd",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"outputId": "c04e0284-4889-4241-d844-960432fe0947"
},
"execution_count": null,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/html": [
"\n",
" <style>\n",
" pre {\n",
" white-space: pre-wrap;\n",
" }\n",
" </style>\n",
" "
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {}
}
]
},
{
"cell_type": "code",
"source": [
"!wget https://v-diffusion.s3.us-west-2.amazonaws.com/cc12m_1_cfg.pth -O /content/diffusion_repo/checkpoints/cc12m_1_cfg.pth\n",
"\n",
"# !wget https://v-diffusion.s3.us-west-2.amazonaws.com/yfcc_2.pth -O /content/diffusion_repo/checkpoints/yfcc_2.pth"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 208
},
"id": "KJzrCQZ5ZbBS",
"outputId": "1695827f-60db-4a77-f2d1-da2299d2c1db"
},
"execution_count": null,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/html": [
"\n",
" <style>\n",
" pre {\n",
" white-space: pre-wrap;\n",
" }\n",
" </style>\n",
" "
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {}
},
{
"output_type": "stream",
"name": "stdout",
"text": [
"--2022-02-12 00:28:39-- https://v-diffusion.s3.us-west-2.amazonaws.com/cc12m_1_cfg.pth\n",
"Resolving v-diffusion.s3.us-west-2.amazonaws.com (v-diffusion.s3.us-west-2.amazonaws.com)... 52.218.246.193\n",
"Connecting to v-diffusion.s3.us-west-2.amazonaws.com (v-diffusion.s3.us-west-2.amazonaws.com)|52.218.246.193|:443... connected.\n",
"HTTP request sent, awaiting response... 200 OK\n",
"Length: 2411747217 (2.2G) [application/octet-stream]\n",
"Saving to: ‘/content/diffusion_repo/checkpoints/cc12m_1_cfg.pth’\n",
"\n",
"/content/diffusion_ 100%[===================>] 2.25G 22.2MB/s in 1m 45s \n",
"\n",
"2022-02-12 00:30:25 (21.8 MB/s) - ‘/content/diffusion_repo/checkpoints/cc12m_1_cfg.pth’ saved [2411747217/2411747217]\n",
"\n"
]
}
]
},
{
"cell_type": "markdown",
"source": [
"# run script"
],
"metadata": {
"id": "AfflkCEEjxut"
}
},
{
"cell_type": "code",
"source": [
"#@markdown get inputs\n",
"import random\n",
"text_prompt = \"The Museum of Unconditional Love. concept art.\" #@param {type:\"string\"}\n",
"NUM_IMAGES = 2#@param {type:\"integer\"}\n",
"BATCH_SIZE = 1#@param {type:\"integer\"}\n",
"OPT_STEPS = 500#@param {type:\"integer\"}\n",
"IMG_DIM = 512#@param {type:\"integer\"}\n",
"SEED = random.SystemRandom().randint(1, 10**8)\n",
"print(f'generated seed: {SEED}')\n",
"import torch\n",
"torch.cuda.empty_cache()"
],
"metadata": {
"id": "ePCrdUTpnR25",
"cellView": "form",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 34
},
"outputId": "3b451f2a-03ee-41a1-87fb-79a10a14cc5b"
},
"execution_count": null,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/html": [
"\n",
" <style>\n",
" pre {\n",
" white-space: pre-wrap;\n",
" }\n",
" </style>\n",
" "
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {}
},
{
"output_type": "stream",
"name": "stdout",
"text": [
"generated seed: 10714310\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"!./cfg_sample.py $text_prompt:4 -n $NUM_IMAGES -bs $BATCH_SIZE \\\n",
" --steps $OPT_STEPS --seed $SEED \\\n",
" --size $IMG_DIM $IMG_DIM \\\n",
" --model cc12m_1_cfg --checkpoint /content/diffusion_repo/checkpoints/cc12m_1_cfg.pth"
],
"metadata": {
"id": "0MZYhEolhGFq"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"import os\n",
"import shutil\n",
"import zipfile\n",
"from os.path import join\n",
"from datetime import datetime\n",
"\n",
"#@markdown define `export_imgs` and helpers\n",
"def get_timestamp():\n",
" \"\"\"\n",
" get_timestamp - returns the current timestamp in the format YYYY-MM-DD-HH-MM-SS\n",
" :return: the current timestamp\n",
" \"\"\"\n",
" return datetime.now().strftime('%Y-%m-%d-%H-%M-%S')\n",
"\n",
"def export_imgs(img_dir, out_dir=None, \n",
" prompt:str=\"generated-img\", img_ext='.png',\n",
" ):\n",
" \"\"\"\n",
" export_imgs - finds all images in a directory, zips them together into a zip file, then deletes the originals\n",
" :param img_dir: the directory to find images in\n",
" :return: filepath to the created zip\n",
" \"\"\"\n",
" out_dir = img_dir if out_dir is None else out_dir\n",
" # get all files in the directory\n",
" files = os.listdir(img_dir)\n",
"\n",
" # create a zip file\n",
" zip_name = '{}-{}.zip'.format(prompt, get_timestamp())\n",
" zip_file = zipfile.ZipFile(join(out_dir, zip_name), 'w')\n",
"\n",
" # add all files to the zip file\n",
" img_files = [f for f in files if f.endswith('.jpg') or f.endswith(img_ext)]\n",
" renamed_imgs = []\n",
" for i, file in enumerate(img_files):\n",
" _src_img = join(img_dir, file)\n",
" _ext = file.split('.')[-1]\n",
" _new_name = f\"{prompt}_{i}.{_ext}\"\n",
" _target = join(img_dir, _new_name)\n",
" os.rename(_src_img, _target)\n",
" renamed_imgs.append(_target)\n",
" zip_file.write(_target)\n",
"\n",
" # close the zip file\n",
" zip_file.close()\n",
"\n",
" # delete the original files\n",
" for file in renamed_imgs:\n",
" os.remove(file)\n",
"\n",
" # return the path of the zip file\n",
" return join(out_dir, zip_name)"
],
"metadata": {
"cellView": "form",
"id": "rHXM860jYpQp"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"from google.colab import files\n",
"results = export_imgs(img_dir='/content/diffusion_repo',\n",
" out_dir='/content',\n",
" prompt=text_prompt,\n",
" )\n"
],
"metadata": {
"id": "ny3hoSAIiw_n"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"from google.colab import files\n",
"files.download(results)"
],
"metadata": {
"id": "5qbwqhzkjJA1"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
""
],
"metadata": {
"id": "bNP_sCiUkRMJ"
},
"execution_count": null,
"outputs": []
}
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment