fix

app.py

@@ -50,20 +50,17 @@ device_type = 'cuda' if use_cuda else 'cpu'
 ## Load model
 
 ### Captioning 
-config = 'configs/image/ePALM_caption.yaml'
-# config = yaml.load(open(config, 'r'), Loader=yaml.Loader)
+config = 'configs/audio/ePALM_audio_caption.yaml'
 config = yaml.load(open(config, 'r'))
 
 text_model = 'facebook/opt-2.7b' 
 vision_model_name = 'vit_base_patch16_224'
 
-# text_model = 'facebook/opt-6.7b' 
-# vision_model_name = 'vit_large_patch16_224'
 
 start_layer_idx = 19
 end_layer_idx = 31
 low_cpu = True 
-model = ePALM(opt_model_name=text_model, 
+MODEL = ePALM(opt_model_name=text_model, 
                vision_model_name=vision_model_name, 
                use_vis_prefix=True, 
                start_layer_idx=start_layer_idx, 
@@ -73,62 +70,20 @@ model = ePALM(opt_model_name=text_model,
                low_cpu=low_cpu
 )
 print("Model Built")
-model.to(device)
+MODEL.to(device)
 
 checkpoint_path = 'checkpoints/float32/ePALM_caption/checkpoint_best.pth'
-# checkpoint_path = '/data/mshukor/logs/eplam/models/accelerate/ePALM_pt_L_acc_caption/checkpoint_best.pth'
 checkpoint = torch.load(checkpoint_path, map_location='cpu') 
 state_dict = checkpoint['model']
-msg = model.load_state_dict(state_dict,strict=False)  
-
-model.bfloat16()
-
-# ###### VQA
-# config = 'configs/image/ePALM_vqa.yaml'
-# config = yaml.load(open(config, 'r'))
-
-# start_layer_idx = 19
-# end_layer_idx = 31
-# low_cpu = True 
-# model_vqa = ePALM(opt_model_name=text_model, 
-#                vision_model_name=vision_model_name, 
-#                use_vis_prefix=True, 
-#                start_layer_idx=start_layer_idx, 
-#                end_layer_idx=end_layer_idx, 
-#                return_hidden_state_vision=True, 
-#                config=config,
-#                low_cpu=low_cpu
-# )
-# print("Model Built")
-# model_vqa.to(device)
-
-
-checkpoint_path = 'checkpoints/float32/ePALM_vqa/checkpoint_best.pth'
-checkpoint = torch.load(checkpoint_path, map_location='cpu') 
-state_dict_vqa = checkpoint['model']
-# msg = model_vqa.load_state_dict(state_dict,strict=False)  
-
+msg = MODEL.load_state_dict(state_dict,strict=False)  
 
-# model_vqa.bfloat16()
+MODEL.bfloat16()
 
 
 
-# Video Captioning
-checkpoint_path = 'checkpoints/float32/ePALM_video_caption_msrvtt/checkpoint_best.pth'
-# checkpoint_path = '/data/mshukor/logs/eplam/models/accelerate/ePALM_pt_L_acc_caption/checkpoint_best.pth'
-checkpoint = torch.load(checkpoint_path, map_location='cpu') 
-state_dict_video_caption = checkpoint['model']
-
-# Video QA
-checkpoint_path = 'checkpoints/float32/ePALM_video_qa_msrvtt/checkpoint_best.pth'
-# checkpoint_path = '/data/mshukor/logs/eplam/models/accelerate/ePALM_pt_L_acc_caption/checkpoint_best.pth'
-checkpoint = torch.load(checkpoint_path, map_location='cpu') 
-state_dict_video_qa = checkpoint['model']
-
 
 # Audio Captioning
 checkpoint_path = 'checkpoints/float32/ePALM_audio_caption/checkpoint_best.pth'
-# checkpoint_path = '/data/mshukor/logs/eplam/models/accelerate/ePALM_pt_L_acc_caption/checkpoint_best.pth'
 checkpoint = torch.load(checkpoint_path, map_location='cpu') 
 state_dict_audio_caption = checkpoint['model']
 
@@ -146,33 +101,8 @@ special_tokens_dict = {'additional_special_tokens': [special_answer_token]}
 tokenizer.add_special_tokens(special_tokens_dict)
 
 
-image_size = 224
-normalize = transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
-
-transform = transforms.Compose([
-            transforms.Resize((image_size,image_size),interpolation=Image.BICUBIC),
-            transforms.ToTensor(),
-            normalize,
-            ])  
 
-type_transform = transforms.Lambda(lambda x: x.float().div(255.0))
-test_transform = transforms.Compose([
-    transforms.Resize((image_size,image_size),interpolation=Image.BICUBIC),
-    type_transform,
-    normalize,
-    ])  
-from dataset.video_utils import VIDEO_READER_FUNCS
-video_reader = VIDEO_READER_FUNCS['decord']
-
-def read_video(path, num_frames=16):
-        
 
-    frames, frame_indices, video_duration = video_reader(
-        path, num_frames, 'rand', max_num_frames=-1
-    )
-    video = test_transform(frames)
-
-    return video
 
 def read_audio(path):
         
@@ -237,37 +167,18 @@ max_length=30
 
 
 
-def inference(image, audio, video, task_type, instruction):
+def inference(image, task_type):
 
-    if task_type == 'Image Captioning':
-        text = ['']  
-        text_input = tokenizer(text, padding='longest', return_tensors="pt").to(device) 
-    elif task_type == 'Video Captioning':
-        text = ['']  
-        text_input = tokenizer(text, padding='longest', return_tensors="pt").to(device) 
-        model = model.load_state_dict(state_dict_video_caption,strict=False)  
-    elif task_type == 'Audio Captioning':
+    if task_type == 'Audio Captioning':
         text = ['']  
         text_input = tokenizer(text, padding='longest', return_tensors="pt").to(device) 
-        model = model.load_state_dict(state_dict_audio_caption,strict=False)  
-    elif task_type == 'Visual Question Answering':
-        question = instruction+'?'+special_answer_token
-        text_input = tokenizer(question, padding='longest', return_tensors="pt").to(device) 
-        model = model.load_state_dict(state_dict_vqa,strict=False)  
-    elif task_type == 'Visual Question Answering':
-        question = instruction+'?'+special_answer_token
-        text_input = tokenizer(question, padding='longest', return_tensors="pt").to(device) 
-        model = model.load_state_dict(state_dict_video_qa,strict=False)  
+        model = MODEL
     else:
         raise NotImplemented
 
-    if "Video" in task_type:
-        image = read_video(image)
-    elif "Audio" in task_type:
-        image = read_audio(image)
-    else:
-        image = transform(image)
-        image = image.to(device,non_blocking=True).unsqueeze(0)
+
+    image = read_audio(image)
+
 
 
 
@@ -290,25 +201,15 @@ def inference(image, audio, video, task_type, instruction):
     return response
 
 
-inputs = [gr.inputs.Image(type='pil'), gr.Audio(source="upload", type="filepath"), gr.Video(source="upload", type="filepath"), gr.inputs.Radio(choices=['Image Captioning', 'Video Captioning', 'Audio Captioning', "Visual Question Answering", "Visual Grounding", "General", "General Video"], type="value", default="Image Captioning", label="Task"), gr.inputs.Textbox(lines=1, label="Instruction")]
+inputs = [gr.Audio(source="upload", type="filepath"),  gr.inputs.Radio(choices=['Audio Captioning'], type="value", default="Image Captioning", label="Task")]
 outputs = ['text']
 examples = [
-    ['examples/images/soccer.jpg', None, None, 'Image Captioning', None],
-    ['examples/images/ski.jpg', None, None, 'Visual Question Answering', 'what does the woman wearing black do?'],
-    ['examples/images/banana.jpg', None, None, 'Image Captioning', None],
-    ['examples/images/skateboard.jpg', None, None, 'Visual Question Answering', 'what is on top of the skateboard?'],
-    ['examples/images/baseball.jpg', None, None, 'Image Captioning', None],
-    [None, None, 'examples/videos/video7014.mp4', 'Video Captioning', None], 
-    [None, None, 'examples/videos/video7017.mp4', 'Video Captioning', None], 
-    [None, None, 'examples/videos/video7019.mp4', 'Video Captioning', None], 
-    [None, None, 'examples/videos/video7021.mp4', 'Video Captioning', None], 
-    [None, None, 'examples/videos/video7021.mp4', 'Video Captioning', None], 
-    [None, 'examples/audios/6cS0FsUM-cQ.wav', None, 'Audio Captioning', None],
-    [None, 'examples/audios/AJtNitYMa1I.wav', None, 'Audio Captioning', None],
+    ['examples/audios/6cS0FsUM-cQ.wav', 'Audio Captioning', None],
+    ['examples/audios/AJtNitYMa1I.wav', 'Audio Captioning', None],
 ]
 
-title = "eP-ALM"
-description = "Gradio Demo for eP-ALM: "
+title = "eP-ALM for Audio-Text tasks"
+description = "Gradio Demo for eP-ALM. For this demo, we use 2.7B OPT. As the model runs on CPUs and float16 mixed precision is not supported on CPUs, the generation can take up to 2 mins."
 article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2303.11403' target='_blank'>Paper</a> | <a href='https://github.com/mshukor/eP-ALM' target='_blank'>Github Repo</a></p>"
 
 io = gr.Interface(fn=inference, inputs=inputs, outputs=outputs,