Update requirements to support CUDA>=11.0

.gitignore

@@ -1,2 +1,4 @@
 stylegan/
-!stylegan/.gitkeep
+!stylegan/.gitkeep
+__pycache__
+.stylegan2-cache

Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04
+FROM nvidia/cuda:11.2.0-cudnn8-devel-ubuntu18.04
 
 RUN apt-get update && apt-get install python3.7 python3-dev libpython3.7-dev python3-pip cmake g++ gnupg -y
 RUN apt-get install libgl1-mesa-glx libegl1-mesa libxrandr2 libxrandr2 libxss1 libxcursor1 libxcomposite1 libasound2 libxi6 libxtst6 libsm6 libxext6 libxrender-dev x11-apps libqt5x11extras5 -y
@@ -13,7 +13,9 @@ ENV LANG=en_US.UTF-8 \
 RUN python3.7 -m pip install --upgrade pip setuptools
 
 WORKDIR /usr/app
+RUN python3.7 -m pip install torch==1.8.2 torchvision==0.9.2 --extra-index-url https://download.pytorch.org/whl/lts/1.8/cu111
 COPY ./requirements.txt /usr/app
-RUN python3.7 -m pip install -r requirements.txt
+RUN --mount=type=cache,target=/root/.cache \
+        python3.7 -m pip install -r requirements.txt
 
 COPY . /usr/app

README.md

@@ -31,13 +31,17 @@ git clone https://github.com/RameenAbdal/StyleFlow.git
 cd StyleFlow/
 ```
 
+StyleGAN2 relies on custom TensorFlow ops that are compiled on the fly using [NVCC](https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html). To correctly setup the StyleGAN2 generator follow the **Requirements** in [this repo](https://github.com/NVlabs/stylegan2).
+```bash
+# This packages are enough to fulfill all requirements (at least for WSL Ubuntu>=20.04 installation)
+sudo apt install cmake g++ gnupg nvidia-cuda-toolkit
+```
+
 This code requires PyTorch, TensorFlow, Torchdiffeq, Python 3+ and Pyqt5. Please install dependencies by
 ```bash
 conda env create -f environment.yml
 ```
 
-StyleGAN2 relies on custom TensorFlow ops that are compiled on the fly using [NVCC](https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html). To correctly setup the StyleGAN2 generator follow the **Requirements** in [this repo](https://github.com/NVlabs/stylegan2).
-
 ## Installation (Docker)
 
 Clone this repo.
@@ -47,7 +51,7 @@ git clone https://github.com/RameenAbdal/StyleFlow.git
 cd StyleFlow/
 ```
 
