Open In Colab

Using xtal2png with imagen-pytorch

[1]:

!nvidia-smi

Sat Oct  8 06:26:18 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  A100-SXM4-40GB      Off  | 00000000:00:04.0 Off |                    0 |
| N/A   29C    P0    43W / 400W |      0MiB / 40536MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

[2]:

dummy = False #@param {type:"boolean"}

[3]:

from google.colab import drive
drive.mount('/content/drive')
%pip install mp-time-split xtal2png matbench-genmetrics imagen-pytorch==1.11.4

Mounted at /content/drive
Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/
Collecting mp-time-split
  Downloading mp_time_split-0.2.0-py3-none-any.whl (38 kB)
Collecting xtal2png
  Downloading xtal2png-0.9.4-py3-none-any.whl (34 kB)
Collecting matbench-genmetrics
  Downloading matbench_genmetrics-0.6.0-py3-none-any.whl (17 kB)
Collecting imagen-pytorch==1.11.4
  Downloading imagen_pytorch-1.11.4-py3-none-any.whl (55 kB)
     |████████████████████████████████| 55 kB 3.4 MB/s
Requirement already satisfied: torch>=1.6 in /usr/local/lib/python3.7/dist-packages (from imagen-pytorch==1.11.4) (1.12.1+cu113)
Collecting pytorch-lightning
  Downloading pytorch_lightning-1.7.7-py3-none-any.whl (708 kB)
     |████████████████████████████████| 708 kB 10.3 MB/s
Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from imagen-pytorch==1.11.4) (4.64.1)
Collecting pytorch-warmup
  Downloading pytorch-warmup-0.1.0.tar.gz (314 kB)
     |████████████████████████████████| 314 kB 75.9 MB/s
Collecting ema-pytorch>=0.0.3
  Downloading ema_pytorch-0.0.10-py3-none-any.whl (4.1 kB)
Collecting einops-exts
  Downloading einops_exts-0.0.3-py3-none-any.whl (3.8 kB)
Requirement already satisfied: pillow in /usr/local/lib/python3.7/dist-packages (from imagen-pytorch==1.11.4) (7.1.2)
Collecting einops>=0.4
  Downloading einops-0.5.0-py3-none-any.whl (36 kB)
Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from imagen-pytorch==1.11.4) (21.3)
Requirement already satisfied: click in /usr/local/lib/python3.7/dist-packages (from imagen-pytorch==1.11.4) (7.1.2)
Collecting kornia
  Downloading kornia-0.6.7-py2.py3-none-any.whl (565 kB)
     |████████████████████████████████| 565 kB 62.1 MB/s
Requirement already satisfied: pydantic in /usr/local/lib/python3.7/dist-packages (from imagen-pytorch==1.11.4) (1.9.2)
Collecting transformers
  Downloading transformers-4.22.2-py3-none-any.whl (4.9 MB)
     |████████████████████████████████| 4.9 MB 78.0 MB/s
Collecting accelerate
  Downloading accelerate-0.13.1-py3-none-any.whl (148 kB)
     |████████████████████████████████| 148 kB 79.4 MB/s
Requirement already satisfied: torchvision in /usr/local/lib/python3.7/dist-packages (from imagen-pytorch==1.11.4) (0.13.1+cu113)
Collecting sentencepiece
  Downloading sentencepiece-0.1.97-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)
     |████████████████████████████████| 1.3 MB 73.2 MB/s
Requirement already satisfied: fsspec in /usr/local/lib/python3.7/dist-packages (from imagen-pytorch==1.11.4) (2022.8.2)
Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from imagen-pytorch==1.11.4) (1.21.6)
Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.6->imagen-pytorch==1.11.4) (4.1.1)
Requirement already satisfied: scikit-learn in /usr/local/lib/python3.7/dist-packages (from mp-time-split) (1.0.2)
Collecting pybtex
  Downloading pybtex-0.24.0-py2.py3-none-any.whl (561 kB)
     |████████████████████████████████| 561 kB 75.4 MB/s
Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from mp-time-split) (5.0.0)
Collecting matminer
  Downloading matminer-0.7.8-py3-none-any.whl (1.4 MB)
     |████████████████████████████████| 1.4 MB 74.7 MB/s
Collecting pymatgen
  Downloading pymatgen-2022.0.17.tar.gz (40.6 MB)
     |████████████████████████████████| 40.6 MB 1.3 MB/s
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Installing backend dependencies ... done
    Preparing wheel metadata ... done
