Enable all the code

app.py

@@ -1,181 +1,181 @@
-# 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 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
 
-# from diffusers import LDMTextToImagePipeline
-# # import PIL.Image
+from diffusers import LDMTextToImagePipeline
+# import PIL.Image
 import random
-# import os
+import os
 
 
-# # 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)
+# 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'
+device = 'cpu' #@param ["cuda", "cpu"] {allow-input: true}
+reward = 'clips_grammar'
 
-# cfg = f'./configs/phase2/clipRN50_{reward}.yml'
+cfg = f'./configs/phase2/clipRN50_{reward}.yml'
 
-# print("Loading cfg from", cfg)
+print("Loading cfg from", cfg)
 
-# opt = opts.parse_opt(parse=False, cfg=cfg)
+opt = opts.parse_opt(parse=False, cfg=cfg)
 
-# import gdown
+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/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/")
+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())
+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)
+ix_to_word = dict_json['ix_to_word']
+vocab_size = len(ix_to_word)
+print('vocab size:', vocab_size)
 
-# seq_length = 1
+seq_length = 1
 
-# opt.vocab_size = vocab_size
-# opt.seq_length = seq_length
+opt.vocab_size = vocab_size
+opt.seq_length = seq_length
 
-# opt.batch_size = 1
-# opt.vocab = ix_to_word
+opt.batch_size = 1
+opt.vocab = ix_to_word
 
-# model = models.setup(opt)
-# del opt.vocab
+model = models.setup(opt)
+del opt.vocab
 
-# ckpt_path = opt.checkpoint_path + '-last.ckpt'
+ckpt_path = opt.checkpoint_path + '-last.ckpt'
 
-# print("Loading checkpoint from", ckpt_path)
-# raw_state_dict = torch.load(
-#     ckpt_path,
-#     map_location=device)
+print("Loading checkpoint from", ckpt_path)
+raw_state_dict = torch.load(
+    ckpt_path,
+    map_location=device)
 
-# strict = True
+strict = True
 
-# state_dict = raw_state_dict['state_dict']
+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)
+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();
+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
+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)
+clip_model, clip_transform = clip.load("RN50", jit=False, device=device)
 
-# preprocess = Compose([
-#     Resize((448, 448), interpolation=Image.BICUBIC),
-#     CenterCrop((448, 448)),
-#     ToTensor()
-# ])
+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)
+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
+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
-# print('Loading the model: CompVis/ldm-text2im-large-256')
-# ldm_pipeline = LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256")
+# End below
+print('Loading the model: CompVis/ldm-text2im-large-256')
+ldm_pipeline = LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256")
 
-# def generate_image_from_text(prompt, steps=100, seed=42, guidance_scale=6.0):
-#     print('RUN: generate_image_from_text')
-#     torch.cuda.empty_cache()
-#     generator = torch.manual_seed(seed)
-#     images = ldm_pipeline([prompt], generator=generator, num_inference_steps=steps, eta=0.3, guidance_scale=guidance_scale)["sample"]
-#     return images[0]
-
-# def generate_text_from_image(img):
-#     print('RUN: generate_text_from_image')
-#     with torch.no_grad():
-#         image = preprocess(img)
-#         image = torch.tensor(np.stack([image])).to(device)
-#         image -= image_mean
-#         image /= image_std
+def generate_image_from_text(prompt, steps=100, seed=42, guidance_scale=6.0):
+    print('RUN: generate_image_from_text')
+    torch.cuda.empty_cache()
+    generator = torch.manual_seed(seed)
+    images = ldm_pipeline([prompt], generator=generator, num_inference_steps=steps, eta=0.3, guidance_scale=guidance_scale)["sample"]
+    return images[0]
+
+def generate_text_from_image(img):
+    print('RUN: generate_text_from_image')
+    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]
+        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
+        att_feat = tmp_att
+        fc_feat = tmp_fc
       
-#     # Inference configurations
-#     eval_kwargs = {}
-#     eval_kwargs.update(vars(opt))
+    # 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)
+    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)
+    # 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
+    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
+        # 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)
+        sents = utils.decode_sequence(model.vocab, seq)
 
-#         return sents[0]
+        return sents[0]
 
 
-# def generate_drawing_from_image(img, steps=100, seed=42, guidance_scale=6.0):
-#     print('RUN: generate_drawing_from_image')
-#     caption = generate_text_from_image(img)
-#     gen_image = generate_image_from_text(caption, steps=steps, seed=seed, guidance_scale=guidance_scale)
-#     return gen_image
+def generate_drawing_from_image(img, steps=100, seed=42, guidance_scale=6.0):
+    print('RUN: generate_drawing_from_image')
+    caption = generate_text_from_image(img)
+    gen_image = generate_image_from_text(caption, steps=steps, seed=seed, guidance_scale=guidance_scale)
+    return gen_image
 
 
 random_seed = random.randint(0, 2147483647)
@@ -183,10 +183,9 @@ random_seed = random.randint(0, 2147483647)
 def test_fn(**kwargs):
     return None
 
-
 gr.Interface(
-#   generate_drawing_from_image,
-    test_fn,
+    generate_drawing_from_image,
+    # test_fn,
     inputs=[
         gr.Image(type="pil"),
         gr.inputs.Slider(1, 100, label='Inference Steps', default=50, step=1),