aiden
/
DragGAN
mirror of https://github.com/XingangPan/DragGAN
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into main

Browse Source
pull/112/head
venshine 2 years ago committed by GitHub
parent eace07043e 1a2d37d03d
commit 80f1c25738
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 149 additions and 52 deletions
Show Stats Download Patch File Download Diff File

2
.gitignore vendored
Unescape Escape View File

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

28
Dockerfile
Unescape Escape View File

@ -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"]

26
README.md
Unescape Escape View File

@ -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/).

14
requirements.txt
Unescape Escape View File

@ -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

23
scripts/download_model.bat
Unescape Escape View File

@ -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

78
scripts/download_model.py
Unescape Escape View File

@ -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()

19
scripts/download_model.sh
Unescape Escape View File

@ -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

10
scripts/download_models.json
Unescape Escape View File

@ -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"
}

1
visualizer_drag_gradio.py
Unescape Escape View File

@ -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

Write Preview
Loading…
Cancel
Save