Collecting element-coder
  Downloading element_coder-0.0.6-py3-none-any.whl (1.1 MB)
     |████████████████████████████████| 1.1 MB 73.2 MB/s
Collecting kaleido
  Downloading kaleido-0.2.1-py2.py3-none-manylinux1_x86_64.whl (79.9 MB)
     |████████████████████████████████| 79.9 MB 1.1 MB/s
Requirement already satisfied: plotly in /usr/local/lib/python3.7/dist-packages (from xtal2png) (5.5.0)
Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from matbench-genmetrics) (1.7.3)
Collecting pystow
  Downloading pystow-0.4.6-py3-none-any.whl (35 kB)
Requirement already satisfied: psutil in /usr/local/lib/python3.7/dist-packages (from accelerate->imagen-pytorch==1.11.4) (5.4.8)
Requirement already satisfied: pyyaml in /usr/local/lib/python3.7/dist-packages (from accelerate->imagen-pytorch==1.11.4) (6.0)
Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->imagen-pytorch==1.11.4) (3.0.9)
Collecting loguru
  Downloading loguru-0.6.0-py3-none-any.whl (58 kB)
     |████████████████████████████████| 58 kB 6.3 MB/s
Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->mp-time-split) (3.8.1)
Collecting matminer
  Downloading matminer-0.7.7.tar.gz (5.2 MB)
     |████████████████████████████████| 5.2 MB 68.9 MB/s
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Installing backend dependencies ... done
    Preparing wheel metadata ... done
  Downloading matminer-0.7.6-py3-none-any.whl (1.4 MB)
     |████████████████████████████████| 1.4 MB 55.7 MB/s
  Downloading matminer-0.7.5-py3-none-any.whl (1.4 MB)
     |████████████████████████████████| 1.4 MB 65.9 MB/s
  Downloading matminer-0.7.4-py3-none-any.whl (1.4 MB)
     |████████████████████████████████| 1.4 MB 69.8 MB/s
Collecting requests>=2.26.0
  Downloading requests-2.28.1-py3-none-any.whl (62 kB)
     |████████████████████████████████| 62 kB 1.6 MB/s
Requirement already satisfied: pandas>=1.3.1 in /usr/local/lib/python3.7/dist-packages (from matminer->mp-time-split) (1.3.5)
Collecting future>=0.18.2
  Downloading future-0.18.2.tar.gz (829 kB)
     |████████████████████████████████| 829 kB 69.3 MB/s
Requirement already satisfied: pymongo>=3.12.0 in /usr/local/lib/python3.7/dist-packages (from matminer->mp-time-split) (4.2.0)
Collecting pint>=0.17
  Downloading Pint-0.18-py2.py3-none-any.whl (209 kB)
     |████████████████████████████████| 209 kB 75.3 MB/s
Collecting six>=1.16.0
  Downloading six-1.16.0-py2.py3-none-any.whl (11 kB)
Collecting monty>=2021.7.8
  Downloading monty-2022.9.9-py3-none-any.whl (66 kB)
     |████████████████████████████████| 66 kB 5.8 MB/s
Collecting sympy>=1.8
  Downloading sympy-1.10.1-py3-none-any.whl (6.4 MB)
     |████████████████████████████████| 6.4 MB 53.2 MB/s
Requirement already satisfied: jsonschema>=3.2.0 in /usr/local/lib/python3.7/dist-packages (from matminer->mp-time-split) (4.3.3)
Requirement already satisfied: importlib-resources>=1.4.0 in /usr/local/lib/python3.7/dist-packages (from jsonschema>=3.2.0->matminer->mp-time-split) (5.9.0)
Requirement already satisfied: pyrsistent!=0.17.0,!=0.17.1,!=0.17.2,>=0.14.0 in /usr/local/lib/python3.7/dist-packages (from jsonschema>=3.2.0->matminer->mp-time-split) (0.18.1)
Requirement already satisfied: attrs>=17.4.0 in /usr/local/lib/python3.7/dist-packages (from jsonschema>=3.2.0->matminer->mp-time-split) (22.1.0)
Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.7/dist-packages (from pandas>=1.3.1->matminer->mp-time-split) (2.8.2)
Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas>=1.3.1->matminer->mp-time-split) (2022.4)
Collecting spglib>=1.9.9.44
  Downloading spglib-2.0.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (514 kB)
     |████████████████████████████████| 514 kB 79.1 MB/s
