|
|
| import gradio as gr |
| import argparse |
| import numpy as np |
| import torch |
| from torch import nn |
|
|
| from data.process_image import load_and_transform_image, get_image_transform |
| from main import SET_GLOBAL_VALUE |
| from model.build_model import create_vat_model |
| from data.process_audio import load_and_transform_audio, get_audio_transform |
| from data.process_video import load_and_transform_video, get_video_transform |
| from data.process_depth import load_and_transform_depth, get_depth_transform |
| from data.process_thermal import load_and_transform_thermal, get_thermal_transform |
| from data.process_text import load_and_transform_text |
| from open_clip import get_tokenizer |
| from open_clip.factory import HF_HUB_PREFIX |
|
|
|
|
|
|
|
|
|
|
|
|
| class LanguageBind(nn.Module): |
| def __init__(self, args, no_temp=False): |
| super(LanguageBind, self).__init__() |
| self.no_temp = no_temp |
| MODEL_DICT = {"ViT-L-14": "laion/CLIP-ViT-L-14-DataComp.XL-s13B-b90K", |
| "ViT-H-14": "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"} |
| args.pretrained = False |
| args.model = MODEL_DICT["ViT-L-14"] |
| args.cache_dir = 'D:/Omni-modal-valdt-audio' |
| args.video_decode_backend = 'decord' |
| |
| args.device = 'cuda:0' |
| device = torch.device(args.device) |
| args.precision = None |
| args.init_temp = 0 |
| args.force_patch_dropout = 0.0 |
| args.add_time_attn = False |
| args.convert_to_lora = True |
| args.lora_r = 2 |
| args.lora_alpha = 16 |
| args.lora_dropout = 0.0 |
| args.num_frames = 8 |
| args.clip_type = 'vl' |
| args.num_mel_bins = 1008 |
| args.target_length = 112 |
| args.audio_sample_rate = 16000 |
| args.audio_mean = 4.5689974 |
| args.audio_std = -4.2677393 |
| args.max_depth = 10 |
| args.image_size = 224 |
| args.rank = 0 |
| SET_GLOBAL_VALUE('PATCH_DROPOUT', args.force_patch_dropout) |
| SET_GLOBAL_VALUE('NUM_FRAMES', args.num_frames) |
| args.clip_type = ['il', 'vl', 'al', 'dl', 'tl'] |
|
|
|
|
| temp_clip_type = args.clip_type |
| self.modality_encoder = {} |
| self.modality_proj = {} |
| self.modality_scale = {} |
| for c in temp_clip_type: |
| args.clip_type = c |
| if c == 'il': |
| args.convert_to_lora = False |
| model = create_vat_model(args) |
| args.convert_to_lora = True |
| elif c == 'vl': |
| args.lora_r = 64 |
| args.add_time_attn = True |
| model = create_vat_model(args) |
| args.add_time_attn = False |
| args.lora_r = 2 |
| elif c == 'al': |
| args.lora_r = 8 |
| model = create_vat_model(args) |
| args.lora_r = 2 |
| else: |
| model = create_vat_model(args) |
| ''' |
| state_dict = torch.load(f'model_zoo/{c}.pt', map_location='cpu') |
| if state_dict.get('state_dict', None) is not None: |
| state_dict = state_dict['state_dict'] |
| if next(iter(state_dict.items()))[0].startswith('module'): |
| state_dict = {k[7:]: v for k, v in state_dict.items()} |
| msg = model.load_state_dict(state_dict, strict=False) |
| print(f'load {c}, {msg}') |
| ''' |
| if c == 'vl': |
| self.modality_encoder['video'] = model.vision_model |
| self.modality_proj['video'] = model.visual_projection |
| self.modality_scale['video'] = model.logit_scale |
| elif c == 'al': |
| self.modality_encoder['audio'] = model.vision_model |
| self.modality_proj['audio'] = model.visual_projection |
| self.modality_scale['audio'] = model.logit_scale |
| elif c == 'dl': |
| self.modality_encoder['depth'] = model.vision_model |
| self.modality_proj['depth'] = model.visual_projection |
| self.modality_scale['depth'] = model.logit_scale |
| elif c == 'tl': |
| self.modality_encoder['thermal'] = model.vision_model |
| self.modality_proj['thermal'] = model.visual_projection |
| self.modality_scale['thermal'] = model.logit_scale |
| elif c == 'il': |
| self.modality_encoder['image'] = model.vision_model |
| self.modality_proj['image'] = model.visual_projection |
| self.modality_scale['image'] = model.logit_scale |
| else: |
| raise NameError(f'No clip_type of {c}') |
| self.modality_encoder['language'] = model.text_model |
| self.modality_proj['language'] = model.text_projection |
|
|
| self.modality_encoder = nn.ModuleDict(self.modality_encoder) |
| self.modality_proj = nn.ModuleDict(self.modality_proj) |
|
|
| def forward(self, inputs): |
| outputs = {} |
| for key, value in inputs.items(): |
| value = self.modality_encoder[key](**value)[1] |
| value = self.modality_proj[key](value) |
| value = value / value.norm(p=2, dim=-1, keepdim=True) |
| if not self.no_temp: |
| if key != 'language': |
| value = value * self.modality_scale[key].exp() |
| outputs[key] = value |
| return outputs |
|
|
|
|
|
|
| def stack_dict(x, device): |
| if len(x) == 0: |
| return None |
| out_dict = {} |
| keys = list(x[0].keys()) |
| for key in keys: |
| out_dict[key] = torch.stack([i[key] for i in x]).to(device) |
| return out_dict |
