Merge branch 'main' into main

.gitignore

@@ -99,7 +99,7 @@ ENV/
 [Ll]ib64
 [Ll]ocal
 [Ss]cripts
-!scripts\download_model.bat
+!scripts/
 pyvenv.cfg
 .venv
 pip-selfcheck.json

Dockerfile

@@ -0,0 +1,28 @@
+FROM nvcr.io/nvidia/pytorch:23.05-py3
+
+ENV PYTHONDONTWRITEBYTECODE 1
+ENV PYTHONUNBUFFERED 1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        make \
+        pkgconf \
+        xz-utils \
+        xorg-dev \
+        libgl1-mesa-dev \
+        libglu1-mesa-dev \
+        libxrandr-dev \
+        libxinerama-dev \
+        libxcursor-dev \
+        libxi-dev \
+        libxxf86vm-dev \
+        && rm -rf /var/lib/apt/lists/*
+
+RUN pip install --no-cache-dir --upgrade pip
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+WORKDIR /workspace
+
+RUN (printf '#!/bin/bash\nexec \"$@\"\n' >> /entry.sh) && chmod a+x /entry.sh
+ENTRYPOINT ["/entry.sh"]

README.md

@@ -47,7 +47,7 @@ If you have CUDA graphic card, please follow the requirements of [NVlabs/stylega
 The usual installation steps involve the following commands, they should set up the correct CUDA version and all the python packages
 
 ```
-conda env create python=3.7 -f environment.yml
+conda env create -f environment.yml
 conda activate stylegan3
 ```
 
@@ -69,11 +69,24 @@ conda activate stylegan3
 export PYTORCH_ENABLE_MPS_FALLBACK=1
 ```
 
+## Run Gradio visualizer in Docker 
+
+Provided docker image is based on NGC PyTorch repository. To quickly try out visualizer in Docker, run the following:  
+
+```sh
+docker build . -t draggan:latest  
+docker run -p 7860: 7860 -v "$PWD":/workspace/src -it draggan:latest bash  
+cd src && python visualizer_drag_gradio.py --listen
+```
+Now you can open a shared link from Gradio (printed in the terminal console).   
+Beware the Docker image takes about 25GB of disk space!
+
 ## Download pre-trained StyleGAN2 weights
 
 To download pre-trained weights, simply run:
-```sh
+
-sh scripts/download_model.sh
+```
+python scripts/download_model.py
 ```
 If you want to try StyleGAN-Human and the Landscapes HQ (LHQ) dataset, please download weights from these links: [StyleGAN-Human](https://drive.google.com/file/d/1dlFEHbu-WzQWJl7nBBZYcTyo000H9hVm/view?usp=sharing), [LHQ](https://drive.google.com/file/d/16twEf0T9QINAEoMsWefoWiyhcTd-aiWc/view?usp=sharing), and put them under `./checkpoints`.
 
@@ -85,10 +98,14 @@ To start the DragGAN GUI, simply run:
 ```sh
 sh scripts/gui.sh
 ```
+If you are using windows, you can run:
+```
+.\scripts\gui.bat
+```
 
 This GUI supports editing GAN-generated images. To edit a real image, you need to first perform GAN inversion using tools like [PTI](https://github.com/danielroich/PTI). Then load the new latent code and model weights to the GUI.
 
-You can run DragGAN Gradio demo as well:
+You can run DragGAN Gradio demo as well, this is universal for both windows and linux:
 ```sh
 python visualizer_drag_gradio.py
 ```
@@ -105,6 +122,7 @@ python visualizer_drag_gradio.py --port=8888
 
 This code is developed based on [StyleGAN3](https://github.com/NVlabs/stylegan3). Part of the code is borrowed from [StyleGAN-Human](https://github.com/stylegan-human/StyleGAN-Human).
 
+(cheers to the community as well)
 ## License
 
 The code related to the DragGAN algorithm is licensed under [CC-BY-NC](https://creativecommons.org/licenses/by-nc/4.0/).

requirements.txt

@@ -1,9 +1,13 @@
-torch
+torch>=2.0.0
-torchvision
+scipy==1.11.0
-Ninja
+Ninja==1.10.2
-gradio
+gradio>=3.35.2
+imageio-ffmpeg>=0.4.3
 huggingface_hub
 hf_transfer
 pyopengl
 imgui
-glfw
+glfw==2.6.1
+pillow>=9.4.0
+torchvision>=0.15.2
+imageio>=2.9.0

scripts/download_model.bat

@@ -1,23 +0,0 @@
-@echo off
-mkdir checkpoints
-cd checkpoints
-
-powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://storage.googleapis.com/self-distilled-stylegan/lions_512_pytorch.pkl', 'lions_512_pytorch.pkl')"
-ren lions_512_pytorch.pkl stylegan2_lions_512_pytorch.pkl
-
-powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://storage.googleapis.com/self-distilled-stylegan/dogs_1024_pytorch.pkl', 'dogs_1024_pytorch.pkl')"
-ren dogs_1024_pytorch.pkl stylegan2_dogs_1024_pytorch.pkl
-
-powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://storage.googleapis.com/self-distilled-stylegan/horses_256_pytorch.pkl', 'horses_256_pytorch.pkl')"
-ren horses_256_pytorch.pkl stylegan2_horses_256_pytorch.pkl
-
-powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://storage.googleapis.com/self-distilled-stylegan/elephants_512_pytorch.pkl', 'elephants_512_pytorch.pkl')"
-ren elephants_512_pytorch.pkl stylegan2_elephants_512_pytorch.pkl
-
-powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-512x512.pkl', 'stylegan2-ffhq-512x512.pkl')"
-powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-afhqcat-512x512.pkl', 'stylegan2-afhqcat-512x512.pkl')"
-powershell -Command "(New-Object System.Net.WebClient).DownloadFile('http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-car-config-f.pkl', 'stylegan2-car-config-f.pkl')"
-powershell -Command "(New-Object System.Net.WebClient).DownloadFile('http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-cat-config-f.pkl', 'stylegan2-cat-config-f.pkl')"
-
-echo "Done"
-pause

scripts/download_model.py

@@ -0,0 +1,78 @@
+import os
+import sys
+import json
+import requests
+from tqdm import tqdm
+
+def download_file(url: str, filename: str, download_dir: str):
+    """Download a file if it does not already exist."""
+
+    try:
+        filepath = os.path.join(download_dir, filename)
+        content_length = int(requests.head(url).headers.get("content-length", 0))
+
+        # If file already exists and size matches, skip download
+        if os.path.isfile(filepath) and os.path.getsize(filepath) == content_length:
+            print(f"{filepath} already exists. Skipping download.")
+            return
+        if os.path.isfile(filepath) and os.path.getsize(filepath) != content_length:
+            print(f"{filepath} already exists but size does not match. Redownloading.")
+        else:
+            print(f"Downloading {filename} from {url}")
+
+        # Start download, stream=True allows for progress tracking
+        response = requests.get(url, stream=True)
+
+        # Check if request was successful
+        response.raise_for_status()
+
+        # Create progress bar
+        total_size = int(response.headers.get('content-length', 0))
+        progress_bar = tqdm(
+            total=total_size, 
+            unit='iB', 
+            unit_scale=True, 
+            ncols=70, 
+            file=sys.stdout
+        )
+
+        # Write response content to file
+        with open(filepath, 'wb') as f:
+            for data in response.iter_content(chunk_size=1024):
+                f.write(data)
+                progress_bar.update(len(data))  # Update progress bar
+
+        # Close progress bar
+        progress_bar.close()
+
+        # Error handling for incomplete downloads
+        if total_size != 0 and progress_bar.n != total_size:
+            print("ERROR, something went wrong while downloading")
+            raise Exception()
+
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+def main():
+    """Main function to download files from URLs in a config file."""
+    
+    # Get JSON config file path
+    script_dir = os.path.dirname(os.path.realpath(__file__))
+    config_file_path = os.path.join(script_dir, "download_models.json")
+
+    # Set download directory
+    download_dir = "checkpoints"
+    os.makedirs(download_dir, exist_ok=True)
+
+    # Load URL and filenames from JSON
+    with open(config_file_path, "r") as f:
+        config = json.load(f)
+
+    # Download each file specified in config
+    for url, filename in config.items():
+        download_file(url, filename, download_dir)
+
+
+if __name__ == "__main__":
+    main()

scripts/download_model.sh

@@ -1,19 +0,0 @@
-mkdir checkpoints
-cd checkpoints
-
-wget https://storage.googleapis.com/self-distilled-stylegan/lions_512_pytorch.pkl
-mv lions_512_pytorch.pkl stylegan2_lions_512_pytorch.pkl
-
-wget https://storage.googleapis.com/self-distilled-stylegan/dogs_1024_pytorch.pkl
-mv dogs_1024_pytorch.pkl stylegan2_dogs_1024_pytorch.pkl
-
-wget https://storage.googleapis.com/self-distilled-stylegan/horses_256_pytorch.pkl
-mv horses_256_pytorch.pkl stylegan2_horses_256_pytorch.pkl
-
-wget https://storage.googleapis.com/self-distilled-stylegan/elephants_512_pytorch.pkl
-mv elephants_512_pytorch.pkl stylegan2_elephants_512_pytorch.pkl
-
-wget https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-512x512.pkl
-wget https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-afhqcat-512x512.pkl
-wget http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-car-config-f.pkl
-wget http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-cat-config-f.pkl

scripts/download_models.json

@@ -0,0 +1,10 @@
+{
+  "https://storage.googleapis.com/self-distilled-stylegan/lions_512_pytorch.pkl": "stylegan2_lions_512_pytorch.pkl",
+  "https://storage.googleapis.com/self-distilled-stylegan/dogs_1024_pytorch.pkl": "stylegan2_dogs_1024_pytorch.pkl",
+  "https://storage.googleapis.com/self-distilled-stylegan/horses_256_pytorch.pkl": "stylegan2_horses_256_pytorch.pkl",
+  "https://storage.googleapis.com/self-distilled-stylegan/elephants_512_pytorch.pkl": "stylegan2_elephants_512_pytorch.pkl",
+  "https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-512x512.pkl": "stylegan2-ffhq-512x512.pkl",
+  "https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-afhqcat-512x512.pkl": "stylegan2-afhqcat-512x512.pkl",
+  "http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-car-config-f.pkl": "stylegan2-car-config-f.pkl",
+  "http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-cat-config-f.pkl": "stylegan2-cat-config-f.pkl"
+}

visualizer_drag_gradio.py

@@ -22,6 +22,7 @@ parser.add_argument('--host', type=str,
                     help="launch gradio with given server name", default=None)
 parser.add_argument('--port', type=int,
                     help="launch gradio with given server port", default=None)
+
 args = parser.parse_args()
 
 cache_dir = args.cache_dir