Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| import json | |
| import captioning.utils.opts as opts | |
| import captioning.models as models | |
| import captioning.utils.misc as utils | |
| import pytorch_lightning as pl | |
| import gradio as gr | |
| # Checkpoint class | |
| class ModelCheckpoint(pl.callbacks.ModelCheckpoint): | |
| def on_keyboard_interrupt(self, trainer, pl_module): | |
| # Save model when keyboard interrupt | |
| filepath = os.path.join(self.dirpath, self.prefix + 'interrupt.ckpt') | |
| self._save_model(filepath) | |
| device = 'cpu' #@param ["cuda", "cpu"] {allow-input: true} | |
| reward = 'clips_grammar' #@param ["mle", "cider", "clips", "cider_clips", "clips_grammar"] {allow-input: true} | |
| if reward == 'mle': | |
| cfg = f'./configs/phase1/clipRN50_{reward}.yml' | |
| else: | |
| cfg = f'./configs/phase2/clipRN50_{reward}.yml' | |
| print("Loading cfg from", cfg) | |
| opt = opts.parse_opt(parse=False, cfg=cfg) | |
| import gdown | |
| if reward == "mle": | |
| url = "https://drive.google.com/drive/folders/1hfHWDn5iXsdjB63E5zdZBAoRLWHQC3LD" | |
| elif reward == "cider": | |
| url = "https://drive.google.com/drive/folders/1MnSmCd8HFnBvQq_4K-q4vsVkzEw0OIOs" | |
| elif reward == "clips": | |
| url = "https://drive.google.com/drive/folders/1toceycN-qilHsbYjKalBLtHJck1acQVe" | |
| elif reward == "cider_clips": | |
| url = "https://drive.google.com/drive/folders/1toceycN-qilHsbYjKalBLtHJck1acQVe" | |
| elif reward == "clips_grammar": | |
| url = "https://drive.google.com/drive/folders/1nSX9aS7pPK4-OTHYtsUD_uEkwIQVIV7W" | |
| gdown.download_folder(url, quiet=True, use_cookies=False, output="save/") | |
| url = "https://drive.google.com/uc?id=1HNRE1MYO9wxmtMHLC8zURraoNFu157Dp" | |
| gdown.download(url, quiet=True, use_cookies=False, output="data/") | |
| dict_json = json.load(open('./data/cocotalk.json')) | |
| print(dict_json.keys()) | |
| ix_to_word = dict_json['ix_to_word'] | |
| vocab_size = len(ix_to_word) | |
| print('vocab size:', vocab_size) | |
| seq_length = 1 | |
| opt.vocab_size = vocab_size | |
| opt.seq_length = seq_length | |
| opt.batch_size = 1 | |
| opt.vocab = ix_to_word | |
| # opt.use_grammar = False | |
| model = models.setup(opt) | |
| del opt.vocab | |
| ckpt_path = opt.checkpoint_path + '-last.ckpt' | |
| print("Loading checkpoint from", ckpt_path) | |
| raw_state_dict = torch.load( | |
| ckpt_path, | |
| map_location=device) | |
| strict = True | |
| state_dict = raw_state_dict['state_dict'] | |
| if '_vocab' in state_dict: | |
| model.vocab = utils.deserialize(state_dict['_vocab']) | |
| del state_dict['_vocab'] | |
| elif strict: | |
| raise KeyError | |
| if '_opt' in state_dict: | |
| saved_model_opt = utils.deserialize(state_dict['_opt']) | |
| del state_dict['_opt'] | |
| # Make sure the saved opt is compatible with the curren topt | |
| need_be_same = ["caption_model", | |
| "rnn_type", "rnn_size", "num_layers"] | |
| for checkme in need_be_same: | |
| if getattr(saved_model_opt, checkme) in ['updown', 'topdown'] and \ | |
| getattr(opt, checkme) in ['updown', 'topdown']: | |
| continue | |
| assert getattr(saved_model_opt, checkme) == getattr( | |
| opt, checkme), "Command line argument and saved model disagree on '%s' " % checkme | |
| elif strict: | |
| raise KeyError | |
| res = model.load_state_dict(state_dict, strict) | |
| print(res) | |
| model = model.to(device) | |
| model.eval(); | |
| import clip | |
| from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize | |
| from PIL import Image | |
| from timm.models.vision_transformer import resize_pos_embed | |
| clip_model, clip_transform = clip.load("RN50", jit=False, device=device) | |
| preprocess = Compose([ | |
| Resize((448, 448), interpolation=Image.BICUBIC), | |
| CenterCrop((448, 448)), | |
| ToTensor() | |
| ]) | |
| image_mean = torch.Tensor([0.48145466, 0.4578275, 0.40821073]).to(device).reshape(3, 1, 1) | |
| image_std = torch.Tensor([0.26862954, 0.26130258, 0.27577711]).to(device).reshape(3, 1, 1) | |
| num_patches = 196 #600 * 1000 // 32 // 32 | |
| pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, clip_model.visual.attnpool.positional_embedding.shape[-1], device=device),) | |
| pos_embed.weight = resize_pos_embed(clip_model.visual.attnpool.positional_embedding.unsqueeze(0), pos_embed) | |
| clip_model.visual.attnpool.positional_embedding = pos_embed | |
| def inference(img): | |
| with torch.no_grad(): | |
| image = preprocess(img) | |
| image = torch.tensor(np.stack([image])).to(device) | |
| image -= image_mean | |
| image /= image_std | |
| tmp_att, tmp_fc = clip_model.encode_image(image) | |
| tmp_att = tmp_att[0].permute(1, 2, 0) | |
| tmp_fc = tmp_fc[0] | |
| att_feat = tmp_att | |
| fc_feat = tmp_fc | |
| # Inference configurations | |
| eval_kwargs = {} | |
| eval_kwargs.update(vars(opt)) | |
| verbose = eval_kwargs.get('verbose', True) | |
| verbose_beam = eval_kwargs.get('verbose_beam', 0) | |
| verbose_loss = eval_kwargs.get('verbose_loss', 1) | |
| # dataset = eval_kwargs.get('dataset', 'coco') | |
| beam_size = eval_kwargs.get('beam_size', 1) | |
| sample_n = eval_kwargs.get('sample_n', 1) | |
| remove_bad_endings = eval_kwargs.get('remove_bad_endings', 0) | |
| with torch.no_grad(): | |
| fc_feats = torch.zeros((1,0)).to(device) | |
| att_feats = att_feat.view(1, 196, 2048).float().to(device) | |
| att_masks = None | |
| # forward the model to also get generated samples for each image | |
| # Only leave one feature for each image, in case duplicate sample | |
| tmp_eval_kwargs = eval_kwargs.copy() | |
| tmp_eval_kwargs.update({'sample_n': 1}) | |
| seq, seq_logprobs = model( | |
| fc_feats, att_feats, att_masks, opt=tmp_eval_kwargs, mode='sample') | |
| seq = seq.data | |
| sents = utils.decode_sequence(model.vocab, seq) | |
| return sents | |
| demo = gr.Blocks() | |
| with demo: | |
| gr.Markdown( | |
| """ | |
| # Demo for CLIP-Caption-Reward | |
| """) | |
| inp = gr.Image(type="pil") | |
| out = gr.Textbox() | |
| image_button = gr.Button("Run") | |
| image_button.click(fn=inference, | |
| inputs=inp, | |
| outputs=out) | |
| demo.launch() |