Requirement already satisfied: palettable>=3.1.1 in /usr/local/lib/python3.7/dist-packages (from pymatgen->mp-time-split) (3.3.0)
Requirement already satisfied: networkx>=2.2 in /usr/local/lib/python3.7/dist-packages (from pymatgen->mp-time-split) (2.6.3)
Collecting uncertainties>=3.1.4
  Downloading uncertainties-3.1.7-py2.py3-none-any.whl (98 kB)
     |████████████████████████████████| 98 kB 8.1 MB/s
Requirement already satisfied: matplotlib>=1.5 in /usr/local/lib/python3.7/dist-packages (from pymatgen->mp-time-split) (3.2.2)
Requirement already satisfied: tabulate in /usr/local/lib/python3.7/dist-packages (from pymatgen->mp-time-split) (0.8.10)
Collecting ruamel.yaml>=0.15.6
  Downloading ruamel.yaml-0.17.21-py3-none-any.whl (109 kB)
     |████████████████████████████████| 109 kB 82.9 MB/s
Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib>=1.5->pymatgen->mp-time-split) (0.11.0)
Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib>=1.5->pymatgen->mp-time-split) (1.4.4)
Requirement already satisfied: tenacity>=6.2.0 in /usr/local/lib/python3.7/dist-packages (from plotly->xtal2png) (8.1.0)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests>=2.26.0->matminer->mp-time-split) (2022.9.24)
Requirement already satisfied: charset-normalizer<3,>=2 in /usr/local/lib/python3.7/dist-packages (from requests>=2.26.0->matminer->mp-time-split) (2.1.1)
Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests>=2.26.0->matminer->mp-time-split) (2.10)
Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests>=2.26.0->matminer->mp-time-split) (1.24.3)
Collecting ruamel.yaml.clib>=0.2.6
  Downloading ruamel.yaml.clib-0.2.6-cp37-cp37m-manylinux1_x86_64.whl (546 kB)
     |████████████████████████████████| 546 kB 73.2 MB/s
Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->mp-time-split) (1.2.0)
Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->mp-time-split) (3.1.0)
Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.7/dist-packages (from sympy>=1.8->matminer->mp-time-split) (1.2.1)
Collecting pyxtal
  Downloading pyxtal-0.5.3-py3-none-any.whl (2.7 MB)
     |████████████████████████████████| 2.7 MB 51.1 MB/s
Collecting latexcodec>=1.0.4
  Downloading latexcodec-2.0.1-py2.py3-none-any.whl (18 kB)
Collecting pickle5
  Downloading pickle5-0.0.12-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.whl (256 kB)
     |████████████████████████████████| 256 kB 86.5 MB/s
Collecting pyDeprecate>=0.3.1
  Downloading pyDeprecate-0.3.2-py3-none-any.whl (10 kB)
Collecting tensorboard>=2.9.1
  Downloading tensorboard-2.10.1-py3-none-any.whl (5.9 MB)
     |████████████████████████████████| 5.9 MB 47.2 MB/s
Collecting torchmetrics>=0.7.0
  Downloading torchmetrics-0.10.0-py3-none-any.whl (529 kB)
     |████████████████████████████████| 529 kB 86.5 MB/s
Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /usr/local/lib/python3.7/dist-packages (from fsspec->imagen-pytorch==1.11.4) (3.8.3)
Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.7/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec->imagen-pytorch==1.11.4) (1.2.0)
Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.7/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec->imagen-pytorch==1.11.4) (1.3.1)
Requirement already satisfied: asynctest==0.13.0 in /usr/local/lib/python3.7/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec->imagen-pytorch==1.11.4) (0.13.0)
Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.7/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec->imagen-pytorch==1.11.4) (6.0.2)
Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.7/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec->imagen-pytorch==1.11.4) (1.8.1)
Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /usr/local/lib/python3.7/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec->imagen-pytorch==1.11.4) (4.0.2)
Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (0.4.6)
Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (0.6.1)
Requirement already satisfied: protobuf<3.20,>=3.9.2 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (3.17.3)
Requirement already satisfied: grpcio>=1.24.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (1.49.1)
Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (1.35.0)
Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (0.37.1)
Requirement already satisfied: werkzeug>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (1.0.1)
Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (3.4.1)
Requirement already satisfied: absl-py>=0.4 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (1.2.0)
Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (57.4.0)
Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (1.8.1)
Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (4.9)
Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (4.2.4)
Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (0.2.8)
Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (1.3.1)
Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.7/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (0.4.8)
Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.9.1->pytorch-lightning->imagen-pytorch==1.11.4) (3.2.1)
Collecting pyshtools==4.10
  Downloading pyshtools-4.10-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (12.5 MB)
     |████████████████████████████████| 12.5 MB 79.2 MB/s
