public-data / TADNE

TADNE (This Anime Does Not Exist) model

The original TADNE site is https://thisanimedoesnotexist.ai/ .

Original TensorFlow model

The original TADNE model is provided in this site under CC-0 license. ( Google Drive )

Model Conversion

The model in the models directory is converted with the following repo: https://github.com/rosinality/stylegan2-pytorch

Apply patches

--- a/model.py
+++ b/model.py
@@ -395,6 +395,7 @@ class Generator(nn.Module):
         style_dim,
         n_mlp,
         channel_multiplier=2,
+        additional_multiplier=2,
         blur_kernel=[1, 3, 3, 1],
         lr_mlp=0.01,
     ):
@@ -426,6 +427,9 @@ class Generator(nn.Module):
             512: 32 * channel_multiplier,
             1024: 16 * channel_multiplier,
         }
+        if additional_multiplier > 1:
+            for k in list(self.channels.keys()):
+                self.channels[k] *= additional_multiplier

         self.input = ConstantInput(self.channels[4])
         self.conv1 = StyledConv(
@@ -518,7 +522,7 @@ class Generator(nn.Module):
                     getattr(self.noises, f"noise_{i}") for i in range(self.num_layers)
                 ]

-        if truncation < 1:
+        if truncation_latent is not None:
             style_t = []

             for style in styles:

--- a/convert_weight.py
+++ b/convert_weight.py
@@ -221,6 +221,7 @@ if __name__ == "__main__":
         default=2,
         help="channel multiplier factor. config-f = 2, else = 1",
     )
+    parser.add_argument("--additional_multiplier", type=int, default=2)
     parser.add_argument("path", metavar="PATH", help="path to the tensorflow weights")

     args = parser.parse_args()
@@ -243,7 +244,8 @@ if __name__ == "__main__":
         if layer[0].startswith('Dense'):
             n_mlp += 1

-    g = Generator(size, 512, n_mlp, channel_multiplier=args.channel_multiplier)
+    style_dim = 512 * args.additional_multiplier
+    g = Generator(size, style_dim, n_mlp, channel_multiplier=args.channel_multiplier, additional_multiplier=args.additional_multiplier)
     state_dict = g.state_dict()
     state_dict = fill_statedict(state_dict, g_ema.vars, size, n_mlp)

@@ -254,7 +256,7 @@ if __name__ == "__main__":
     ckpt = {"g_ema": state_dict, "latent_avg": latent_avg}

     if args.gen:
-        g_train = Generator(size, 512, n_mlp, channel_multiplier=args.channel_multiplier)
+        g_train = Generator(size, style_dim, n_mlp, channel_multiplier=args.channel_multiplier, additional_multiplier=args.additional_multiplier)
         g_train_state = g_train.state_dict()
         g_train_state = fill_statedict(g_train_state, generator.vars, size, n_mlp)
         ckpt["g"] = g_train_state
@@ -271,9 +273,12 @@ if __name__ == "__main__":
     batch_size = {256: 16, 512: 9, 1024: 4}
     n_sample = batch_size.get(size, 25)

+    if args.additional_multiplier > 1:
+        n_sample = 2
+
     g = g.to(device)

