Update generate.ipynb

drboog · Mar 27, 2022 · 5b48186 · 5b48186
1 parent a69b916
commit 5b48186
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/generate.ipynb b/generate.ipynb
@@ -34,19 +34,14 @@
     "        self.name = 'generator'\n",
     "        self.model = self.load_model(device, path)\n",
     "        self.device = device\n",
-    "        if device == 'cpu':\n",
-    "            self.force_32 = True\n",
-    "        else:\n",
-    "            self.force_32 = False\n",
+    "        self.force_32 = False\n",
     "        \n",
     "    def load_model(self, device, path):\n",
     "        with dnnlib.util.open_url(path) as f:\n",
     "            network= legacy.load_network_pkl(f)\n",
     "            self.G_ema = network['G_ema'].to(device)\n",
     "            self.D = network['D'].to(device)\n",
     "#                 self.G = network['G'].to(device)\n",
-    "            if device == 'cpu':\n",
-    "                self.G_ema = self.G_ema.float()\n",
     "            return self.G_ema\n",
     "        \n",
     "    def generate(self, z, c, fts, noise_mode='const', return_styles=True):\n",
@@ -76,7 +71,7 @@
    "source": [
     "with torch.no_grad():\n",
     "\n",
-    "    device = 'cuda:0'\n",
+    "    device = 'cuda:0' # please use GPU, do not use CPU\n",
     "    path = './some_pre-trained_models.pkl'  # pre-trained model\n",
     "    generator = Generator(device=device, path=path)\n",
     "    clip_model, _ = clip.load(\"ViT-B/32\", device=device)\n",