Requirement already satisfied: numba>=0.55.2 in /usr/local/lib/python3.7/dist-packages (from pyxtal->mp-time-split) (0.56.2)
Collecting ase>=3.18.0
  Downloading ase-3.22.1-py3-none-any.whl (2.2 MB)
     |████████████████████████████████| 2.2 MB 60.4 MB/s
Collecting py3Dmol>=0.8.0
  Downloading py3Dmol-1.8.1-py2.py3-none-any.whl (6.5 kB)
Requirement already satisfied: pooch>=1.1 in /usr/local/lib/python3.7/dist-packages (from pyshtools==4.10->pyxtal->mp-time-split) (1.6.0)
Collecting matplotlib>=1.5
  Downloading matplotlib-3.5.3-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.whl (11.2 MB)
     |████████████████████████████████| 11.2 MB 60.3 MB/s
Requirement already satisfied: xarray in /usr/local/lib/python3.7/dist-packages (from pyshtools==4.10->pyxtal->mp-time-split) (0.20.2)
Requirement already satisfied: astropy>=4.0 in /usr/local/lib/python3.7/dist-packages (from pyshtools==4.10->pyxtal->mp-time-split) (4.3.1)
Requirement already satisfied: pyerfa>=1.7.3 in /usr/local/lib/python3.7/dist-packages (from astropy>=4.0->pyshtools==4.10->pyxtal->mp-time-split) (2.0.0.1)
Collecting fonttools>=4.22.0
  Downloading fonttools-4.37.4-py3-none-any.whl (960 kB)
     |████████████████████████████████| 960 kB 81.0 MB/s
Requirement already satisfied: llvmlite<0.40,>=0.39.0dev0 in /usr/local/lib/python3.7/dist-packages (from numba>=0.55.2->pyxtal->mp-time-split) (0.39.1)
Requirement already satisfied: appdirs>=1.3.0 in /usr/local/lib/python3.7/dist-packages (from pooch>=1.1->pyshtools==4.10->pyxtal->mp-time-split) (1.4.4)
Requirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from transformers->imagen-pytorch==1.11.4) (3.8.0)
Collecting tokenizers!=0.11.3,<0.13,>=0.11.1
  Downloading tokenizers-0.12.1-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (6.6 MB)
     |████████████████████████████████| 6.6 MB 70.5 MB/s
Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.7/dist-packages (from transformers->imagen-pytorch==1.11.4) (2022.6.2)
Collecting huggingface-hub<1.0,>=0.9.0
  Downloading huggingface_hub-0.10.0-py3-none-any.whl (163 kB)
     |████████████████████████████████| 163 kB 84.2 MB/s
Building wheels for collected packages: future, pymatgen, pytorch-warmup
  Building wheel for future (setup.py) ... done
  Created wheel for future: filename=future-0.18.2-py3-none-any.whl size=491070 sha256=1e93095e5e16d1063c9d391f8c716beb11a4e887ed664ad7ff4d136ce9eb860b
  Stored in directory: /root/.cache/pip/wheels/56/b0/fe/4410d17b32f1f0c3cf54cdfb2bc04d7b4b8f4ae377e2229ba0
  Building wheel for pymatgen (PEP 517) ... done
  Created wheel for pymatgen: filename=pymatgen-2022.0.17-cp37-cp37m-linux_x86_64.whl size=41840996 sha256=bb86fef844e1545066942909cab2d4418d8532846023d8ad937ae4c8fff276c8
  Stored in directory: /root/.cache/pip/wheels/cf/f6/22/58a9be23c5f1b452770e02ff42047175eaf0f9c2f15219fc76
  Building wheel for pytorch-warmup (setup.py) ... done
  Created wheel for pytorch-warmup: filename=pytorch_warmup-0.1.0-py3-none-any.whl size=5802 sha256=d3def681fb97dfbb403ec42277fd2a3dd52940fd57847ea5f152d9c83556532c
  Stored in directory: /root/.cache/pip/wheels/c7/5d/24/9475e442daa1e9332c122c79fb5131b9e4e91946009365902e
