Update app.py

app.py

@@ -0,0 +1,171 @@
+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'
+
+cfg = f'./configs/phase2/clipRN50_{reward}.yml'
+
+print("Loading cfg from", cfg)
+
+opt = opts.parse_opt(parse=False, cfg=cfg)
+
+import gdown
+
+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
+
+
+# End below
+
+def generate_image(img, steps=100, seed=42, guidance_scale=6.0):
+
+  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[0]
+
+gr.Interface(
+  generate_image,
+  inputs=[
+      gr.Image(type="pil"),
+      gr.inputs.Slider(1, 100, label='Inference Steps', default=50, step=1),
+      gr.inputs.Slider(0, 2147483647, label='Seed', default=random_seed, step=1),
+      gr.inputs.Slider(1.0, 20.0, label='Guidance Scale - how much the prompt will influence the results', default=6.0, step=0.1),
+  ],
+  outputs=gr.Image(shape=[256,256], type="pil", elem_id="output_image"),
+  css="#output_image{width: 256px}",
+).launch()