Add Cog file to build Docker image

Author: andreasjansson

README.md

@@ -11,6 +11,8 @@ I found bugs in the implementation thanks to @adambielski and @TropComplique! (h
 
 Implementation of A Style-Based Generator Architecture for Generative Adversarial Networks (https://arxiv.org/abs/1812.04948) in PyTorch
 
+* [Demo and Docker image on Replicate](https://replicate.ai/rosinality/style-based-gan-pytorch)
+
 Usage:
 
 You should prepare lmdb dataset

cog.yaml

@@ -0,0 +1,9 @@
+predict: predict.py:Predictor
+build:
+  python_version: 3.8
+  python_packages:
+    - torch==1.7.0
+    - torchvision==0.8.1
+    - tqdm==4.59.0
+    - pillow==8.1.2
+    - lmdb==1.1.1

predict.py

@@ -0,0 +1,40 @@
+import os
+import math
+import tempfile
+from pathlib import Path
+import torch
+from torchvision import utils
+import cog
+
+from generate import sample, get_mean_style
+from model import StyledGenerator
+
+SIZE = 1024
+
+
+class Predictor(cog.Predictor):
+    def setup(self):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.generator = StyledGenerator(512).to(self.device)
+        print("Loading checkpoint")
+        self.generator.load_state_dict(
+            torch.load(
+                "stylegan-1024px-new.model",
+                map_location=self.device,
+            )["g_running"],
+        )
+        self.generator.eval()
+
+    @cog.input("seed", type=int, default=-1, help="Random seed, -1 for random")
+    def predict(self, seed):
+        if seed < 0:
+            seed = int.from_bytes(os.urandom(2), "big")
+        torch.manual_seed(seed)
+        print(f"seed: {seed}")
+
+        mean_style = get_mean_style(self.generator, self.device)
+        step = int(math.log(SIZE, 2)) - 2
+        img = sample(self.generator, step, mean_style, 1, self.device)
+        output_path = Path(tempfile.mkdtemp()) / "output.png"
+        utils.save_image(img, output_path, normalize=True)
+        return output_path