diff --git a/.gitignore b/.gitignore index 24b5204..8b626aa 100644 --- a/.gitignore +++ b/.gitignore @@ -99,7 +99,7 @@ ENV/ [Ll]ib64 [Ll]ocal [Ss]cripts -!scripts\download_model.bat +!scripts/ pyvenv.cfg .venv pip-selfcheck.json diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..2c6ad40 --- /dev/null +++ b/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"] diff --git a/README.md b/README.md index f8cd72d..0008993 100644 --- a/README.md +++ b/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/). diff --git a/requirements.txt b/requirements.txt index e71044d..556e127 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,13 @@ -torch -torchvision -Ninja -gradio +torch>=2.0.0 +scipy==1.11.0 +Ninja==1.10.2 +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 diff --git a/scripts/download_model.bat b/scripts/download_model.bat deleted file mode 100644 index 67101d4..0000000 --- a/scripts/download_model.bat +++ /dev/null @@ -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 diff --git a/scripts/download_model.py b/scripts/download_model.py new file mode 100644 index 0000000..b951b05 --- /dev/null +++ b/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() diff --git a/scripts/download_model.sh b/scripts/download_model.sh deleted file mode 100644 index 4d2dccc..0000000 --- a/scripts/download_model.sh +++ /dev/null @@ -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 diff --git a/scripts/download_models.json b/scripts/download_models.json new file mode 100644 index 0000000..637d553 --- /dev/null +++ b/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" +} diff --git a/visualizer_drag_gradio.py b/visualizer_drag_gradio.py index 418fe81..3fd2774 100644 --- a/visualizer_drag_gradio.py +++ b/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