Spaces:
Running
Running
File size: 3,859 Bytes
891b88f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
import os
from functools import lru_cache
from typing import Any, Callable, Dict, List, Optional, Tuple
import attr
import numpy as np
import torch
import torch.nn as nn
import yaml
from glide_text2im.tokenizer.simple_tokenizer import SimpleTokenizer
from .encoders import ImageEncoder, TextEncoder
@lru_cache()
def default_config_path() -> str:
return os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.yaml")
@attr.s
class CLIPModel:
config: Dict[str, Any] = attr.ib()
text_encoder: nn.Module = attr.ib()
image_encoder: nn.Module = attr.ib()
logit_scale: torch.Tensor = attr.ib()
device: torch.device = attr.ib()
tokenizer: SimpleTokenizer = attr.ib()
def encode_prompts(self, prompts: List[str]) -> Tuple[torch.Tensor, torch.Tensor]:
tokens = []
lens = []
for prompt in prompts:
sub_tokens, sub_len = self.tokenizer.padded_tokens_and_len(
self.tokenizer.encode(prompt), self.text_encoder.max_text_len
)
tokens.append(sub_tokens)
lens.append(sub_len)
return (
torch.tensor(tokens).to(dtype=torch.long, device=self.device),
torch.tensor(lens).to(dtype=torch.long, device=self.device),
)
def text_embeddings(self, prompts: List[str]) -> torch.Tensor:
tokens, lens = self.encode_prompts(prompts)
z_t = self.text_encoder(tokens, lens)
return z_t / (torch.linalg.norm(z_t, dim=-1, keepdim=True) + 1e-12)
def image_embeddings(self, images: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
z_i = self.image_encoder((images + 1) * 127.5, t)
return z_i / (torch.linalg.norm(z_i, dim=-1, keepdim=True) + 1e-12)
def cond_fn(self, prompts: List[str], grad_scale: float) -> Callable[..., torch.Tensor]:
with torch.no_grad():
z_t = self.text_embeddings(prompts)
def cond_fn(x, t, grad_scale=grad_scale, **kwargs):
with torch.enable_grad():
x_var = x.detach().requires_grad_(True)
z_i = self.image_embeddings(x_var, t)
loss = torch.exp(self.logit_scale) * (z_t * z_i).sum()
grad = torch.autograd.grad(loss, x_var)[0].detach()
return grad * grad_scale
return cond_fn
def create_clip_model(
config_path: Optional[str] = None,
device: Optional[torch.device] = None,
tokenizer: Optional[SimpleTokenizer] = None,
) -> CLIPModel:
if config_path is None:
config_path = default_config_path()
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if tokenizer is None:
tokenizer = SimpleTokenizer()
with open(config_path, "r") as f:
config = yaml.load(f, Loader=yaml.SafeLoader)
text_encoder = TextEncoder(
n_bpe_vocab=config["n_vocab"],
max_text_len=config["max_text_len"],
n_embd=config["n_embd"],
n_head=config["n_head_text"],
n_xf_blocks=config["n_xf_blocks_text"],
n_head_state=config["n_head_state_text"],
device=device,
)
image_encoder = ImageEncoder(
image_size=config["image_size"],
patch_size=config["patch_size"],
n_embd=config["n_embd"],
n_head=config["n_head_image"],
n_xf_blocks=config["n_xf_blocks_image"],
n_head_state=config["n_head_state_image"],
n_timestep=config["n_timesteps"],
device=device,
)
logit_scale = torch.tensor(
np.log(config["logit_scale"]),
dtype=torch.float32,
device=device,
requires_grad=False,
)
return CLIPModel(
config=config,
text_encoder=text_encoder,
image_encoder=image_encoder,
logit_scale=logit_scale,
device=device,
tokenizer=tokenizer,
)
|