Spaces:
Runtime error
Runtime error
| # Modified from https://huggingface.co/spaces/akhaliq/CLIP_prefix_captioning/blob/main/app.py | |
| import os | |
| os.system("gdown https://drive.google.com/uc?id=1_8v2ZUUaf9hhXP35jESXJ_Hgzl_rmufh") | |
| import clip | |
| import os | |
| from torch import nn | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as nnf | |
| import sys | |
| from typing import Tuple, List, Union, Optional | |
| from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup | |
| from tqdm import tqdm, trange | |
| import skimage.io as io | |
| import PIL.Image | |
| import gradio as gr | |
| N = type(None) | |
| V = np.array | |
| ARRAY = np.ndarray | |
| ARRAYS = Union[Tuple[ARRAY, ...], List[ARRAY]] | |
| VS = Union[Tuple[V, ...], List[V]] | |
| VN = Union[V, N] | |
| VNS = Union[VS, N] | |
| T = torch.Tensor | |
| TS = Union[Tuple[T, ...], List[T]] | |
| TN = Optional[T] | |
| TNS = Union[Tuple[TN, ...], List[TN]] | |
| TSN = Optional[TS] | |
| TA = Union[T, ARRAY] | |
| D = torch.device | |
| CPU = torch.device('cpu') | |
| def get_device(device_id: int) -> D: | |
| if not torch.cuda.is_available(): | |
| return CPU | |
| device_id = min(torch.cuda.device_count() - 1, device_id) | |
| return torch.device(f'cuda:{device_id}') | |
| CUDA = get_device | |
| class MLP(nn.Module): | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.model(x) | |
| def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh): | |
| """Project clip output to embedding of first prefix_length tokens""" | |
| super(MLP, self).__init__() | |
| layers = [] | |
| for i in range(len(sizes) - 1): | |
| layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias)) | |
| if i < len(sizes) - 2: | |
| layers.append(act()) | |
| # added some dropout here | |
| layers.append(nn.Dropout(p=0.2)) | |
| self.model = nn.Sequential(*layers) | |
| class ClipCaptionModel(nn.Module): | |
| def get_dummy_token(self, batch_size: int, device: torch.device) -> torch.Tensor: | |
| """Generate prefix tokens, shape Bxprefix_length""" | |
| return torch.zeros( | |
| batch_size, self.prefix_length, dtype=torch.int64, device=device | |
| ) | |
| def forward( | |
| self, | |
| tokens: torch.Tensor, | |
| prefix: torch.Tensor, | |
| mask: Optional[torch.Tensor] = None, | |
| labels: Optional[torch.Tensor] = None, | |
| ): | |
| embedding_text = self.gpt.transformer.wte(tokens) | |
| prefix_projections = self.clip_project(prefix).view( | |
| -1, self.prefix_length, self.gpt_embedding_size | |
| ) | |
| embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1) | |
| if labels is not None: | |
| dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device) | |
| labels = torch.cat((dummy_token, tokens), dim=1) | |
| out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask) | |
| return out | |
| def __init__(self, prefix_length: int = 10, prefix_size: int = 512): | |
| super(ClipCaptionModel, self).__init__() | |
| self.prefix_length = prefix_length | |
| self.gpt = GPT2LMHeadModel.from_pretrained("imthanhlv/gpt2news") | |
| self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1] | |
| self.clip_project = MLP( | |
| ( | |
| prefix_size, | |
| (self.gpt_embedding_size * prefix_length) // 2, | |
| self.gpt_embedding_size * prefix_length, | |
| ) | |
| ) | |
| class ClipCaptionPrefix(ClipCaptionModel): | |
| def parameters(self, recurse: bool = True): | |
| return self.clip_project.parameters() | |
| def train(self, mode: bool = True): | |
| super(ClipCaptionPrefix, self).train(mode) | |
| self.gpt.eval() | |
| return self | |
| #@title Caption prediction | |
| def generate_beam(model, tokenizer, beam_size: int = 5, prompt=None, embed=None, | |
| entry_length=67, temperature=1., stop_token: str = '.'): | |
| model.eval() | |
| stop_token_index = tokenizer.encode(stop_token)[0] | |
| tokens = None | |
| scores = None | |
| device = next(model.parameters()).device | |
| seq_lengths = torch.ones(beam_size, device=device) | |
| is_stopped = torch.zeros(beam_size, device=device, dtype=torch.bool) | |
| with torch.no_grad(): | |
| if embed is not None: | |
| generated = embed | |
| if prompt is not None: | |
| tokens = torch.tensor(tokenizer.encode(prompt)) | |
| tokens = tokens.unsqueeze(0).to(device) | |
| prompt_tokens = model.gpt.transformer.wte(tokens) | |
| generated = torch.cat((generated, prompt_tokens), dim=1) | |
| for i in range(entry_length): | |
| outputs = model.gpt(inputs_embeds=generated) | |
| logits = outputs.logits | |
| logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0) | |
| logits = logits.softmax(-1).log() | |
| if scores is None: | |
| scores, next_tokens = logits.topk(beam_size, -1) | |
| generated = generated.expand(beam_size, *generated.shape[1:]) | |
| next_tokens, scores = next_tokens.permute(1, 0), scores.squeeze(0) | |
| if tokens is None: | |
| tokens = next_tokens | |
| else: | |
| tokens = tokens.expand(beam_size, *tokens.shape[1:]) | |
| tokens = torch.cat((tokens, next_tokens), dim=1) | |
| else: | |
| logits[is_stopped] = -float(np.inf) | |
| logits[is_stopped, 0] = 0 | |
| scores_sum = scores[:, None] + logits | |
| seq_lengths[~is_stopped] += 1 | |
| scores_sum_average = scores_sum / seq_lengths[:, None] | |
| scores_sum_average, next_tokens = scores_sum_average.view(-1).topk(beam_size, -1) | |
| next_tokens_source = next_tokens // scores_sum.shape[1] | |
| seq_lengths = seq_lengths[next_tokens_source] | |
| next_tokens = next_tokens % scores_sum.shape[1] | |
| next_tokens = next_tokens.unsqueeze(1) | |
| tokens = tokens[next_tokens_source] | |
| tokens = torch.cat((tokens, next_tokens), dim=1) | |
| generated = generated[next_tokens_source] | |
| scores = scores_sum_average * seq_lengths | |
| is_stopped = is_stopped[next_tokens_source] | |
| next_token_embed = model.gpt.transformer.wte(next_tokens.squeeze()).view(generated.shape[0], 1, -1) | |
| generated = torch.cat((generated, next_token_embed), dim=1) | |
| is_stopped = is_stopped + next_tokens.eq(stop_token_index).squeeze() | |
| if is_stopped.all(): | |
| break | |
| scores = scores / seq_lengths | |
| output_list = tokens.cpu().numpy() | |
| output_texts = [tokenizer.decode(output[:int(length)]) for output, length in zip(output_list, seq_lengths)] | |
| order = scores.argsort(descending=True) | |
| output_texts = [output_texts[i] for i in order] | |
| return output_texts | |
| def generate2( | |
| model, | |
| tokenizer, | |
| tokens=None, | |
| prompt=None, | |
| embed=None, | |
| entry_count=1, | |
| entry_length=67, # maximum number of words | |
| top_p=0.8, | |
| temperature=1., | |
| stop_token: str = '.', | |
| ): | |
| model.eval() | |
| generated_num = 0 | |
| generated_list = [] | |
| stop_token_index = tokenizer.encode(stop_token)[0] | |
| filter_value = -float("Inf") | |
| device = next(model.parameters()).device | |
| with torch.no_grad(): | |
| for entry_idx in trange(entry_count): | |
| if embed is not None: | |
| generated = embed | |
| else: | |
| if tokens is None: | |
| tokens = torch.tensor(tokenizer.encode(prompt)) | |
| tokens = tokens.unsqueeze(0).to(device) | |
| generated = model.gpt.transformer.wte(tokens) | |
| for i in range(entry_length): | |
| outputs = model.gpt(inputs_embeds=generated) | |
| logits = outputs.logits | |
| logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0) | |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) | |
| cumulative_probs = torch.cumsum(nnf.softmax(sorted_logits, dim=-1), dim=-1) | |
| sorted_indices_to_remove = cumulative_probs > top_p | |
| sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[ | |
| ..., :-1 | |
| ].clone() | |
| sorted_indices_to_remove[..., 0] = 0 | |
| indices_to_remove = sorted_indices[sorted_indices_to_remove] | |
| logits[:, indices_to_remove] = filter_value | |
| next_token = torch.argmax(logits, -1).unsqueeze(0) | |
| next_token_embed = model.gpt.transformer.wte(next_token) | |
| if tokens is None: | |
| tokens = next_token | |
| else: | |
| tokens = torch.cat((tokens, next_token), dim=1) | |
| generated = torch.cat((generated, next_token_embed), dim=1) | |
| if stop_token_index == next_token.item(): | |
| break | |
| output_list = list(tokens.squeeze().cpu().numpy()) | |
| output_text = tokenizer.decode(output_list) | |
| generated_list.append(output_text) | |
| return generated_list[0] | |
| is_gpu = False | |
| device = CUDA(0) if is_gpu else "cpu" | |
| clip_model, preprocess = clip.load("ViT-B/16", device=device, jit=False) | |
| from transformers import AutoTokenizer | |
| tokenizer = AutoTokenizer.from_pretrained("imthanhlv/gpt2news") | |
| def inference(img, text, is_translation, prompt=None): | |
| prefix_length = 10 | |
| model = ClipCaptionModel(prefix_length) | |
| model_path = 'sat_019.pt' | |
| model.load_state_dict(torch.load(model_path, map_location=CPU)) | |
| model = model.eval() | |
| device = CUDA(0) if is_gpu else "cpu" | |
| model = model.to(device) | |
| promt = prompt.strip() | |
| if not prompt: | |
| prompt = None | |
| if is_translation: | |
| # encode text | |
| if text is None: | |
| return "No text provided" | |
| text = clip.tokenize([text]).to(device) | |
| with torch.no_grad(): | |
| prefix = clip_model.encode_text(text).to(device, dtype=torch.float32) | |
| prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1) | |
| generated_text_prefix = generate_beam(model, tokenizer, embed=prefix_embed, prompt=prompt)[0] | |
| else: | |
| if img is None: | |
| return "No image" | |
| image = io.imread(img.name) | |
| pil_image = PIL.Image.fromarray(image) | |
| image = preprocess(pil_image).unsqueeze(0).to(device) | |
| with torch.no_grad(): | |
| prefix = clip_model.encode_image(image).to(device, dtype=torch.float32) | |
| prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1) | |
| generated_text_prefix = generate_beam(model, tokenizer, embed=prefix_embed, prompt=prompt)[0] | |
| return generated_text_prefix | |
| title = "CLIP Dual encoder" | |
| description = "You can translate English to Vietnamese or generate Vietnamese caption from image" | |
| examples=[ | |
| ["examples/drug.jpg","", False, "Một bức ảnh về"], | |
| ["examples/harry.jpeg","", False, "Một bức ảnh về"], | |
| ["examples/OldTrafford.jpeg","", False, "Một bức ảnh về"], | |
| ["examples/coffee.jpg","", False, "Một bức ảnh về"], | |
| ["", "What is your name?", True, ""] | |
| ] | |
| inputs = [ | |
| gr.inputs.Image(type="file", label="Image to generate Vietnamese caption", optional=True), | |
| gr.inputs.Textbox(lines=2, placeholder="English sentence for translation"), | |
| gr.inputs.Checkbox(), | |
| gr.inputs.Textbox(lines=1, placeholder="Prompt [Optional]") | |
| ] | |
| gr.Interface( | |
| inference, | |
| inputs, | |
| gr.outputs.Textbox(label="Vietnamese sentence"), | |
| title=title, | |
| description=description, | |
| enable_queue=True, | |
| examples=examples | |
| ).launch(debug=True) | |