Successfully built future pymatgen pytorch-warmup
Installing collected packages: six, ruamel.yaml.clib, future, fonttools, uncertainties, sympy, spglib, ruamel.yaml, requests, monty, matplotlib, pymatgen, pint, latexcodec, pyshtools, pybtex, py3Dmol, matminer, ase, torchmetrics, tokenizers, tensorboard, pyxtal, pyDeprecate, pickle5, mp-time-split, loguru, huggingface-hub, einops, transformers, sentencepiece, pytorch-warmup, pytorch-lightning, pystow, kornia, kaleido, ema-pytorch, element-coder, einops-exts, accelerate, xtal2png, matbench-genmetrics, imagen-pytorch
  Attempting uninstall: six
    Found existing installation: six 1.15.0
    Uninstalling six-1.15.0:
      Successfully uninstalled six-1.15.0
  Attempting uninstall: future
    Found existing installation: future 0.16.0
    Uninstalling future-0.16.0:
      Successfully uninstalled future-0.16.0
  Attempting uninstall: sympy
    Found existing installation: sympy 1.7.1
    Uninstalling sympy-1.7.1:
      Successfully uninstalled sympy-1.7.1
  Attempting uninstall: requests
    Found existing installation: requests 2.23.0
    Uninstalling requests-2.23.0:
      Successfully uninstalled requests-2.23.0
  Attempting uninstall: matplotlib
    Found existing installation: matplotlib 3.2.2
    Uninstalling matplotlib-3.2.2:
      Successfully uninstalled matplotlib-3.2.2
  Attempting uninstall: tensorboard
    Found existing installation: tensorboard 2.8.0
    Uninstalling tensorboard-2.8.0:
      Successfully uninstalled tensorboard-2.8.0
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
ipython 7.9.0 requires jedi>=0.10, which is not installed.
tensorflow 2.8.2+zzzcolab20220929150707 requires tensorboard<2.9,>=2.8, but you have tensorboard 2.10.1 which is incompatible.
Successfully installed accelerate-0.13.1 ase-3.22.1 einops-0.5.0 einops-exts-0.0.3 element-coder-0.0.6 ema-pytorch-0.0.10 fonttools-4.37.4 future-0.18.2 huggingface-hub-0.10.0 imagen-pytorch-1.11.4 kaleido-0.2.1 kornia-0.6.7 latexcodec-2.0.1 loguru-0.6.0 matbench-genmetrics-0.6.0 matminer-0.7.4 matplotlib-3.5.3 monty-2022.9.9 mp-time-split-0.2.0 pickle5-0.0.12 pint-0.18 py3Dmol-1.8.1 pyDeprecate-0.3.2 pybtex-0.24.0 pymatgen-2022.0.17 pyshtools-4.10 pystow-0.4.6 pytorch-lightning-1.7.7 pytorch-warmup-0.1.0 pyxtal-0.5.3 requests-2.28.1 ruamel.yaml-0.17.21 ruamel.yaml.clib-0.2.6 sentencepiece-0.1.97 six-1.16.0 spglib-2.0.1 sympy-1.10.1 tensorboard-2.10.1 tokenizers-0.12.1 torchmetrics-0.10.0 transformers-4.22.2 uncertainties-3.1.7 xtal2png-0.9.4

Data type cannot be displayed: application/vnd.colab-display-data+json

[4]:

try:
  from google.colab import output
  output.enable_custom_widget_manager()
except Exception as e:
  print(e)
  pass

[5]:

from os import path
from pathlib import Path
from uuid import uuid4
from tqdm import tqdm
import math

import numpy as np
import torch
from imagen_pytorch import ImagenConfig, ElucidatedImagenConfig, ImagenTrainer
from matbench_genmetrics.core import MPTSMetrics10, MPTSMetrics1000

from xtal2png.core import XtalConverter

Helper

[6]:

from PIL import Image
def image_grid(imgs, rows, cols):
    """https://stackoverflow.com/questions/37921295/python-pil-image-make-3x3-grid-from-sequence-images
    https://stackoverflow.com/a/65583584/13697228
    """
    assert len(imgs) == rows*cols

    w, h = imgs[0].size
    grid = Image.new('RGB', size=(cols*w, rows*h))
    grid_w, grid_h = grid.size

    for i, img in enumerate(imgs):
        grid.paste(img, box=(i%cols*w, i//cols*h))
    return grid

[7]:

import gc
from torch.cuda import empty_cache

def deallocate():
  if "trainer" in locals():
    try:
      del trainer
    except Exception as e:
      print(e)
      pass
  if "imagen" in locals():
    try:
      del imagen
    except Exception as e:
      print(e)
      pass
  if "training_images" in locals():
    try:
      del training_images
    except Exception as e:
      print(e)
      pass

  gc.collect()
  empty_cache()

Data and xtal2png images

[8]:

if dummy:
  mptm = MPTSMetrics10(dummy=dummy, verbose=True)
else:
  mptm = MPTSMetrics1000(dummy=dummy, verbose=True)

fold = 0
train_val_inputs = mptm.get_train_and_val_data(fold)

xc = XtalConverter(encode_cell_type=None, decode_cell_type=None)
arrays, id_data, id_mapper = xc.structures_to_arrays(train_val_inputs.tolist(), rgb_scaling=False)

Reading file /usr/local/lib/python3.7/dist-packages/mp_time_split/utils/mp_time_summary.json.gz: 0it [01:21, ?it/s]
Decoding objects from /usr/local/lib/python3.7/dist-packages/mp_time_split/utils/mp_time_summary.json.gz: 100%|##########| 40476/40476 [01:21<00:00, 497.36it/s]
  0%|          | 0/5400 [00:00<?, ?it/s]2022-10-08 06:30:07.574 | WARNING  | element_coder.data.coding_data:get_coding_dict:92 - This coding is not unique for certain elements. This will cause problems when decoding.
100%|██████████| 5400/5400 [01:23<00:00, 65.02it/s]

[9]:

gdrive_dir = '/content/drive/MyDrive/sparks-baird/xtal2png/'
print(path.isdir(gdrive_dir))
if dummy:
  gdrive_dir = path.join(gdrive_dir, "dummy")

True

Training

[14]:

if dummy:
  dummy_timesteps = 10 #@param {type:"integer"}
  dummy_epochs = 3 #@param {type:"integer"}
  dummy_checkpoint_rate = 2 #@param {type:"integer"}
  dummy_max_batch_size = 1 #@param {type:"integer"}
  dummy_dim = 16 #@param {type:"integer"}
  timesteps = dummy_timesteps
  epochs = dummy_epochs
  checkpoint_rate = dummy_checkpoint_rate
  max_batch_size = dummy_max_batch_size
  dim = dummy_dim
else:
  timesteps = 1000 #@param {type:"integer"}
  epochs = 1000 #@param {type:"integer"}
  checkpoint_rate = 100 #@param {type:"integer"}
  # max_batch_size suggested by lucidrains to be at least 32
  max_batch_size = 64 #@param {type:"integer"}
  dim = 128 #@param {type:"integer"}

[15]:

results_folder = path.join(
    gdrive_dir,
    "data",
    "interim",
    "elucidated-imagen-pytorch",
    f"fold={fold}",
    # str(uuid4())[0:4],
)

checkpoint_folder = path.join(results_folder, "checkpoint")
Path(checkpoint_folder).mkdir(exist_ok=True, parents=True)

[16]:

from glob import glob

files = glob(path.join(checkpoint_folder, "epoch=*.pt"))
checkpoint_epochs = [int(Path(file).stem.split("epoch=")[-1]) for file in files]
latest_chkpt_num = max(checkpoint_epochs) if len(checkpoint_epochs) > 0 else 0
latest_chkpt_num

[16]:

0

[17]:

from imagen_pytorch import load_imagen_from_checkpoint, ImagenTrainer

deallocate() # free up CUDA memory of previously defined `imagen` and `trainer` if existing

training_images = torch.from_numpy(arrays).float().cuda()

if latest_chkpt_num != 0:
  latest_checkpoint = path.join(checkpoint_folder, f"epoch={latest_chkpt_num}.pt")
  imagen = load_imagen_from_checkpoint(latest_checkpoint)
else:
  imagen = ElucidatedImagenConfig(
      condition_on_text=False,
      unets = [
          dict(dim = dim, dim_mults = (1, 2, 4, 8), memory_efficient=True),
          dict(dim = dim, dim_mults = (1, 2, 4, 8), memory_efficient=True),
      ],
      image_sizes = (32, 64),
      channels=1,
      cond_drop_prob = 0.1,
      num_sample_steps = 32,       # number of sample steps (default)
      sigma_min = 0.002,           # min noise level (default)
      sigma_max = 80,              # max noise level (default)
      sigma_data = 0.5,            # standard deviation of data distribution (default)
      rho = 7,                     # controls the sampling schedule (default)
      P_mean = -1.2,               # mean of log-normal distribution from which noise is drawn for training (default)
      P_std = 1.2,                 # standard deviation of log-normal distribution from which noise is drawn for training (default)
      S_churn = 80,                # parameters for stochastic sampling - depends on dataset, Table 5 in apper (default)
      S_tmin = 0.05,
      S_tmax = 50,
      S_noise = 1.003,
      # timesteps = timesteps, # Imagen
  ).create()

trainer = ImagenTrainer(imagen)

[18]:

%%capture nvidia_smi_output_instantiated
!nvidia-smi

[27]:

tmppath = path.join(results_folder, f"nvidia-smi-instantiated-checkpoint-{latest_chkpt_num}.txt")
with open(tmppath, "w") as f:
  f.write(nvidia_smi_output_instantiated.stdout)

[29]:

!nvidia-smi

Sat Oct  8 06:42:10 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  A100-SXM4-40GB      Off  | 00000000:00:04.0 Off |                    0 |
| N/A   29C    P0    48W / 400W |   4456MiB / 40536MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

[ ]:

# capture the tqdm output (which keeps printing newlines) to avoid Colab crash
from IPython.utils import io
import pandas as pd

# train each unet in concert, or separately (recommended) to completion
l_arr = [[],[]]

with tqdm(initial=latest_chkpt_num, total=epochs) as pbar:
  for epoch in range(latest_chkpt_num, epochs):
    for u in (1, 2):
      with io.capture_output() as captured:
        loss = trainer(training_images, unet_number=u, max_batch_size=max_batch_size)
        trainer.update(unet_number=u)
        trainer.save('./checkpoint.pt') # for switching back-and-forth b/t unets
        imagen = load_imagen_from_checkpoint('./checkpoint.pt')
        trainer = ImagenTrainer(imagen)
      l_arr[u-1].append(loss)
    pbar.update()
    if (epoch % checkpoint_rate == 0) or (epoch == epochs - 1):
      with io.capture_output() as captured:
        images = trainer.sample(batch_size=16, return_pil_images=True)

      grid = image_grid(images, 4, 4)
      display(grid)
      grid.save(path.join(checkpoint_folder, f"sample_4x4,epoch={epoch}.png"))
      # if not math.isnan(loss) and (epoch != epochs - 1):
      trainer.save(path.join(checkpoint_folder, f"epoch={epoch}.pt"))
      pd.DataFrame(dict(u1=l_arr[0], u2=l_arr[1])).to_csv(path.join(checkpoint_folder, f"loss_epoch={epoch}.csv"), index=False)
      print(f"loss is {loss} for epoch={epoch}")

  0%|          | 1/1000 [02:03<34:10:41, 123.16s/it]
../_images/notebooks_3.1-imagen-pytorch_21_1.png 
this checkpoint is commandable from the CLI - "imagen --model /content/drive/MyDrive/sparks-baird/xtal2png/data/interim/elucidated-imagen-pytorch/fold=0/checkpoint/epoch=0.pt "<prompt>""
checkpoint saved to /content/drive/MyDrive/sparks-baird/xtal2png/data/interim/elucidated-imagen-pytorch/fold=0/checkpoint/epoch=0.pt
loss is 1.953648234717548 for epoch=0

 10%|█         | 101/1000 [3:15:41<28:56:07, 115.87s/it]
../_images/notebooks_3.1-imagen-pytorch_21_4.png 
this checkpoint is commandable from the CLI - "imagen --model /content/drive/MyDrive/sparks-baird/xtal2png/data/interim/elucidated-imagen-pytorch/fold=0/checkpoint/epoch=100.pt "<prompt>""
checkpoint saved to /content/drive/MyDrive/sparks-baird/xtal2png/data/interim/elucidated-imagen-pytorch/fold=0/checkpoint/epoch=100.pt
loss is 0.8361775362864137 for epoch=100

 20%|██        | 201/1000 [6:30:17<25:56:53, 116.91s/it]
../_images/notebooks_3.1-imagen-pytorch_21_7.png 
this checkpoint is commandable from the CLI - "imagen --model /content/drive/MyDrive/sparks-baird/xtal2png/data/interim/elucidated-imagen-pytorch/fold=0/checkpoint/epoch=200.pt "<prompt>""
checkpoint saved to /content/drive/MyDrive/sparks-baird/xtal2png/data/interim/elucidated-imagen-pytorch/fold=0/checkpoint/epoch=200.pt
loss is 0.6849522094707936 for epoch=200

 30%|███       | 301/1000 [9:46:09<22:52:13, 117.79s/it]
../_images/notebooks_3.1-imagen-pytorch_21_10.png 
this checkpoint is commandable from the CLI - "imagen --model /content/drive/MyDrive/sparks-baird/xtal2png/data/interim/elucidated-imagen-pytorch/fold=0/checkpoint/epoch=300.pt "<prompt>""
checkpoint saved to /content/drive/MyDrive/sparks-baird/xtal2png/data/interim/elucidated-imagen-pytorch/fold=0/checkpoint/epoch=300.pt
loss is 0.51092101004906 for epoch=300

 35%|███▍      | 348/1000 [11:18:19<21:06:12, 116.52s/it]

Sampling

[ ]:

with io.capture_output() as captured:
  images = trainer.sample(batch_size=16, return_pil_images=True)
  trainer.save(path.join(results_folder, f"epoch={epoch}.pt"))

[ ]:

grid = image_grid(images, 4, 4)
display(grid)
grid.save(path.join(results_folder, f"sample_4x4,epoch={epoch}.png"))

deallocate CUDA memory

https://discuss.pytorch.org/t/how-can-we-release-gpu-memory-cache/14530/28

[ ]:

%%capture nvidia_smi_output
!nvidia-smi

[ ]:

tmppath = path.join(results_folder, f"nvidia-smi-completed-checkpoint-{latest_chkpt_num}.txt")
with open(tmppath, "w") as f:
  f.write(nvidia_smi_output.stdout)

[ ]:

deallocate()
!nvidia-smi

[ ]:

num_gen = mptm.num_gen

if dummy:
  n = 1
else:
  n = 100

def chunks(lst, n):
    """Yield successive n-sized chunks from lst."""
    for i in range(0, len(lst), n):
        yield lst[i:i + n]

gen_images = []

with io.capture_output() as captured:
  for chunk in tqdm(chunks(range(num_gen), n)):
    gen_images.append(trainer.sample(batch_size=len(chunk), return_pil_images=True))

[ ]:

gen_structures = xc.png2xtal(gen_images)
mptm.evaluate_and_record(fold, gen_structures)
print(mptm.recorded_metrics)

[ ]:

import json
import pickle
with open(path.join(results_folder, f"metrics_epoch={epoch}.json"),"w") as f:
    json.dump(mptm.recorded_metrics,f)

with open(path.join(results_folder, f"mptm_epoch={epoch}.pkl"), "wb") as f:
  pickle.dump(mptm, f)

[ ]:

from IPython.display import display, Javascript
display(Javascript('IPython.notebook.save_checkpoint();')) # save and pin revision, check

[ ]:

import time

timestr = time.strftime("%Y%m%d-%H%M%S")
notebook_savepath = path.join(results_folder, f"notebook-{timestr}.ipynb")
print(notebook_savepath)

[ ]:

# doesn't save output
%notebook -e notebook.ipynb
!mv notebook.ipynb $notebook_savepath

[ ]:

from google.colab import runtime
runtime.unassign()

Code Graveyard

[ ]:

# %pip install  ase nglview ipywidgets python-vtk

# import crystal_toolkit.components as ctc
# from ase.io import write
# from pymatgen.io.ase import AseAtomsAdaptor
# from pymatgen.vis.structure_vtk import StructureVis
# from PIL import Image

# def get_pil_image(atoms):
#   fname = 'image.png'
#   write(fname, gen_atoms[0])
#   return Image.open(fname)

# aaa = AseAtomsAdaptor()
# sv = StructureVis()

# def get_pil_from_struct(structure):
#   fname = 'image.png'
#   sv.set_structure(structure)
#   sv.write_image(fname)
#   return Image.open(fname)

        # arrs = np.expand_dims(np.stack([np.asarray(im) for im in images]), 1)
        # gen_structures = xc.png2xtal(images)
        # gen_atoms = [aaa.get_atoms(s) for s in gen_structures]
        # gen_views = [get_pil_image(ga) for ga in gen_atoms]
        # gen_views = [get_pil_from_struct(s) for s in gen_structures]

      # views_grid = image_grid(gen_views, 4, 4)
      # display(views_grid)
      # views_grid.save(path.join(checkpoint_folder, f"xtal_sample_4x4,epoch={epoch}.png"))

# gen_structs = xc.png2xtal(images)
# gen_atoms = [aaa.get_atoms(s) for s in gen_structures]
# gen_views = [get_pil_image(ga) for ga in gen_atoms]

[ ]:

# gen_arrays = torch.concat(gen_subarrays, axis=0)
# gen_structures = xc.arrays_to_structures(gen_arrays.cpu().numpy(), id_data=id_data, id_mapper=id_mapper, rgb_scaling=False)