{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "name": "StyleGAN3 Sampling + Comet.ipynb", "provenance": [], "collapsed_sections": [] }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "markdown", "metadata": { "id": "aHzltei9Ue7I" }, "source": [ "#Sample from StyleGAN3 and log results to Comet" ] }, { "cell_type": "markdown", "metadata": { "id": "WVF2P1OlT92b" }, "source": [ "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/comet-ml/comet-examples/blob/ml-art/notebooks/Comet_StyleGAN3_Sampling.ipynb)", "\n", "\n", "**Overview**\n", "\n", "This notebook expands on a notebook by [crimeacs](https://twitter.com/EarthML1). We have added some functionality to demostrate how Comet Panels can be used to track parameters and model configurations as well as log/view images and video for deep learning and computer vision projects and ML art. \n", "\n", "First, we're embedding some panels from an existing Comet project to display our own previous outputs from running this notebook. Then, we added a cell to install and initialize Comet so that you can log your results to your own Comet project. Then, we reproduce the original notebook code, but with Comet logging (optionally) enabled. Finally, we include insructions for adding an Image Viewer or Video Viewer panel to your Comet project and embedding it in a notebook like this one. \n", "\n", "
\n", "\n", "\n", "**Original intro by [crimeacs](https://twitter.com/EarthML1)**\n", "\n", "This Colab is to try out **StyleGAN3** (aka Alias-Free GAN) released in [this repo](https://github.com/NVlabs/stylegan3) by NVidia. Colab produced by [crimeacs](https://twitter.com/EarthML1). Also check out my [telegram channel](https://t.me/nn_for_science) about Deep Learning for Science.\n", "\n" ] }, { "cell_type": "markdown", "metadata": { "id": "YC4Y2sOMeGSI" }, "source": [ "## Example Results - Comet:" ] }, { "cell_type": "code", "metadata": { "id": "HEaeYEdWiKPY", "cellView": "form" }, "source": [ "#@title Generated Images\n", "#@title Display utility functions\n", "from IPython import display\n", "def make_iframe(src):\n", " return f''''''\n", "\n", "\n", "iframe = make_iframe('https://www.comet.ml/embedded/?instanceId=PSMnawkGLivaljLupbmXYeptA&projectId=1dc7a6ddf9a444178f408f49a530035f&templateId=TrYNymVHvwNPgS84AzLhvtPEl&viewId=5IbNaVEfZFWBjdPQTHJkqpHXZ')\n", "\n", "display.HTML(iframe)" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "gKwFuK79eY4N", "cellView": "form" }, "source": [ "#@title Generate Videos\n", "\n", "iframe = make_iframe(\"https://www.comet.ml/embedded/?instanceId=qvKKiuyRQLXAkgwnwJOu4R5g5&projectId=1dc7a6ddf9a444178f408f49a530035f&templateId=xDJ93UCEHhXQ82bzdJnPjidQH&viewId=5IbNaVEfZFWBjdPQTHJkqpHXZ\")\n", "display.HTML(iframe)" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "iMcpPb23jAIF" }, "source": [ "## Make your own!" ] }, { "cell_type": "code", "metadata": { "id": "d99HsTjTQRzg", "cellView": "form" }, "source": [ "#@title Install StyleGAN3 dependencies\n", "from IPython.display import clear_output\n", "\n", "!git clone https://github.com/NVlabs/stylegan3.git\n", "%cd stylegan3\n", "!wget -O mini.sh https://repo.anaconda.com/miniconda/Miniconda3-py38_4.8.2-Linux-x86_64.sh\n", "!chmod +x mini.sh\n", "!bash ./mini.sh -b -f -p /usr/local\n", "!conda install -q -y --prefix /usr/local jupyter\n", "!python -m ipykernel install --name \"py38\" --user\n", "!pip install click -q\n", "!pip install numpy -q\n", "!pip install pillow -q\n", "!pip install torch -q\n", "!pip install scipy -q\n", "!pip install Ninja -q\n", "!pip install imageio -q\n", "!pip install imageio-ffmpeg -q\n", "clear_output()\n" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "zikJHylIjCfV", "cellView": "form" }, "source": [ "#@title Install / Initialize Comet\n", "\n", "#@markdown Running this cell will install comet and run `comet_ml.init()`, which will prompt you to log in and copy/paste your API key, or to create an account if you do not already have one.\n", "#@markdown If you have ever synced Google Drive to a Colab notebook, this process should feel familiar. \n", "\n", "#@markdown **Note:** You can use this notebook without installing and signing up for Comet. You simply won't be able to log your images/videos to view later through the Comet UI. If you want to proceed without Comet, uncheck the checkbox below before running this cell.\n", "\n", "\n", "use_comet = True #@param {type:\"boolean\"}\n", "\n", "if use_comet:\n", " !pip3 install --quiet comet_ml\n", " import comet_ml\n", " comet_ml.init()\n", "\n", "comet_project_name = \"stylegan3-colab-demo\" #@param {type:\"string\"}\n", "\n", "\n" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "IkBzNIQ9QsFB", "cellView": "form" }, "source": [ "#@title Generate an image\n", "#@markdown StyleGAN3 pre-trained models for config T (translation equiv.) and config R (translation and rotation equiv.)\n", "seed = 3454 #@param {type:\"slider\", min:0, max:9999, step:1}\n", "\n", "\n", "baselink ='https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/'\n", "model = \"stylegan3-t-metfaces-1024x1024.pkl\" #@param [\"stylegan3-r-afhqv2-512x512.pkl\", \"stylegan3-r-ffhq-1024x1024.pkl\", \"stylegan3-r-ffhqu-1024x1024.pkl\",\"stylegan3-r-ffhqu-256x256.pkl\",\"stylegan3-r-metfaces-1024x1024.pkl\",\"stylegan3-r-metfacesu-1024x1024.pkl\",\"stylegan3-t-afhqv2-512x512.pkl\",\"stylegan3-t-ffhq-1024x1024.pkl\",\"stylegan3-t-ffhqu-1024x1024.pkl\",\"stylegan3-t-ffhqu-256x256.pkl\",\"stylegan3-t-metfaces-1024x1024.pkl\",\"stylegan3-t-metfacesu-1024x1024.pkl\"]\n", "\n", "if use_comet:\n", " experiment = comet_ml.Experiment(project_name=comet_project_name)\n", " params = {\n", " 'model':model,\n", " 'seed':seed,\n", " 'task':'image'\n", " }\n", " experiment.log_parameters(params)\n", " \n", "# Generate an image using pre-trained AFHQv2 model (\"Ours\" in Figure 1, left).\n", "!python gen_images.py --outdir=/content/stylegan3/out --trunc=1 \\\n", "--seeds=$seed --network=$baselink$model\n", "\n", "if use_comet:\n", " image_path = f'{model.replace(\".pkl\", \"\")}__seed{seed:04d}.png'\n", " experiment.log_image(image_data='/content/stylegan3/out/seed%04d.png' % seed, name=image_path.split('/')[-1])\n", " experiment.end()\n", "\n", "import matplotlib.pyplot as plt\n", "from PIL import Image\n", "\n", "plt.figure(figsize=(10,10))\n", "img = Image.open('/content/stylegan3/out/seed%04d.png' % seed);\n", "plt.imshow(img);\n", "plt.axis('off');\n" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "ZREJuYfjbV3G", "cellView": "form" }, "source": [ "#@title Generate an interpolation video\n", "%cd /content/stylegan3\n", "\n", "start_seed = 4#@param {type:\"number\"}\n", "stop_seed = 2235 #@param {type:\"number\"}\n", "n_cols = 2#@param {type:\"number\"}\n", "n_rows = 1#@param {type:\"number\"}\n", "\n", "#@markdown How many key frames to have?\n", "num_keyframes = 3#@param {type:\"number\"}\n", "\n", "#@markdown How many frames for interpolation?\n", "w_frames = 20#@param {type:\"number\"}\n", "\n", "#@markdown Total length in frames is `num_keyframes`*`w_frames`\n", "\n", "assert stop_seed > start_seed, 'Stop_seed should be larger then start_seed'\n", "baselink ='https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/'\n", "model = \"stylegan3-r-metfacesu-1024x1024.pkl\" #@param [\"stylegan3-r-afhqv2-512x512.pkl\", \"stylegan3-r-ffhq-1024x1024.pkl\", \"stylegan3-r-ffhqu-1024x1024.pkl\",\"stylegan3-r-ffhqu-256x256.pkl\",\"stylegan3-r-metfaces-1024x1024.pkl\",\"stylegan3-r-metfacesu-1024x1024.pkl\",\"stylegan3-t-afhqv2-512x512.pkl\",\"stylegan3-t-ffhq-1024x1024.pkl\",\"stylegan3-t-ffhqu-1024x1024.pkl\",\"stylegan3-t-ffhqu-256x256.pkl\",\"stylegan3-t-metfaces-1024x1024.pkl\",\"stylegan3-t-metfacesu-1024x1024.pkl\"]\n", "if use_comet:\n", " experiment = comet_ml.Experiment(project_name=comet_project_name)\n", " params = {\n", " 'start_seed':start_seed,\n", " 'stop_seed':stop_seed,\n", " 'n_cols':n_cols,\n", " 'n_rows':n_rows,\n", " 'num_keyframes':num_keyframes,\n", " 'w_frames':w_frames,\n", " 'baselink':baselink,\n", " 'model':model,\n", " 'task':'video'\n", " }\n", " experiment.log_parameters(params)\n", "\n", "# Generate an image using pre-trained AFHQv2 model (\"Ours\" in Figure 1, left).\n", "output_video = f'start_{start_seed}_stop_{stop_seed}.mp4'\n", "\n", "# Render a grid of interpolations for seeds N through K.\n", "!python gen_video.py --output=$output_video --trunc=1 --seeds=$start_seed-$stop_seed --grid={n_cols}x{n_rows} \\\n", " --network=$baselink$model --num-keyframes=$num_keyframes \\\n", " --w-frames=$w_frames\n", "\n", "if use_comet:\n", " experiment.log_asset(output_video, file_name=f\"{model.replace('.pkl', '')}_{output_video}\")\n", " experiment.end()\n", "\n", "\n", "\n", "from IPython.display import HTML\n", "from base64 import b64encode\n", "mp4 = open(output_video,'rb').read()\n", "\n", "\n", "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n", "HTML(\"\"\"\n", "\n", "\"\"\" % data_url)" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "FgQfYTOUN0qu" }, "source": [ "## Create your own Comet panels! " ] }, { "cell_type": "markdown", "metadata": { "id": "0GBCL3WVOyEu" }, "source": [ "1. Run the cell below to display the Comet UI in an IFrame\n", "2. Click the button that says \"Add Panel\"\n", "3. Click the \"Public\" tab, then locate the row for \"Comet Image Viewer\" or \"Comet Video Viewer - alpha\" and click \"Add\". \n", "4. If you'd like, fill in panel options in the next screen. For example, the \"Comet Image Viewer\" panel accepts the following two options:\n", " - `nColumns`: The number of columns in the image grid.\n", " - `nImages`: The number of Images to Display inside the Panel. \n", " \n", " So you could fill in something like this, for example:\n", "```\n", " {\n", " \"nColumns\": \"4\",\n", " \"nImages\": \"24\"\n", " }\n", "```\n", "5. Click \"Done\" to add the column to your view and close the pop up modal. \n", "6. Press \"Save View\" in the upper-right region of the Project view.\n", "7. Now, you can view your updated Panel at any time at the Project view page, and you can also embed this Panel in this notebook (or somewhere else!) if you would like. Click the three dots in the top-right corner of the panel to bring up a popup menu, then select \"Share\" or \"Embed\" depending on what you want to do. To embed the panel in a notebook, you can do the following: \n", "\n", "```python\n", "from IPython.display import HTML\n", "\n", "HTML(embed_code) \n", "```\n", "where `embed_code` is the HTML you copy/pasted from the Embed popup modal. " ] }, { "cell_type": "code", "metadata": { "id": "VFNk6XCw9K5n" }, "source": [ "\n", "experiment.display_project()" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "0v08imk6jgbq" }, "source": [ "" ], "execution_count": null, "outputs": [] } ] }