-You must have CUDA (>=10.0 && <11.0) and [nvidia-docker2](https://github.com/NVIDIA/nvidia-docker) installed first !
+You must have CUDA >=10.0 and [nvidia-docker2](https://github.com/NVIDIA/nvidia-docker) installed first !
 
 Then, run :
 
@@ -84,7 +88,7 @@ xhost -local:docker
 
 2. Run the Attribute Transfer UI
 	```bash
-   python main_attribute.py 
+   python main_attribute.py
     ```
 
 ## Web UI (Beta)
@@ -99,7 +103,7 @@ Create `./data_numpy/` in the main folder and extract the above data or create y
 
 Train your model:
 ```bash
-   python train_flow.py 
+   python train_flow.py
 ```
 ## Projection
 
@@ -136,8 +140,8 @@ keywords = {image editing, Generative adversarial networks}
 ```
 @INPROCEEDINGS{9008515,
   author={Abdal, Rameen and Qin, Yipeng and Wonka, Peter},
-  booktitle={2019 IEEE/CVF International Conference on Computer Vision (ICCV)}, 
-  title={Image2StyleGAN: How to Embed Images Into the StyleGAN Latent Space?}, 
+  booktitle={2019 IEEE/CVF International Conference on Computer Vision (ICCV)},
+  title={Image2StyleGAN: How to Embed Images Into the StyleGAN Latent Space?},
   year={2019},
   volume={},
   number={},

dnnlib/tflib/autosummary.py

@@ -23,7 +23,11 @@
 
 from collections import OrderedDict
 import numpy as np
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+    tf.disable_v2_behavior()
+except:
+    import tensorflow as tf
 from tensorboard import summary as summary_lib
 from tensorboard.plugins.custom_scalar import layout_pb2
 

dnnlib/tflib/custom_ops.py

@@ -13,7 +13,11 @@
 import hashlib
 import tempfile
 import shutil
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+    tf.disable_v2_behavior()
+except:
+    import tensorflow as tf
 from tensorflow.python.client import device_lib # pylint: disable=no-name-in-module
 
 #----------------------------------------------------------------------------
@@ -126,7 +130,7 @@ def get_plugin(cuda_file):
             compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.lib')
         elif os.name == 'posix':
             compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.so')
-            compile_opts += ' --compiler-options \'-fPIC -D_GLIBCXX_USE_CXX11_ABI=0\''
+            compile_opts += ' --compiler-options \'-fPIC -D_GLIBCXX_USE_CXX11_ABI=1\''
         else:
             assert False # not Windows or Linux, w00t?
         compile_opts += ' --gpu-architecture=%s' % _get_cuda_gpu_arch_string()

dnnlib/tflib/network.py

@@ -12,7 +12,11 @@
 import uuid
 import sys
 import numpy as np
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+    tf.disable_v2_behavior()
+except:
+    import tensorflow as tf
 
 from collections import OrderedDict
 from typing import Any, List, Tuple, Union

dnnlib/tflib/ops/fused_bias_act.py

@@ -8,7 +8,14 @@
 
 import os
 import numpy as np
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+    tf.disable_v2_behavior()
+    force_ref_impl = True
+except:
+    import tensorflow as tf
+    force_ref_impl = False
+
 from .. import custom_ops
 from ...util import EasyDict
 
@@ -65,7 +72,7 @@ def fused_bias_act(x, b=None, axis=1, act='linear', alpha=None, gain=None, impl=
         'ref':  _fused_bias_act_ref,
         'cuda': _fused_bias_act_cuda,
     }
-    return impl_dict[impl](x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain)
+    return impl_dict['ref' if force_ref_impl else impl](x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain)
 
 #----------------------------------------------------------------------------
 

dnnlib/tflib/ops/upfirdn_2d.py

@@ -8,7 +8,14 @@
 
 import os
 import numpy as np
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+    tf.disable_v2_behavior()
+    force_ref_impl = True
+except:
+    import tensorflow as tf
+    force_ref_impl = False
+
 from .. import custom_ops
 
 def _get_plugin():
@@ -59,7 +66,7 @@ def upfirdn_2d(x, k, upx=1, upy=1, downx=1, downy=1, padx0=0, padx1=0, pady0=0,
         'ref':  _upfirdn_2d_ref,
         'cuda': _upfirdn_2d_cuda,
     }
-    return impl_dict[impl](x=x, k=k, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1)
+    return impl_dict['ref' if force_ref_impl else impl](x=x, k=k, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1)
 
 #----------------------------------------------------------------------------
 

dnnlib/tflib/optimizer.py

@@ -7,7 +7,11 @@
 """Helper wrapper for a Tensorflow optimizer."""
 
 import numpy as np
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+    tf.disable_v2_behavior()
+except:
+    import tensorflow as tf
 
 from collections import OrderedDict
 from typing import List, Union

dnnlib/tflib/tfutil.py

@@ -8,13 +8,15 @@
 
 import os
 import numpy as np
-import tensorflow as tf
+try:
+    import tensorflow.compat.v1 as tf
+    tf.disable_v2_behavior()
+except:
+    import tensorflow as tf
 
 # Silence deprecation warnings from TensorFlow 1.13 onwards
 import logging
 logging.getLogger('tensorflow').setLevel(logging.ERROR)
-import tensorflow.contrib   # requires TensorFlow 1.x!
-tf.contrib = tensorflow.contrib
 
 from typing import Any, Iterable, List, Union
 
@@ -36,7 +38,7 @@ def is_tf_expression(x: Any) -> bool:
     return isinstance(x, (tf.Tensor, tf.Variable, tf.Operation))
 
 
-def shape_to_list(shape: Iterable[tf.Dimension]) -> List[Union[int, None]]:
+def shape_to_list(shape) -> List[Union[int, None]]:
     """Convert a Tensorflow shape to a list of ints. Retained for backwards compatibility -- use TensorShape.as_list() in new code."""
     return [dim.value for dim in shape]
 

docker-compose.yml

@@ -5,9 +5,14 @@ services:
   interface:
     build: .
     command: python3.7 /usr/app/main.py
-    volumes: 
+    volumes:
       - /tmp/.X11-unix:/tmp/.X11-unix
     environment:
       DISPLAY: $DISPLAY
       QT_X11_NO_MITSHM: 1
       QT_DEBUG_PLUGINS: 1
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]

environment.yml

@@ -2,6 +2,7 @@ name: StyleFlow
 channels:
   - anaconda
   - defaults
+  - nvidia
   - conda-forge
 dependencies:
   - _libgcc_mutex=0.1=conda_forge
@@ -10,6 +11,8 @@ dependencies:
   - blas=1.0=mkl
   - bzip2=1.0.8=h7f98852_4
   - ca-certificates=2020.12.8=h06a4308_0
+  - cudatoolkit=11.1.74=h6bb024c_0
+  - cudnn=8.1.0.77=h90431f1_0
   - certifi=2020.12.5=py37h06a4308_0
   - ffmpeg=4.3.1=h3215721_1
   - freetype=2.9.1=h8a8886c_1
@@ -54,7 +57,10 @@ dependencies:
   - zlib=1.2.11=h7b6447c_3
   - zstd=1.3.7=h0b5b093_0
   - pip:
-    - absl-py==0.7.1
+    - --extra-index-url https://download.pytorch.org/whl/lts/1.8/cu111
+    - torch==1.8.2
+    - torchvision==0.9.2
+    - absl-py==0.10
     - appdirs==1.4.4
     - astor==0.8.0
     - astunparse==1.6.3
@@ -64,6 +70,7 @@ dependencies:
     - cachetools==4.1.0
     - cffi==1.12.3
     - chardet==3.0.4
+    - clang==5.0
     - cloudpickle==1.2.1
     - cycler==0.10.0
     - cytoolz==0.9.0.1
@@ -76,12 +83,13 @@ dependencies:
     - dominate==2.3.5
     - easydict==1.9
     - entrypoints==0.3
-    - gast==0.2.2
+    - flatbuffers==1.12.0
+    - gast==0.4.0
     - google-auth==1.14.3
     - google-auth-oauthlib==0.4.1
     - google-pasta==0.2.0
-    - grpcio==1.22.0
-    - h5py==2.10.0
+    - grpcio==1.37.0
+    - h5py==3.1.0
     - helpdev==0.6.10
     - idna==2.8
     - imageio==2.5.0
@@ -98,7 +106,7 @@ dependencies:
     - jupyter-client==5.2.4
     - jupyter-console==6.0.0
     - jupyter-core==4.5.0
-    - keras==2.2.4
+    - keras==2.6
     - keras-applications==1.0.8
     - keras-preprocessing==1.1.0
     - kiwisolver==1.1.0
@@ -113,21 +121,21 @@ dependencies:
     - notebook==5.7.8
     - oauthlib==3.1.0
     - opencv-python==4.1.0.25
-    - opt-einsum==3.2.1
+    - opt-einsum==3.3.0
     - pandocfilters==1.4.2
     - parso==0.5.0
     - pexpect==4.7.0
     - pickleshare==0.7.5
     - pillow==6.0.0
     - prometheus-client==0.7.1
     - prompt-toolkit==2.0.9
-    - protobuf==3.8.0
+    - protobuf==3.9.2
     - psutil==5.6.3
     - ptyprocess==0.6.0
     - pyasn1==0.4.8
     - pyasn1-modules==0.2.8
     - pycparser==2.19
-    - pycuda==2019.1.2
+    - pycuda==2022.1
     - pygments==2.4.2
     - pyparsing==2.4.0
     - pyqt5==5.13.0
@@ -151,26 +159,23 @@ dependencies:
     - scipy==1.2.1
     - send2trash==1.5.0
     - shiboken2==5.13.0
-    - six==1.12.0
-    - tensorboard==1.15.0
-    - tensorboard-plugin-wit==1.6.0.post3
-    - tensorflow-estimator==1.15.1
-    - tensorflow-gpu==1.15.0
+    - six==1.15.0
+    - tensorboard==2.11.0
+    - tensorboard-plugin-wit==1.8.1
+    - tensorflow==2.6.0
+    - tensorflow-estimator==2.11.0
     - termcolor==1.1.0
     - terminado==0.8.2
     - testpath==0.4.2
     - toolz==0.9.0
-    - torch==1.1.0
     - torchdiffeq==0.0.1
-    - torchvision==0.3.0
     - tornado==6.0.3
     - tqdm==4.32.1
     - traitlets==4.3.2
     - urllib3==1.25.3
     - wcwidth==0.1.7
     - webencodings==0.5.1
-    - werkzeug==0.15.4
+    - werkzeug==1.0.1
     - widgetsnbextension==3.4.2
-    - wrapt==1.11.2
+    - wrapt==1.12.1
     - zipp==0.5.2
-