-    z = np.random.RandomState(0).randn(n_sample, 512).astype("float32")
+    z = np.random.RandomState(0).randn(n_sample, style_dim).astype("float32")

     with torch.no_grad():
         img_pt, _ = g(

Build Docker image

FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04

ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update -y && \
    apt-get install -y --no-install-recommends \
    git \
    ninja-build \
    # pyenv dependencies \
    make \
    build-essential \
    libssl-dev \
    zlib1g-dev \
    libbz2-dev \
    libreadline-dev \
    libsqlite3-dev \
    wget \
    curl \
    llvm \
    libncursesw5-dev \
    xz-utils \
    tk-dev \
    libxml2-dev \
    libxmlsec1-dev \
    libffi-dev \
    liblzma-dev && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

ARG PYTHON_VERSION=3.7.12
ENV PYENV_ROOT /opt/pyenv
ENV PATH ${PYENV_ROOT}/shims:${PYENV_ROOT}/bin:${PATH}
RUN curl https://pyenv.run | bash
RUN pyenv install ${PYTHON_VERSION} && \
    pyenv global ${PYTHON_VERSION}
RUN pip install --no-cache-dir -U requests tqdm opencv-python-headless
RUN pip install --no-cache-dir -U tensorflow-gpu==1.15.4
RUN pip install --no-cache-dir -U torch==1.10.2+cu102 torchvision==0.11.3+cu102 -f https://download.pytorch.org/whl/torch/ -f https://download.pytorch.org/whl/torchvision/
RUN rm -rf ${HOME}/.cache/pip

WORKDIR /work
ENV PYTHONPATH /work/:${PYTHONPATH}

docker build . -t stylegan2_pytorch

Convert

git clone https://github.com/NVLabs/stylegan2
docker run --rm -it -u $(id -u):$(id -g) -e XDG_CACHE_HOME=/work --ipc host --gpus all -w /work -v `pwd`:/work stylegan2_pytorch python convert_weight.py --repo stylegan2 aydao-anime-danbooru2019s-512-5268480.pkl

Usage

Apply patch

--- a/generate.py
+++ b/generate.py
@@ -6,21 +6,25 @@ from model import Generator
 from tqdm import tqdm


-def generate(args, g_ema, device, mean_latent):
+def generate(args, g_ema, device, mean_latent, randomize_noise):

     with torch.no_grad():
         g_ema.eval()
         for i in tqdm(range(args.pics)):
-            sample_z = torch.randn(args.sample, args.latent, device=device)
+            samples = []
+            for _ in range(args.split):
+                sample_z = torch.randn(args.sample // args.split, args.latent, device=device)

-            sample, _ = g_ema(
-                [sample_z], truncation=args.truncation, truncation_latent=mean_latent
-            )
+                sample, _ = g_ema(
+                        [sample_z], truncation=args.truncation, truncation_latent=mean_latent,
+                    randomize_noise=randomize_noise
+                )
+                samples.extend(sample)

             utils.save_image(
-                sample,
-                f"sample/{str(i).zfill(6)}.png",
-                nrow=1,
+                samples,
+                f"{args.output_dir}/{str(i).zfill(6)}.{args.ext}",
+                nrow=args.ncol,
                 normalize=True,
                 range=(-1, 1),
             )
@@ -30,6 +34,8 @@ if __name__ == "__main__":
     device = "cuda"

     parser = argparse.ArgumentParser(description="Generate samples from the generator")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--output-dir", '-o', type=str, required=True)

     parser.add_argument(
         "--size", type=int, default=1024, help="output image size of the generator"
@@ -37,11 +43,14 @@ if __name__ == "__main__":
     parser.add_argument(
         "--sample",
         type=int,
-        default=1,
+        default=100,
         help="number of samples to be generated for each image",
     )
+    parser.add_argument("--ncol", type=int, default=10)
+    parser.add_argument("--split", type=int, default=4)
+    parser.add_argument("--ext", type=str, default='png')
     parser.add_argument(
-        "--pics", type=int, default=20, help="number of images to be generated"
+        "--pics", type=int, default=1, help="number of images to be generated"
     )
     parser.add_argument("--truncation", type=float, default=1, help="truncation ratio")
     parser.add_argument(
@@ -62,23 +71,31 @@ if __name__ == "__main__":
         default=2,
         help="channel multiplier of the generator. config-f = 2, else = 1",
     )
+    parser.add_argument("--additional_multiplier", type=int, default=1)
+    parser.add_argument("--load_latent_vec", action='store_true')
+    parser.add_argument("--no-randomize-noise", dest='randomize_noise', action='store_false')
+    parser.add_argument("--n_mlp", type=int, default=8)

     args = parser.parse_args()

-    args.latent = 512
-    args.n_mlp = 8
+    seed = args.seed
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    args.latent = 512 * args.additional_multiplier

     g_ema = Generator(
-        args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier
+        args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier,
+        additional_multiplier=args.additional_multiplier
     ).to(device)
     checkpoint = torch.load(args.ckpt)

-    g_ema.load_state_dict(checkpoint["g_ema"])
+    g_ema.load_state_dict(checkpoint["g_ema"], strict=True)

-    if args.truncation < 1:
+    if not args.load_latent_vec:
         with torch.no_grad():
             mean_latent = g_ema.mean_latent(args.truncation_mean)
     else:
-        mean_latent = None
+        mean_latent = checkpoint['latent_avg'].to(device)

-    generate(args, g_ema, device, mean_latent)
+    generate(args, g_ema, device, mean_latent, randomize_noise=args.randomize_noise)

Run

python generate.py --ckpt aydao-anime-danbooru2019s-512-5268480.pt --size 512 --n_mlp 4 --additional_multiplier 2 --load_latent_vec --no-randomize-noise -o out_images --truncation 0.6 --seed 333 --pics 1 --sample 48 --ncol 8 --ext jpg