Skip to content

Instantly share code, notes, and snippets.

@afiaka87

afiaka87/zeta_quantize.ipynb

Last active November 27, 2023 04:43
Show Gist options
  • Save afiaka87/a97cca3b54c02209b94ff805224f9eb5 to your computer and use it in GitHub Desktop.
Save afiaka87/a97cca3b54c02209b94ff805224f9eb5 to your computer and use it in GitHub Desktop.
Download ZIP
zeta_quantize.ipynb
Display the source blob
Display the rendered blob
Raw
zeta_quantize.ipynb
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "zeta_quantize.ipynb",
"private_outputs": true,
"provenance": [],
"collapsed_sections": [
"CppIQlPhhwhs",
"U_M1vTTyFjtH"
],
"include_colab_link": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"language_info": {
"name": "python"
},
"accelerator": "GPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/afiaka87/a97cca3b54c02209b94ff805224f9eb5/zeta_quantize.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "CppIQlPhhwhs"
},
"source": [
"# CLIP-guided VQGAN (z-quantize method) by [RiversHaveWings](https://twitter.com/RiversHaveWings)\n",
"\n",
"## `status` - Runs. May need to tinker with settings.\n",
"\n",
"## Generate Images From Text Using Several Available Pretrained VQGAN Models\n",
"- https://github.com/crowsonkb\n",
"- https://twitter.com/RiversHaveWings"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AB0yOX1sJrf2"
},
"source": [
"Notebook's presentation modified by [afiaka87](https://github.com/afiaka87). \n",
"\n",
"Other contributors include:\n",
"- Eleiber#8347 - tweaks and documentation\n",
"- Crimeacs#8222 - Pooling trick\n",
"- Abulafia#3734 - GUI Help"
]
},
{
"cell_type": "code",
"metadata": {
"id": "WhYxLaY72Oqf",
"cellView": "form"
},
"source": [
"#@title Licensed under the MIT License\n",
"# Copyright (c) 2021 Katherine Crowson\n",
"# Permission is hereby granted, free of charge, to any person obtaining a copy\n",
"# of this software and associated documentation files (the \"Software\"), to deal\n",
"# in the Software without restriction, including without limitation the rights\n",
"# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
"# copies of the Software, and to permit persons to whom the Software is\n",
"# furnished to do so, subject to the following conditions:\n",
"\n",
"# The above copyright notice and this permission notice shall be included in\n",
"# all copies or substantial portions of the Software.\n",
"\n",
"# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
"# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
"# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
"# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
"# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
"# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\n",
"# THE SOFTWARE."
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "U_M1vTTyFjtH"
},
"source": [
"# (Required) Setup"
]
},
{
"cell_type": "code",
"metadata": {
"id": "fDhVPH7Y0QJm",
"cellView": "form"
},
"source": [
"#@markdown # Setup\n",
"#@markdown \n",
"#@markdown ## Download a pretrained `vqgan`\n",
"#@markdown \n",
"#@markdown A `vqgan` efficiently represents images.\n",
"#@markdown `imagenet-16384` is very good.\n",
"%pip install kornia\n",
"%pip install einops\n",
"\n",
"\n",
"!wget --continue 'https://heibox.uni-heidelberg.de/d/8088892a516d4e3baf92/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1' -O vqgan_imagenet_f16_1024.yaml\n",
"!wget --continue 'https://heibox.uni-heidelberg.de/d/8088892a516d4e3baf92/files/?p=%2Fckpts%2Flast.ckpt&dl=1' -O vqgan_imagenet_f16_1024.ckpt\n",
"!wget --continue 'https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1' -O vqgan_imagenet_f16_16384.yaml\n",
"!wget --continue 'https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/files/?p=%2Fckpts%2Flast.ckpt&dl=1' -O vqgan_imagenet_f16_16384.ckpt\n",
"!wget --continue 'https://heibox.uni-heidelberg.de/seafhttp/files/a6b39678-39cc-4f72-a29e-9918b0f7365f/last.ckpt' -O gumbelvqgan_oi_f8_8192.ckpt\n",
"!wget --continue 'https://heibox.uni-heidelberg.de/seafhttp/files/50892827-06a5-4716-b1a2-7af99e890871/model.yaml' -O 'gumbelvqgan_oi_f8_8192.yaml'\n",
"\n",
"from IPython.display import clear_output\n",
"clear_output()"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "3Kj79sSaIp54",
"cellView": "form"
},
"source": [
"#@title # Install CLIP and VQGAN codebase and dependencies\n",
"from IPython.display import clear_output\n",
"%rm -r /content/CLIP/\n",
"%rm -r /content/taming-transformers/\n",
"\n",
"!git clone https://github.com/openai/CLIP\n",
"!git clone https://github.com/CompVis/taming-transformers\n",
"%pip install --upgrade -e /content/taming-transformers/\n",
"# %pip install --upgrade -e /content/CLIP/\n",
"\n",
"# TODO Use these after PR merges.\n",
"# %pip install clip-anytorch\n",
"# %pip install taming-transformers-rom1504\n",
"%pip install ftfy regex tqdm omegaconf pytorch-lightning\n",
"from IPython.display import clear_output\n",
"clear_output()\n"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "HqGbT9_h2LFf",
"cellView": "form"
},
"source": [
"#@markdown # `z+quantize` sampling via CLIP (code)\n",
"\n",
"%pip install imageio-ffmpeg\n",
"import kornia.augmentation as K\n",
"import imageio\n",
"import math\n",
"from pathlib import Path\n",
"import sys\n",
"from copy import deepcopy\n",
"from urllib.request import urlopen\n",
"import numpy as np\n",
"sys.path.append('./taming-transformers')\n",
"from IPython import display\n",
"from omegaconf import OmegaConf\n",
"from PIL import Image\n",
"from taming.models import cond_transformer, vqgan\n",
"import torch\n",
"from torch import nn, optim\n",
"from torch.nn import functional as F\n",
"from torchvision import transforms\n",
"from torchvision.transforms import functional as TF\n",
"from tqdm.notebook import tqdm\n",
"\n",
"from CLIP import clip\n",
"\n",
"class AddGaussianNoise(object):\n",
" def __init__(self, mean=0., std=1.):\n",
" self.std = std\n",
" self.mean = mean\n",
" \n",
" def __call__(self, tensor):\n",
" return tensor + torch.randn(tensor.size()) * self.std + self.mean\n",
" \n",
" def __repr__(self):\n",
" return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std)\n",
"\n",
"def sinc(x):\n",
" return torch.where(x != 0, torch.sin(math.pi * x) / (math.pi * x), x.new_ones([]))\n",
"\n",
"\n",
"def lanczos(x, a):\n",
" cond = torch.logical_and(-a < x, x < a)\n",
" out = torch.where(cond, sinc(x) * sinc(x/a), x.new_zeros([]))\n",
" return out / out.sum()\n",
"\n",
"\n",
"def ramp(ratio, width):\n",
" n = math.ceil(width / ratio + 1)\n",
" out = torch.empty([n])\n",
" cur = 0\n",
" for i in range(out.shape[0]):\n",
" out[i] = cur\n",
" cur += ratio\n",
" return torch.cat([-out[1:].flip([0]), out])[1:-1]\n",
"\n",
"\n",
"def resample(input, size, align_corners=True):\n",
" n, c, h, w = input.shape\n",
" dh, dw = size\n",
"\n",
" input = input.view([n * c, 1, h, w])\n",
"\n",
" if dh < h:\n",
" kernel_h = lanczos(ramp(dh / h, 2), 2).to(input.device, input.dtype)\n",
" pad_h = (kernel_h.shape[0] - 1) // 2\n",
" input = F.pad(input, (0, 0, pad_h, pad_h), 'reflect')\n",
" input = F.conv2d(input, kernel_h[None, None, :, None])\n",
"\n",
" if dw < w:\n",
" kernel_w = lanczos(ramp(dw / w, 2), 2).to(input.device, input.dtype)\n",
" pad_w = (kernel_w.shape[0] - 1) // 2\n",
" input = F.pad(input, (pad_w, pad_w, 0, 0), 'reflect')\n",
" input = F.conv2d(input, kernel_w[None, None, None, :])\n",
"\n",
" input = input.view([n, c, h, w])\n",
" return F.interpolate(input, size, mode='bicubic', align_corners=align_corners)\n",
"\n",
"\n",
"class ReplaceGrad(torch.autograd.Function):\n",
" @staticmethod\n",
" def forward(ctx, x_forward, x_backward):\n",
" ctx.shape = x_backward.shape\n",
" return x_forward\n",
"\n",
" @staticmethod\n",
" def backward(ctx, grad_in):\n",
" return None, grad_in.sum_to_size(ctx.shape)\n",
"\n",
"\n",
"replace_grad = ReplaceGrad.apply\n",
"\n",
"\n",
"class ClampWithGrad(torch.autograd.Function):\n",
" @staticmethod\n",
" def forward(ctx, input, min, max):\n",
" ctx.min = min\n",
" ctx.max = max\n",
" ctx.save_for_backward(input)\n",
" return input.clamp(min, max)\n",
"\n",
" @staticmethod\n",
" def backward(ctx, grad_in):\n",
" input, = ctx.saved_tensors\n",
" return grad_in * (grad_in * (input - input.clamp(ctx.min, ctx.max)) >= 0), None, None\n",
"\n",
"\n",
"clamp_with_grad = ClampWithGrad.apply\n",
"\n",
"\n",
"def vector_quantize(x, codebook):\n",
" d = x.pow(2).sum(dim=-1, keepdim=True) + codebook.pow(2).sum(dim=1) - 2 * x @ codebook.T\n",
" indices = d.argmin(-1)\n",
" x_q = F.one_hot(indices, codebook.shape[0]).to(d.dtype) @ codebook\n",
" return replace_grad(x_q, x)\n",
"\n",
"\n",
"class Prompt(nn.Module):\n",
" def __init__(self, embed, weight=1., stop=float('-inf')):\n",
" super().__init__()\n",
" self.register_buffer('embed', embed)\n",
" self.register_buffer('weight', torch.as_tensor(weight))\n",
" self.register_buffer('stop', torch.as_tensor(stop))\n",
"\n",
" def forward(self, input):\n",
" input_normed = F.normalize(input.unsqueeze(1), dim=2)\n",
" embed_normed = F.normalize(self.embed.unsqueeze(0), dim=2)\n",
" dists = input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2)\n",
" dists = dists * self.weight.sign()\n",
" return self.weight.abs() * replace_grad(dists, torch.maximum(dists, self.stop)).mean()\n",
"\n",
"\n",
"def fetch(url_or_path):\n",
" if str(url_or_path).startswith('http://') or str(url_or_path).startswith('https://'):\n",
" r = requests.get(url_or_path)\n",
" r.raise_for_status()\n",
" fd = io.BytesIO()\n",
" fd.write(r.content)\n",
" fd.seek(0)\n",
" return fd\n",
" return open(url_or_path, 'rb')\n",
"\n",
"\n",
"def parse_prompt(prompt):\n",
" if prompt.startswith('http://') or prompt.startswith('https://'):\n",
" vals = prompt.rsplit(':', 3)\n",
" vals = [vals[0] + ':' + vals[1], *vals[2:]]\n",
" else:\n",
" vals = prompt.rsplit(':', 2)\n",
" vals = vals + ['', '1', '-inf'][len(vals):]\n",
" return vals[0], float(vals[1]), float(vals[2])\n",
"\n",
"\n",
"\n",
"class MakeCutouts(nn.Module):\n",
" def __init__(self, cut_size, cutn, cut_pow=1.):\n",
" super().__init__()\n",
" self.cut_size = cut_size\n",
" self.cutn = cutn\n",
" self.cut_pow = cut_pow\n",
"\n",
" def forward(self, input):\n",
" sideY, sideX = input.shape[2:4]\n",
" max_size = min(sideX, sideY)\n",
" min_size = min(sideX, sideY, self.cut_size)\n",
" cutouts = []\n",
" for _ in range(self.cutn):\n",
" size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size)\n",
" offsetx = torch.randint(0, sideX - size + 1, ())\n",
" offsety = torch.randint(0, sideY - size + 1, ())\n",
" cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]\n",
" cutouts.append(resample(cutout, (self.cut_size, self.cut_size)))\n",
" return clamp_with_grad(torch.cat(cutouts, dim=0), 0, 1)\n",
"# class MakeCutouts(nn.Module):\n",
"# def __init__(self, cut_size, cutn, cut_pow=1.):\n",
"# super().__init__()\n",
"# self.cut_size = cut_size\n",
"# self.cutn = cutn\n",
"# self.cut_pow = cut_pow\n",
"# self.augs = transforms.Compose([\n",
"# AddGaussianNoise(0.0, 1.0),\n",
"# transforms.RandomSolarize(threshold=16, p=0.1),\n",
"# transforms.RandomHorizontalFlip(p=0.5),\n",
"# transforms.RandomAdjustSharpness(sharpness_factor=0, p=0.1),\n",
"# transforms.RandomAdjustSharpness(sharpness_factor=2, p=0.1),\n",
"# transforms.ColorJitter(brightness=0.0, contrast=0.0, saturation=0.2, hue=0.2),\n",
"# transforms.RandomResizedCrop(cut_size, interpolation=T.InterpolationMode.LANCZOS),\n",
"# # transforms.Resize(cut_size*2, interpolation=T.InterpolationMode.LANCZOS),\n",
"# ])\n",
"\n",
" \n",
"# def forward(self, input):\n",
"# sideY, sideX = input.shape[2:4]\n",
"# max_size = min(sideX, sideY)\n",
"# min_size = min(sideX, sideY, self.cut_size)\n",
"# cutouts = []\n",
"# for _ in range(self.cutn):\n",
"# size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size)\n",
"# offsetx = torch.randint(0, sideX - size + 1, ())\n",
"# offsety = torch.randint(0, sideY - size + 1, ())\n",
"# cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]\n",
"# cutouts.append(resample(cutout, (self.cut_size, self.cut_size)))\n",
"# batch = self.augs(torch.cat(cutouts, dim=0))\n",
"# return batch\n",
"\n",
"# # def __init__(self, cut_size, cutn, cut_pow=1.):\n",
"# # super().__init__()\n",
"# # self.cut_size = cut_size\n",
"# # self.cutn = cutn\n",
"# # self.cut_pow = cut_pow\n",
"# # self.noise_fac = 0.1\n",
"# # self.av_pool = nn.AdaptiveAvgPool2d((self.cut_size, self.cut_size))\n",
"# # self.max_pool = nn.AdaptiveMaxPool2d((self.cut_size, self.cut_size))\n",
"# # self.augs = nn.Sequential(\n",
"# # K.RandomHorizontalFlip(p=0.2),\n",
"# # K.RandomVerticalFlip(p=0.1),\n",
"# # K.RandomSolarize(0.01, 0.01, p=0.1),\n",
"# # K.RandomSharpness(0.3,p=0.1),\n",
"# # K.RandomAffine(degrees=7, translate=0.1, p=0.1, padding_mode='border'),\n",
"# # K.RandomPerspective(0.7,p=0.2),\n",
"# # K.ColorJitter(hue=0.1, saturation=0.1, p=0.5),\n",
"# # K.RandomErasing((.1, .4), (.3, 1/.3), same_on_batch=True, p=0.1),\n",
"# # )\n",
"\n",
"\n",
"# # def forward(self, input):\n",
"# # sideY, sideX = input.shape[2:4]\n",
"# # max_size = min(sideX, sideY)\n",
"# # min_size = min(sideX, sideY, self.cut_size)\n",
"# # cutouts = []\n",
" \n",
"# # for _ in range(self.cutn):\n",
"# # cutout = (self.av_pool(input) + self.max_pool(input)) / 2\n",
"# # cutouts.append(cutout)\n",
"# # batch = self.augs(torch.cat(cutouts, dim=0))\n",
"# # if self.noise_fac:\n",
"# # facs = batch.new_empty([self.cutn, 1, 1, 1]).uniform_(0, self.noise_fac)\n",
"# # batch = batch + facs * torch.randn_like(batch)\n",
"# # return batch\n",
"\n",
"\n",
"def load_vqgan_model(config_path, checkpoint_path):\n",
" config = OmegaConf.load(config_path)\n",
" if config.model.target == 'taming.models.vqgan.VQModel':\n",
" model = vqgan.VQModel(**config.model.params)\n",
" model.eval().requires_grad_(False)\n",
" model.init_from_ckpt(checkpoint_path)\n",
" elif config.model.target == 'taming.models.cond_transformer.Net2NetTransformer':\n",
" parent_model = cond_transformer.Net2NetTransformer(**config.model.params)\n",
" parent_model.eval().requires_grad_(False)\n",
" parent_model.init_from_ckpt(checkpoint_path)\n",
" model = parent_model.first_stage_model\n",
" elif config.model.target == 'taming.models.vqgan.GumbelVQ':\n",
" model = vqgan.GumbelVQ(**config.model.params)\n",
" model.eval().requires_grad_(False)\n",
" model.init_from_ckpt(checkpoint_path)\n",
" else:\n",
" raise ValueError(f'unknown model type: {config.model.target}')\n",
" del model.loss\n",
" return model\n",
"\n",
"\n",
"def resize_image(image, out_size):\n",
" ratio = image.size[0] / image.size[1]\n",
" area = min(image.size[0] * image.size[1], out_size[0] * out_size[1])\n",
" size = round((area * ratio)**0.5), round((area / ratio)**0.5)\n",
" return image.resize(size, Image.LANCZOS)\n",
"\n",
"from IPython.display import clear_output \n",
"clear_output()"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "7NNQxxa4FuqK"
},
"source": [
"# Configure"
]
},
{
"cell_type": "code",
"metadata": {
"id": "aw7MEJThAfDv",
"cellView": "form"
},
"source": [
"#@title # Config\n",
"from IPython.display import clear_output\n",
"import torchvision.transforms as T\n",
"import torchvision\n",
"import argparse\n",
"# prompts='the first day of the waters', #@param\n",
"prompts = \"an orange hue sunset in the background. a surfer riding a surfboard under a huge wave.\" #@param {type:\"string\"}\n",
"vqgan_model = \"vqgan_imagenet_f16_16384\" #@param [\"vqgan_imagenet_f16_16384\", \"vqgan_imagenet_f16_1024\", \"gumbelvqgan_oi_f8_8192\"]\n",
"noise_prompt_seeds = [] \n",
"noise_prompt_weights = [] \n",
"size_x = 440 #@param {type: \"integer\"}\n",
"size_y = 440#@param {type: \"integer\"}\n",
"\n",
"\n",
"#@markdown `initial_weight` controls how much the image is allowed t\n",
"#@markdown diverge from the initial image, 0 is fully unconstrained and higher means divergence is penalized more\n",
"\n",
"initial_image='' \n",
"initial_weight=0.0\n",
"clip_model = \"ViT-B/16\" #@param [\"ViT-B/16\", \"RN50\", \"RN101\", \"RN50x4\", \"ViT-B/32\"]\n",
"learning_rate=0.5 #@param {type: \"number\"}\n",
"num_random_cutouts=64 #@param {type: 'integer' }\n",
"#@markdown `cutout_power` \"lower = more global coherence\"\n",
"cutout_power=1. #@param {type: 'number' } \n",
"display_frequency=50 #@param {type: 'integer' } \n",
"seed=0 #@param {type: 'integer' }\n",
"\n",
"num_iterations=1550 #@param {type: 'integer'} \n",
"\n",
"vqgan_config=f'{vqgan_model}.yaml'\n",
"vqgan_checkpoint=f'{vqgan_model}.ckpt' \n",
"\n",
"args = argparse.Namespace(\n",
" prompts=[prompts],\n",
" image_prompts=[],\n",
" noise_prompt_seeds=[],\n",
" noise_prompt_weights=[],\n",
" size=[size_x,size_y],\n",
" init_image=initial_image,\n",
" init_weight=initial_weight,\n",
" clip_model=clip_model,\n",
" vqgan_config=vqgan_config,\n",
" vqgan_checkpoint=vqgan_checkpoint,\n",
" step_size=learning_rate,\n",
" cutn=num_random_cutouts,\n",
" cut_pow=cutout_power,\n",
" display_freq=display_frequency,\n",
" seed=seed,\n",
" iterations=num_iterations,\n",
")\n",
"\n",
"clear_output()\n",
"print(args)\n",
"print(\"Run the next cell to generate an image with these settings.\")"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "iqeWW9yCFzBd"
},
"source": [
"# Run generation with current settings."
]
},
{
"cell_type": "code",
"metadata": {
"id": "sXgVarOm3gJI",
"cellView": "form"
},
"source": [
"#@title Generate an Image from a Prompt and/or a starting image.\n",
"#@markdown May take a while depending on settings. \n",
"\n",
"#@markdown Consider using `clip_model=ViT-B/32` and `vqgan_model=vqgan_imagenet_1024_f16` for faster results.\n",
"from IPython.display import clear_output\n",
"device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
"print('Using device:', device)\n",
"\n",
"model = load_vqgan_model(args.vqgan_config, args.vqgan_checkpoint).to(device)\n",
"perceptor = clip.load(args.clip_model, jit=False)[0].eval().requires_grad_(False).to(device)\n",
"\n",
"cut_size = perceptor.visual.input_resolution\n",
"if vqgan_model == \"gumbelvqgan_oi_f8_8192\":\n",
" e_dim = model.quantize.embedding_dim\n",
"else:\n",
" e_dim = model.quantize.e_dim\n",
"f = 2**(model.decoder.num_resolutions - 1)\n",
"make_cutouts = MakeCutouts(cut_size, args.cutn, cut_pow=args.cut_pow)\n",
"if vqgan_model == \"gumbelvqgan_oi_f8_8192\":\n",
" n_toks = model.quantize.n_embed\n",
"else:\n",
" n_toks = model.quantize.n_e\n",
"toksX, toksY = args.size[0] // f, args.size[1] // f\n",
"sideX, sideY = toksX * f, toksY * f\n",
"if vqgan_model == \"gumbelvqgan_oi_f8_8192\":\n",
" z_min = model.quantize.embed.weight.min(dim=0).values[None, :, None, None]\n",
" z_max = model.quantize.embed.weight.max(dim=0).values[None, :, None, None]\n",
"else:\n",
" z_min = model.quantize.embedding.weight.min(dim=0).values[None, :, None, None]\n",
" z_max = model.quantize.embedding.weight.max(dim=0).values[None, :, None, None]\n",
"\n",
"if args.seed is not None:\n",
" torch.manual_seed(args.seed)\n",
"\n",
"if args.init_image:\n",
" pil_image = Image.open(fetch(args.init_image)).convert('RGB')\n",
" pil_image = pil_image.resize((sideX, sideY), Image.LANCZOS)\n",
" z, *_ = model.encode(TF.to_tensor(pil_image).to(device).unsqueeze(0) * 2 - 1)\n",
"else:\n",
" one_hot = F.one_hot(torch.randint(n_toks, [toksY * toksX], device=device), n_toks).float()\n",
" if vqgan_model == \"gumbelvqgan_oi_f8_8192\":\n",
" z = one_hot @ model.quantize.embed.weight\n",
" else:\n",
" z = one_hot @ model.quantize.embedding.weight\n",
" z = z.view([-1, toksY, toksX, e_dim]).permute(0, 3, 1, 2)\n",
"z_orig = z.clone()\n",
"z.requires_grad_(True)\n",
"opt = optim.AdamW([z], lr=args.step_size, betas=(0.9, 0.96), weight_decay=1e-8, amsgrad=True)\n",
"\n",
"normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])\n",
"\n",
"pMs = []\n",
"\n",
"for prompt in args.prompts:\n",
" txt, weight, stop = parse_prompt(prompt)\n",
" embed = perceptor.encode_text(clip.tokenize(txt).to(device)).float()\n",
" pMs.append(Prompt(embed, weight, stop).to(device))\n",
"\n",
"for prompt in args.image_prompts:\n",
" path, weight, stop = parse_prompt(prompt)\n",
" img = resize_image(Image.open(path).convert('RGB'), (sideX, sideY))\n",
" batch = make_cutouts(TF.to_tensor(img).unsqueeze(0).to(device))\n",
" embed = perceptor.encode_image(normalize(batch)).float()\n",
" pMs.append(Prompt(embed, weight, stop).to(device))\n",
"\n",
"for seed, weight in zip(args.noise_prompt_seeds, args.noise_prompt_weights):\n",
" gen = torch.Generator().manual_seed(seed)\n",
" embed = torch.empty([1, perceptor.visual.output_dim]).normal_(generator=gen)\n",
" pMs.append(Prompt(embed, weight).to(device))\n",
"\n",
"def synth(z):\n",
" if vqgan_model == \"gumbelvqgan_oi_f8_8192\":\n",
" z_q = vector_quantize(z.movedim(1, 3), model.quantize.embed.weight).movedim(3, 1)\n",
" else:\n",
" z_q = vector_quantize(z.movedim(1, 3), model.quantize.embedding.weight).movedim(3, 1)\n",
" return clamp_with_grad(model.decode(z_q).add(1).div(2), 0, 1)\n",
"\n",
"@torch.no_grad()\n",
"def checkin(iter, losses):\n",
" losses_str = ', '.join(f'{loss.item():g}' for loss in losses)\n",
" tqdm.write(f'i: {iter}, loss: {sum(losses).item():g}, losses: {losses_str}')\n",
" out = synth(z)\n",
" TF.to_pil_image(out[0].cpu()).save('progress.png')\n",
" display.display(display.Image('progress.png'))\n",
"\n",
"\n",
"def ascend_txt(current_iteration):\n",
" out = synth(z)\n",
" img_embed = perceptor.encode_image(make_cutouts(out)).float()\n",
" result = []\n",
"\n",
" if args.init_weight:\n",
" result.append(F.mse_loss(z, z_orig) * args.init_weight / 2)\n",
" # result.append(F.mse_loss(z, torch.zeros_like(z_orig)) * ((1/torch.tensor(current_iteration*2 + 1))*args.init_weight) / 2)\n",
" for prompt in pMs:\n",
" result.append(prompt(img_embed))\n",
" img = np.array(out.mul(255).clamp(0, 255)[0].cpu().detach().numpy().astype(np.uint8))[:,:,:]\n",
" img = np.transpose(img, (1, 2, 0))\n",
" %mkdir -p '/content/steps/'\n",
" imageio.imwrite('/content/steps/' + str(current_iteration) + '.png', np.array(img))\n",
" return result\n",
"\n",
"\n",
"def train(iter):\n",
" opt.zero_grad()\n",
" lossAll = ascend_txt(iter)\n",
" if iter % args.display_freq == 0:\n",
" clear_output(wait=True)\n",
" checkin(iter, lossAll)\n",
" loss = sum(lossAll)\n",
" loss.backward()\n",
" opt.step()\n",
" with torch.no_grad():\n",
" z.copy_(z.maximum(z_min).minimum(z_max))\n",
"\n",
"\n",
"try:\n",
" for iteration in range(args.iterations):\n",
" train(iteration)\n",
"except KeyboardInterrupt:\n",
" pass"
],
"execution_count": null,
"outputs": []
}
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment