diff --git a/main.py b/main.py deleted file mode 100644 index adcb976..0000000 --- a/main.py +++ /dev/null @@ -1,166 +0,0 @@ -import clip -import torch -import skimage.io as io -import PIL.Image -import numpy as np -import torch.nn.functional as nnf -from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup -from tqdm import tqdm, trange -from clipcap_model import MLP, ClipCaptionModel, ClipCaptionPrefix - -is_gpu = False -device = CUDA(0) if is_gpu else "cpu" -clip_model, preprocess = clip.load("ViT-B/32", device=device, jit=False) -tokenizer = GPT2Tokenizer.from_pretrained("gpt2") -CPU = torch.device('cpu') - - -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] - -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 - 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) - 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 - -prefix_length = 10 - -model = ClipCaptionModel(prefix_length) -model_path = '/Users/zilliz/git/image_captioning/git/clipcap/weights/coco_weights.pt' -model.load_state_dict(torch.load(model_path, map_location=CPU)) -model = model.eval() - -use_beam_search = False #@param {type:"boolean"} -use_beam_search = True #@param {type:"boolean"} - -UPLOADED_FILE = 'einstein.jpg' -image = io.imread(UPLOADED_FILE) -pil_image = PIL.Image.fromarray(image) - -image = preprocess(pil_image).unsqueeze(0).to(device) -with torch.no_grad(): - # if type(model) is ClipCaptionE2E: - # prefix_embed = model.forward_image(image) - # else: - prefix = clip_model.encode_image(image).to(device, dtype=torch.float32) - prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1) -if use_beam_search: - generated_text_prefix = generate_beam(model, tokenizer, embed=prefix_embed)[0] -else: - generated_text_prefix = generate2(model, tokenizer, embed=prefix_embed) - -print(generated_text_prefix) - - diff --git a/models/.utils.py.swp b/models/.utils.py.swp deleted file mode 100644 index ac298e3..0000000 Binary files a/models/.utils.py.swp and /dev/null differ diff --git a/models/clipcap.py b/models/clipcap.py index ef15bb3..cb97a50 100644 --- a/models/clipcap.py +++ b/models/clipcap.py @@ -3,7 +3,6 @@ import torch.nn.functional as nnf #@title Imports from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup -import clip import os from typing import Tuple, List, Union, Optional from torch import nn diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..ade6d9d --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +transformers +torch +towhee>=0.7 +towhee.models>=0.7