Upload 5 files

model/AnomalyGPT_models.py

@@ -0,0 +1,73 @@
+import torch
+from torch import nn
+import numpy as np
+# from datas.dataset_3d import  *
+from torch.nn import functional as F
+
+
+class Normalize(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        return torch.nn.functional.normalize(x, dim=self.dim, p=2)
+
+    
+class LinearLayer(nn.Module):
+    def __init__(self, dim_in, dim_out, k):
+        super(LinearLayer, self).__init__()
+        self.fc = nn.ModuleList([nn.Linear(dim_in, dim_out) for i in range(k)])
+
+    def forward(self, tokens):
+        for i in range(len(tokens)):
+            if len(tokens[i].shape) == 3:
+                tokens[i] = tokens[i].transpose(0,1)
+                tokens[i] = self.fc[i](tokens[i][:, 1:, :])
+            else:
+                B, C, H, W = tokens[i].shape
+                tokens[i] = self.fc[i](tokens[i].view(B, C, -1).permute(0, 2, 1).contiguous())
+        return tokens
+    
+class PromptLearner(nn.Module):
+    def __init__(self, dim_in, dim_out) -> None:
+        super().__init__()
+        self.meta_net = nn.Sequential(
+            nn.Conv2d(dim_in, dim_in * 4, kernel_size=3, padding=1),
+            # nn.BatchNorm2d(dim_in * 4),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2), # 112 * 112
+
+            nn.Conv2d(dim_in * 4, dim_in * 16, kernel_size=3, padding=1),
+            # nn.BatchNorm2d(dim_in * 16),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2), # 56 * 56
+
+            nn.Conv2d(dim_in * 16, dim_in * 64, kernel_size=3, padding=1),
+            # nn.BatchNorm2d(dim_in * 64),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2), # 28 * 28
+
+            nn.Conv2d(dim_in * 64, dim_in * 256, kernel_size=3, padding=1),
+            # nn.BatchNorm2d(dim_in * 256),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2), # 14 * 14
+
+            nn.Conv2d(dim_in * 256, dim_in * 1024, kernel_size=3, padding=1),
+            # nn.BatchNorm2d(dim_in * 1024),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2), # 7 * 7
+
+            nn.Conv2d(dim_in * 1024, dim_out, kernel_size=5, padding=0),
+            # nn.BatchNorm2d(dim_out),
+            # nn.ReLU(inplace=True),
+        )
+        self.base_prompts = nn.Parameter(torch.randn((9, dim_out)),requires_grad=True)
+
+    def forward(self, input):
+        B,C,H,W = input.shape
+        img_prompts = self.meta_net(input)
+        # print(input.shape, img_prompts.shape)
+        img_prompts = img_prompts.reshape(B,4096,9).transpose(-2,-1)
+        output = torch.cat([self.base_prompts.expand(B,-1,-1), img_prompts], dim=1)
+        return output

model/__init__.py

@@ -0,0 +1,11 @@
+from .agent import DeepSpeedAgent
+from .openllama import OpenLLAMAPEFTModel
+# from .openllama_CLIP import OpenLLAMAPEFTModel_CLIP
+from .ImageBind import models
+
+def load_model(args):
+    agent_name = args['models'][args['model']]['agent_name']
+    model_name = args['models'][args['model']]['model_name']
+    model = globals()[model_name](**args)
+    agent = globals()[agent_name](model, args)
+    return agent

model/agent.py

@@ -0,0 +1,81 @@
+from header import *
+
+class DeepSpeedAgent:
+    
+    def __init__(self, model, args):
+        super(DeepSpeedAgent, self).__init__()
+        self.args = args
+        self.model = model
+        self.load_stage_1_parameters(args["delta_ckpt_path"])
+
+
+
+        for name, param in self.model.named_parameters():
+            param.requires_grad = False
+
+        for name, param in self.model.image_decoder.named_parameters():
+            param.requires_grad = True
+
+        for name, param in self.model.prompt_learner.named_parameters():
+            param.requires_grad = True
+
+
+
+
+        # load config parameters of deepspeed
+        ds_params = json.load(open(self.args['ds_config_path']))
+        ds_params['scheduler']['params']['total_num_steps'] = self.args['total_steps']
+        ds_params['scheduler']['params']['warmup_num_steps'] = max(10, int(self.args['total_steps'] * self.args['warmup_rate']))
+        self.ds_engine, self.optimizer, _ , _ = deepspeed.initialize(
+            model=self.model, 
+            model_parameters=self.model.parameters(),
+            config_params=ds_params, 
+            dist_init_required=True,
+            args=types.SimpleNamespace(**args)
+        )
+
+    @torch.no_grad()
+    def predict(self, batch):
+        self.model.eval()
+        string = self.model.generate_one_sample(batch)
+        return string
+
+    def train_model(self, batch, current_step=0, pbar=None):
+        self.ds_engine.module.train()
+        loss, mle_acc = self.ds_engine(batch)
+
+        self.ds_engine.backward(loss)
+        self.ds_engine.step()
+        pbar.set_description(f'[!] loss: {round(loss.item(), 4)}; token_acc: {round(mle_acc*100, 2)}')
+        pbar.update(1)
+        if self.args['local_rank'] == 0 and self.args['log_path'] and current_step % self.args['logging_step'] == 0:
+            elapsed = pbar.format_dict['elapsed']
+            rate = pbar.format_dict['rate']
+            remaining = (pbar.total - pbar.n) / rate if rate and pbar.total else 0
+            remaining = str(datetime.timedelta(seconds=remaining))
+            logging.info(f'[!] progress: {round(pbar.n/pbar.total, 5)}; remaining time: {remaining}; loss: {round(loss.item(), 4)}; token_acc: {round(mle_acc*100, 2)}')
+            
+        mle_acc *= 100
+        return mle_acc
+
+    def save_model(self, path, current_step):
+        # only save trainable model parameters
+        param_grad_dic = {
+            k: v.requires_grad for (k, v) in self.ds_engine.module.named_parameters()
+        }
+        state_dict = self.ds_engine.module.state_dict()
+        checkpoint = OrderedDict()
+        for k, v in self.ds_engine.module.named_parameters():
+            if v.requires_grad:
+                print(k)
+                checkpoint[k] = v
+        torch.save(checkpoint, f'{path}/pytorch_model.pt')
+        # save tokenizer
+        self.model.llama_tokenizer.save_pretrained(path)
+        # save configuration
+        self.model.llama_model.config.save_pretrained(path)
+        print(f'[!] save model into {path}')
+
+    def load_stage_1_parameters(self, path):
+        delta_ckpt = torch.load(path, map_location=torch.device('cpu'))
+        self.model.load_state_dict(delta_ckpt, strict=False)

model/modeling_llama.py

@@ -0,0 +1,755 @@
+# This script is based on https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
+
+""" PyTorch LLaMA model."""
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
+from transformers.models.llama.configuration_llama import LlamaConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "LlamaConfig"
+
+
+# Copied from transformers.models.bart.modeling_bart._make_causal_mask
+def _make_causal_mask(
+    input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
+):
+    """
+    Make causal mask used for bi-directional self-attention.
+    """
+    bsz, tgt_len = input_ids_shape
+    mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)
+    mask_cond = torch.arange(mask.size(-1), device=device)
+    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
+    mask = mask.to(dtype)
+
+    if past_key_values_length > 0:
+        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
+    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
+
+
+# Copied from transformers.models.bart.modeling_bart._expand_mask
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
+    """
+    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+    """
+    bsz, src_len = mask.size()
+    tgt_len = tgt_len if tgt_len is not None else src_len
+
+    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+
+    inverted_mask = 1.0 - expanded_mask
+
+    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
+
+
+class LlamaRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        LlamaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+
+        # convert into half-precision if necessary
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            hidden_states = hidden_states.to(self.weight.dtype)
+
+        return self.weight * hidden_states
+
+
+class LlamaRotaryEmbedding(torch.nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+        # Build here to make `torch.jit.trace` work.
+        self.max_seq_len_cached = max_position_embeddings
+        t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
+        self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
+        if seq_len > self.max_seq_len_cached:
+            self.max_seq_len_cached = seq_len
+            t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            # Different from paper, but it uses a different permutation in order to obtain the same calculation
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+            self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
+            self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
+        return (
+            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
+            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
+        )
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
+    gather_indices = position_ids[:, None, :, None]  # [bs, 1, seq_len, 1]
+    gather_indices = gather_indices.repeat(1, cos.shape[1], 1, cos.shape[3])
+    cos = torch.gather(cos.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices)
+    sin = torch.gather(sin.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class LlamaMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+    ):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.act_fn = ACT2FN[hidden_act]
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LlamaAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: LlamaConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.max_position_embeddings = config.max_position_embeddings
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+        self.rotary_emb = LlamaRotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+        # [bsz, nh, t, hd]
+
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+            attn_weights = torch.max(attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min))
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+class LlamaDecoderLayer(nn.Module):
+    def __init__(self, config: LlamaConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = LlamaAttention(config=config)
+        self.mlp = LlamaMLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+        )
+        self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+        """
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+
+LLAMA_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`LlamaConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",
+    LLAMA_START_DOCSTRING,
+)
+class LlamaPreTrainedModel(PreTrainedModel):
+    config_class = LlamaConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LlamaDecoderLayer"]
+    _keys_to_ignore_on_load_unexpected = [r"decoder\.version"]
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, LlamaModel):
+            module.gradient_checkpointing = value
+
+
+LLAMA_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+            `past_key_values`).
+
+            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+            information on the default strategy.
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.n_positions - 1]`.
+
+            [What are position IDs?](../glossary#position-ids)
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+            `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
+            `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
+
+            Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
+
+            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
+            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
+            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",
+    LLAMA_START_DOCSTRING,
+)
+class LlamaModel(LlamaPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]
+
+    Args:
+        config: LlamaConfig
+    """
+
+    def __init__(self, config: LlamaConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
+    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
+        # create causal mask
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        combined_attention_mask = None
+        if input_shape[-1] > 1:
+            combined_attention_mask = _make_causal_mask(
+                input_shape,
+                inputs_embeds.dtype,
+                device=inputs_embeds.device,
+                past_key_values_length=past_key_values_length,
+            )
+
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
+                inputs_embeds.device
+            )
+            combined_attention_mask = (
+                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+            )
+
+        return combined_attention_mask
+
+    @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        query_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        else:
+            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        if query_embeds is not None:
+            inputs_embeds = torch.cat([query_embeds, inputs_embeds], dim=1)
+            batch_size, seq_length, _ = inputs_embeds.shape
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+            )
+            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+        else:
+            position_ids = position_ids.view(-1, seq_length).long()
+
+        # embed positions
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
+            )
+        attention_mask = self._prepare_decoder_attention_mask(
+            attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+        )
+
+        hidden_states = inputs_embeds
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = () if use_cache else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            past_key_value = past_key_values[idx] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs, output_attentions, None)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(decoder_layer),
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    None,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class LlamaForCausalLM(LlamaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LlamaModel(config)
+
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        query_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, LlamaForCausalLM
+
+        >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+        >>> prompt = "Hey, are you consciours? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you."
+        ```"""
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            query_embeds=query_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, query_embeds=None, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+                query_embeds = None
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "query_embeds": query_embeds,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
+        return reordered_past
+

model/openllama.py

@@ -0,0 +1,755 @@
+from header import *
+import torch.nn.functional as F
+from .ImageBind import *
+from .ImageBind import data
+from .modeling_llama import LlamaForCausalLM
+from .AnomalyGPT_models import LinearLayer, PromptLearner
+from transformers import StoppingCriteria, StoppingCriteriaList
+from utils.loss import FocalLoss, BinaryDiceLoss
+import kornia as K
+
+import torch
+from torch.nn.utils import rnn
+from transformers import AutoConfig, AutoModelForCausalLM
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch, infer_auto_device_map
+
+CLASS_NAMES = ['bottle', 'cable', 'capsule', 'carpet', 'grid', 'hazelnut', 'leather', 'metal nut', 'pill', 'screw', 'tile', 'toothbrush', 'transistor', 'wood', 'zipper', 'object',
+               'candle', 'cashew', 'chewinggum', 'fryum', 'macaroni', 'pcb', 'pipe fryum']
+
+prompt_normal = ['{}', 'flawless {}', 'perfect {}', 'unblemished {}', '{} without flaw', '{} without defect', '{} without damage']
+prompt_abnormal = ['damaged {}', 'broken {}', '{} with flaw', '{} with defect', '{} with damage']
+
+prompt_state = [prompt_normal, prompt_abnormal]
+prompt_templates = ['a photo of a {}.', 'a photo of the {}.']
+# prompt_templates = [
+#                         'a cropped photo of the {}.', 'a cropped photo of a {}.', 'a close-up photo of a {}.', 'a close-up photo of the {}.',
+#                         'a bright photo of the {}.', 'a bright photo of a {}.', 'a dark photo of a {}.', 'a dark photo of the {}.',
+#                         'a dark photo of the {}.', 'a dark photo of a {}.', 'a jpeg corrupted photo of a {}.', 'a jpeg corrupted photo of the {}.',
+#                         'a blurry photo of the {}.', 'a blurry photo of a {}.', 'a photo of a {}.', 'a photo of the {}.',
+#                         'a photo of the small {}.', 'a photo of a small {}.', 'a photo of the large {}.', 'a photo of a large {}.',
+#                         'a photo of the {} for visual insprction.', 'a photo of a {} for visual insprction.',
+#                         'a photo of the {} for anomaly detection.', 'a photo of a {} for anomaly detection.'
+#                         ]
+objs = ['bottle', 'cable', 'capsule', 'carpet', 'grid', 'hazelnut', 'leather', 'metal nut', 'pill', 'screw', 'tile', 'toothbrush', 'transistor', 'wood', 'zipper', 'object',
+        'candle', 'cashew', 'chewinggum', 'fryum', 'macaroni', 'pcb', 'pipe fryum', 'macaroni1', 'macaroni2','pcb1', 'pcb2', 'pcb3', 'pcb4', 'capsules']
+
+prompt_sentences = {}
+
+for obj in objs:
+    prompt_sentence_obj = []
+    for i in range(len(prompt_state)):
+        prompted_state = [state.format(obj) for state in prompt_state[i]]
+        prompted_sentence = []
+        for s in prompted_state:
+            for template in prompt_templates:
+                prompted_sentence.append(template.format(s))
+        prompted_sentence = data.load_and_transform_text(prompted_sentence, torch.cuda.current_device())#torch.cuda.current_device())
+        prompt_sentence_obj.append(prompted_sentence)
+    prompt_sentences[obj] = prompt_sentence_obj
+
+
+
+def encode_text_with_prompt_ensemble(model, obj, device):
+
+    global prompt_sentences
+    normal_sentences = []
+    abnormal_sentences = []
+    for idx in range(len(obj)):
+        sentence = prompt_sentences[obj[idx].replace('_', ' ')]
+        normal_sentences.append(sentence[0])
+        abnormal_sentences.append(sentence[1])
+
+    normal_sentences = torch.cat(normal_sentences).to(device)
+    abnormal_sentences = torch.cat(abnormal_sentences).to(device)
+
+    class_embeddings_normal = model({ModalityType.TEXT: normal_sentences})[ModalityType.TEXT][0]
+    class_embeddings_abnormal = model({ModalityType.TEXT: abnormal_sentences})[ModalityType.TEXT][0]
+    # class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)
+
+    class_embeddings_normal = class_embeddings_normal.reshape((len(obj), len(prompt_templates) * len(prompt_normal), 1024))
+    class_embeddings_normal = class_embeddings_normal.mean(dim=1, keepdim=True)
+    class_embeddings_normal = class_embeddings_normal / class_embeddings_normal.norm(dim=-1, keepdim=True)
+
+    class_embeddings_abnormal = class_embeddings_abnormal.reshape((len(obj), len(prompt_templates) * len(prompt_abnormal), 1024))
+    class_embeddings_abnormal = class_embeddings_abnormal.mean(dim=1, keepdim=True)
+    class_embeddings_abnormal = class_embeddings_abnormal / class_embeddings_abnormal.norm(dim=-1, keepdim=True)
+
+    text_features = torch.cat([class_embeddings_normal, class_embeddings_abnormal], dim=1)
+
+    return text_features
+
+
+
+class StoppingCriteriaSub(StoppingCriteria):
+
+    def __init__(self, stops = [], encounters=1):
+        super().__init__()
+        self.stops = stops
+        self.ENCOUNTERS = encounters
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
+        stop_count = 0
+        for stop in self.stops:
+            stop_count = (stop == input_ids[0]).sum().item()
+        if stop_count >= self.ENCOUNTERS:
+            return True
+        return False
+
+def build_one_instance(tokenizer, conversation):
+    text_list = []
+    turn_num = len(conversation)
+    input_ids, target_ids = [], []
+    for i in range(turn_num):
+        turn = conversation[i]
+        role = turn['from']
+        if i == 0: # the first human turn
+            assert role == 'human'
+            text = turn['value'] + '\n### Assistant:'
+            one_input_id = tokenizer(text, add_special_tokens=False).input_ids
+            input_ids += one_input_id
+            target_ids += [-100]*len(one_input_id) # do not perform loss regression on human prompt
+        else:
+            if role == 'human':
+                text = 'Human: ' + turn['value'] + '\n### Assistant:'
+                one_input_id = tokenizer(text, add_special_tokens=False).input_ids
+                input_ids += one_input_id
+                target_ids += [-100]*len(one_input_id)
+            elif role == 'gpt':
+                text = turn['value'] + '\n###'
+                one_input_id = tokenizer(text, add_special_tokens=False).input_ids
+                input_ids += one_input_id
+                target_ids += one_input_id
+            else:
+                raise Exception('Wrong Role!!!')
+        text_list.append(text)
+        assert len(input_ids) == len(target_ids)
+    return text_list, input_ids, target_ids
+
+def process_batch_instance(tokenizer, batch_of_conversations, max_tgt_len):
+    batch_input_ids, batch_target_ids = [], []
+    for conversation in batch_of_conversations:
+        _, one_input_ids, one_target_ids = build_one_instance(tokenizer, conversation)
+        batch_input_ids.append(torch.LongTensor(one_input_ids))
+        batch_target_ids.append(torch.LongTensor(one_target_ids))
+    input_ids = rnn.pad_sequence(batch_input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
+    target_ids = rnn.pad_sequence(batch_target_ids, batch_first=True, padding_value=-100)
+    assert input_ids.size() == target_ids.size()
+    input_ids = input_ids[:,:max_tgt_len]
+    target_ids = target_ids[:,:max_tgt_len]
+    attention_mask = input_ids.ne(tokenizer.pad_token_id)
+    assert attention_mask.size() == input_ids.size()
+    return input_ids, target_ids, attention_mask.long()
+
+def find_first_file_in_directory(directory_path):
+    try:
+        file_list = os.listdir(directory_path)
+        for item in file_list:
+            item_path = os.path.join(directory_path, item)
+            if os.path.isfile(item_path):
+                return item_path
+        return None
+
+    except OSError as e:
+        print(f"Error while accessing directory: {e}")
+        return None
+
+
+PROMPT_START = '### Human: <Img>'
+class OpenLLAMAPEFTModel(nn.Module):
+
+    '''LoRA for LLaMa model'''
+
+    def __init__(self, **args):
+        super(OpenLLAMAPEFTModel, self).__init__()
+        self.args = args
+        imagebind_ckpt_path = args['imagebind_ckpt_path']
+        vicuna_ckpt_path = args['vicuna_ckpt_path']
+        max_tgt_len = args['max_tgt_len']
+        stage = args['stage']
+
+        self.device = torch.cuda.current_device()
+
+        print (f'Initializing visual encoder from {imagebind_ckpt_path} ...')
+
+        self.visual_encoder, self.visual_hidden_size = imagebind_model.imagebind_huge(args)
+        self.visual_encoder.to(torch.float16).to(self.device)
+        imagebind_ckpt = torch.load(imagebind_ckpt_path, map_location=torch.device('cpu'))
+        self.visual_encoder.load_state_dict(imagebind_ckpt, strict=True)
+        
+
+        self.iter = 0
+
+        self.image_decoder = LinearLayer(1280, 1024, 4).to(torch.float16).to(self.device)
+
+        self.prompt_learner = PromptLearner(1, 4096).to(torch.float16).to(self.device)
+
+        self.loss_focal = FocalLoss()
+        self.loss_dice = BinaryDiceLoss()
+
+
+        # free vision encoder
+        for name, param in self.visual_encoder.named_parameters():
+            param.requires_grad = False
+        self.visual_encoder.eval()
+        print ('Visual encoder initialized.')
+
+        print (f'Initializing language decoder from {vicuna_ckpt_path} ...')
+        
+        # add the lora module
+        peft_config = LoraConfig(
+            task_type=TaskType.CAUSAL_LM, 
+            inference_mode=False, 
+            r=self.args['lora_r'], 
+            lora_alpha=self.args['lora_alpha'], 
+            lora_dropout=self.args['lora_dropout'],
+            target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']
+        )
+
+        # config = AutoConfig.from_pretrained(vicuna_ckpt_path)
+        # with init_empty_weights():
+        #     self.llama_model = AutoModelForCausalLM.from_config(config)
+
+        # # device_map = infer_auto_device_map(self.llama_model, no_split_module_classes=["OPTDecoderLayer"], dtype="float16")
+        # # print(device_map)
+        device_map = {'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layers.4': 0, 'model.layers.5': 0, 'model.layers.6': 0, 'model.layers.7': 0, 'model.layers.8': 0, 'model.layers.9': 0, 'model.layers.10.self_attn': 0, 'model.layers.10.mlp.gate_proj': 0, 'model.layers.10.mlp.down_proj': 'cpu', 'model.layers.10.mlp.up_proj': 'cpu', 'model.layers.10.mlp.act_fn': 'cpu', 'model.layers.10.input_layernorm': 'cpu', 'model.layers.10.post_attention_layernorm': 'cpu', 'model.layers.11': 'cpu', 'model.layers.12': 'cpu', 'model.layers.13': 'cpu', 'model.layers.14': 'cpu', 'model.layers.15': 'cpu', 'model.layers.16': 'cpu', 'model.layers.17': 'cpu', 'model.layers.18': 'cpu', 'model.layers.19': 'cpu', 'model.layers.20': 'cpu', 'model.layers.21': 'cpu', 'model.layers.22': 'cpu', 'model.layers.23': 'cpu', 'model.layers.24': 'disk', 'model.layers.25': 'disk', 'model.layers.26': 'disk', 'model.layers.27': 'disk', 'model.layers.28': 'disk', 'model.layers.29': 'disk', 'model.layers.30': 'disk', 'model.layers.31.self_attn': 'disk', 'model.layers.31.mlp.gate_proj': 'disk', 'model.layers.31.mlp.down_proj': 'disk', 'model.layers.31.mlp.up_proj': 'disk', 'model.layers.31.mlp.act_fn': 'disk', 'model.layers.31.input_layernorm': 'disk', 'model.layers.31.post_attention_layernorm': 'disk', 'model.norm': 'disk', 'lm_head': 'disk'}
+        # # self.llama_model = load_checkpoint_and_dispatch(self.llama_model, vicuna_ckpt_path, device_map=device_map, offload_folder="offload", offload_state_dict = True)
+        # # self.llama_model.to(torch.float16)
+        # # try:
+        self.llama_model = AutoModelForCausalLM.from_pretrained(vicuna_ckpt_path, torch_dtype=torch.float16, device_map='auto', load_in_8bit=True)
+        # # except:
+        #     pass
+        # finally:
+        #     print(self.llama_model.hf_device_map)
+        self.llama_model = get_peft_model(self.llama_model, peft_config)
+        # delta_ckpt = torch.load(args['delta_ckpt_path'], map_location=torch.device('cpu'))
+        # self.llama_model.load_state_dict(delta_ckpt, strict=False)
+        self.llama_model.print_trainable_parameters()
+
+        self.llama_tokenizer = LlamaTokenizer.from_pretrained(vicuna_ckpt_path, use_fast=False, torch_dtype=torch.float16)
+        self.llama_tokenizer.pad_token = self.llama_tokenizer.eos_token
+        self.llama_tokenizer.padding_side = "right"
+        print ('Language decoder initialized.')
+
+        self.llama_proj = nn.Linear(
+            self.visual_hidden_size, self.llama_model.config.hidden_size
+        ).to(torch.float16).to(self.device)
+
+        self.max_tgt_len = max_tgt_len
+        
+
+
+    def rot90_img(self,x,k):
+        # k is 0,1,2,3
+        degreesarr = [0., 90., 180., 270., 360]
+        degrees = torch.tensor(degreesarr[k]).to(self.llama_model.dtype).to(self.device)
+        x = K.geometry.transform.rotate(x, angle = degrees, padding_mode='reflection')
+        return x
+
+    def encode_video(self, video_paths):
+        inputs = {ModalityType.VISION: data.load_and_transform_video_data(video_paths, self.device)}
+        # convert into visual dtype
+        inputs = {key: inputs[key].to(self.llama_model.dtype) for key in inputs}
+        with torch.no_grad():
+            embeddings = self.visual_encoder(inputs)
+            video_embeds = embeddings[ModalityType.VISION][0] # bsz x 1024
+        inputs_llama = self.llama_proj(video_embeds).unsqueeze(1) # bsz x 1 x llama_size
+        atts_llama = torch.ones(inputs_llama.size()[:-1], dtype=torch.long).to(self.device) # bsz x 1
+        return inputs_llama, atts_llama
+
+    def encode_audio(self, audio_paths):
+        inputs = {ModalityType.AUDIO: data.load_and_transform_audio_data(audio_paths, self.device)}
+        # convert into visual dtype
+        inputs = {key: inputs[key].to(self.llama_model.dtype) for key in inputs}
+        with torch.no_grad():
+            embeddings = self.visual_encoder(inputs)
+            audio_embeds = embeddings[ModalityType.AUDIO][0] # bsz x 1024
+        inputs_llama = self.llama_proj(audio_embeds).unsqueeze(1) # bsz x 1 x llama_size
+        atts_llama = torch.ones(inputs_llama.size()[:-1], dtype=torch.long).to(self.device) # bsz x 1
+        return inputs_llama, atts_llama
+
+    def encode_thermal(self, thermal_paths):
+        inputs = {ModalityType.THERMAL: data.load_and_transform_thermal_data(thermal_paths, self.device)}
+        # convert into visual dtype
+        inputs = {key: inputs[key].to(self.llama_model.dtype) for key in inputs}
+        with torch.no_grad():
+            embeddings = self.visual_encoder(inputs)
+            image_embeds = embeddings['thermal'][0] # bsz x 1024
+        inputs_llama = self.llama_proj(image_embeds).unsqueeze(1) # bsz x 1 x llama_size
+        atts_llama = torch.ones(inputs_llama.size()[:-1], dtype=torch.long).to(self.device) # bsz x 1
+        return inputs_llama, atts_llama
+
+    def encode_image(self, image_paths):
+        inputs = {ModalityType.VISION: data.load_and_transform_vision_data(image_paths, self.device)}
+        # convert into visual dtype
+        inputs = {key: inputs[key].to(self.llama_model.dtype) for key in inputs}
+        with torch.no_grad():
+            embeddings = self.visual_encoder(inputs)
+            image_embeds = embeddings['vision'][0] # bsz x 1024
+            patch_features = embeddings['vision'][1] # bsz x h*w x 1280
+        patch_tokens = self.image_decoder(patch_features) # bsz x h*w x 1024
+
+        inputs_llama = self.llama_proj(image_embeds).unsqueeze(1) # bsz x 1 x llama_size
+        atts_llama = torch.ones(inputs_llama.size()[:-1], dtype=torch.long).to(self.device) # bsz x 1
+        return inputs_llama, atts_llama, patch_tokens
+    
+    def encode_image_for_web_demo(self, image_paths):
+        inputs = {ModalityType.VISION: data.load_and_transform_vision_data_for_web_demo(image_paths, self.device)}
+        # convert into visual dtype
+        inputs = {key: inputs[key].to(self.llama_model.dtype) for key in inputs}
+        with torch.no_grad():
+            embeddings = self.visual_encoder(inputs)
+            image_embeds = embeddings['vision'][0] # bsz x 1024
+            patch_features = embeddings['vision'][1] # bsz x h*w x 1280
+        patch_tokens = self.image_decoder(patch_features) # bsz x h*w x 1024
+
+        inputs_llama = self.llama_proj(image_embeds).unsqueeze(1) # bsz x 1 x llama_size
+        atts_llama = torch.ones(inputs_llama.size()[:-1], dtype=torch.long).to(self.device) # bsz x 1
+        return inputs_llama, atts_llama, patch_tokens
+    
+    def encode_image_for_one_shot(self, image_paths):
+        inputs = {ModalityType.VISION: data.load_and_transform_vision_data(image_paths, self.device)}
+        # convert into visual dtype
+        inputs = {key: inputs[key].to(self.llama_model.dtype) for key in inputs}
+        with torch.no_grad():
+            embeddings = self.visual_encoder(inputs)
+            patch_features = embeddings['vision'][1] # bsz x h*w x 1280
+            for i in range(len(patch_features)):
+                patch_features[i] = patch_features[i].transpose(0, 1)[:, 1:, :]
+
+        return patch_features
+    
+    def encode_image_for_one_shot_from_tensor(self, image_tensors):
+        if not isinstance(image_tensors, list):
+            image_tensors = [image_tensors]
+        inputs = {ModalityType.VISION: torch.stack(image_tensors, dim=0).to(self.device)}
+        # convert into visual dtype
+        inputs = {key: inputs[key].to(self.llama_model.dtype) for key in inputs}
+        with torch.no_grad():
+            embeddings = self.visual_encoder(inputs)
+            patch_features = embeddings['vision'][1] # bsz x h*w x 1280
+            for i in range(len(patch_features)):
+                patch_features[i] = patch_features[i].transpose(0, 1)[:, 1:, :]
+
+        return patch_features
+    
+    def encode_image_for_one_shot_with_aug(self, image_paths):
+        image_tensors = data.load_and_transform_vision_data(image_paths, self.device).to(self.llama_model.dtype)
+        B,C,H,W = image_tensors.shape
+        # print(B,C,H,W)
+
+        rotated_images = torch.zeros((4, B, C, H, W)).to(self.llama_model.dtype).to(self.device)
+
+
+        for j, degree in enumerate([0, 1, 2, 3]):
+            rotated_img = self.rot90_img(image_tensors, degree)
+            # 存储旋转后的图像
+            rotated_images[j] = rotated_img
+
+        image_tensors = rotated_images.transpose(0,1).reshape(B * 4, C, H, W)
+
+        inputs = {ModalityType.VISION: image_tensors}
+        # convert into visual dtype
+        inputs = {key: inputs[key] for key in inputs}
+        with torch.no_grad():
+            embeddings = self.visual_encoder(inputs)
+            patch_features = embeddings['vision'][1] # bsz x h*w x 1280
+            for i in range(len(patch_features)):
+                patch_features[i] = patch_features[i].transpose(0, 1)[:, 1:, :].reshape(B,4,256,1280).reshape(B, 4 * 256, 1280)
+
+        return patch_features
+    
+    def encode_image_from_tensor(self, image_tensors):
+        if not isinstance(image_tensors, list):
+            image_tensors = [image_tensors]
+        inputs = {ModalityType.VISION: torch.stack(image_tensors, dim=0).to(self.device)}
+        # convert into visual dtype
+        inputs = {key: inputs[key].to(self.llama_model.dtype) for key in inputs}
+        with torch.no_grad():
+            embeddings = self.visual_encoder(inputs)
+            image_embeds = embeddings['vision'][0] # bsz x 1024
+            patch_features = embeddings['vision'][1] # bsz x h*w x 1024
+        patch_tokens = self.image_decoder(patch_features)
+
+
+        inputs_llama = self.llama_proj(image_embeds).unsqueeze(1) # bsz x 1 x llama_size
+        atts_llama = torch.ones(inputs_llama.size()[:-1], dtype=torch.long).to(self.device) # bsz x 1
+        return inputs_llama, atts_llama, patch_tokens
+    
+    def encode_image_from_tensor_no_patch(self, image_tensors):
+        if not isinstance(image_tensors, list):
+            image_tensors = [image_tensors]
+        inputs = {ModalityType.VISION: torch.stack(image_tensors, dim=0).to(self.device)}
+        # convert into visual dtype
+        inputs = {key: inputs[key].to(self.llama_model.dtype) for key in inputs}
+        with torch.no_grad():
+            embeddings = self.visual_encoder(inputs)
+            image_embeds = embeddings['vision'][0] # bsz x 1024
+
+        inputs_llama = self.llama_proj(image_embeds).unsqueeze(1) # bsz x 1 x llama_size
+        atts_llama = torch.ones(inputs_llama.size()[:-1], dtype=torch.long).to(self.device) # bsz x 1
+        return inputs_llama, atts_llama
+
+
+
+    def prompt_wrap(self, img_embeds, input_ids, target_ids, attention_mask, anomaly_embedding = None):
+        '''
+            input_ids, target_ids, attention_mask: bsz x s2
+        '''
+        input_ids = input_ids.to(self.device) # bsz x s2
+        target_ids = target_ids.to(self.device) # bsz x s2
+        attention_mask = attention_mask.to(self.device) # bsz x s2
+
+        batch_size = img_embeds.shape[0]
+        p_before = PROMPT_START
+        p_before_tokens = self.llama_tokenizer(p_before, 
+            return_tensors="pt", add_special_tokens=False).to(self.device)
+        # peft model need deeper call
+        p_before_embeds = self.llama_model.model.model.embed_tokens(p_before_tokens.input_ids).expand(batch_size, -1, -1) # bsz x s1 x embed_dim
+
+        p_middle = '</Img> '
+        p_middle_tokens = self.llama_tokenizer(p_middle, 
+            return_tensors="pt", add_special_tokens=False).to(self.device)
+        # peft model need deeper call
+        p_middle_embeds = self.llama_model.model.model.embed_tokens(p_middle_tokens.input_ids).expand(batch_size, -1, -1) # bsz x s1 x embed_dim
+
+
+        p_after_embeds = self.llama_model.model.model.embed_tokens(input_ids).expand(batch_size, -1, -1) # bsz x s2 x embed_dim
+        bos = torch.ones([batch_size, 1],
+                         dtype=p_before_tokens.input_ids.dtype,
+                         device=p_before_tokens.input_ids.device) * self.llama_tokenizer.bos_token_id # bsz x 1
+        bos_embeds = self.llama_model.model.model.embed_tokens(bos) # bsz x 1 x embed_dim
+
+        
+
+        if anomaly_embedding != None:
+            inputs_embeds = torch.cat([bos_embeds, p_before_embeds, img_embeds, p_middle_embeds, anomaly_embedding, p_after_embeds], dim=1) # bsz x (1+s1+1+s2) x embed_dim
+            # create targets
+            empty_targets = (
+                torch.ones([batch_size, 1+p_before_embeds.size()[1]+1+p_middle_embeds.size()[1] + anomaly_embedding.size()[1]], # 1 (bos) + s1 + 1 (image vector)
+                        dtype=torch.long).to(self.device).fill_(-100)  
+            ) # bsz x (1 + s1 + 1)
+            targets = torch.cat([empty_targets, target_ids], dim=1) # bsz x (1 + s1 + 1 + s2)
+            assert inputs_embeds.size()[1] == targets.size()[1]
+
+            atts_prefix = torch.ones([batch_size, 1+p_before_embeds.size()[1]+1+p_middle_embeds.size()[1] + anomaly_embedding.size()[1]], dtype=torch.long).to(self.device) # bsz x (1 + s1 +1)
+            attention_mask = torch.cat([atts_prefix, attention_mask], dim=1)
+            assert attention_mask.size() == targets.size() # bsz x (1 + s1 + 1 + s2)
+            return inputs_embeds, targets, attention_mask 
+        else:
+            inputs_embeds = torch.cat([bos_embeds, p_before_embeds, img_embeds, p_middle_embeds, p_after_embeds], dim=1) # bsz x (1+s1+1+s2) x embed_dim
+            # create targets
+            empty_targets = (
+                torch.ones([batch_size, 1+p_before_embeds.size()[1]+1+p_middle_embeds.size()[1]], # 1 (bos) + s1 + 1 (image vector)
+                        dtype=torch.long).to(self.device).fill_(-100)  
+            ) # bsz x (1 + s1 + 1)
+            targets = torch.cat([empty_targets, target_ids], dim=1) # bsz x (1 + s1 + 1 + s2)
+            assert inputs_embeds.size()[1] == targets.size()[1]
+
+            atts_prefix = torch.ones([batch_size, 1+p_before_embeds.size()[1]+1+p_middle_embeds.size()[1]], dtype=torch.long).to(self.device) # bsz x (1 + s1 +1)
+            attention_mask = torch.cat([atts_prefix, attention_mask], dim=1)
+            assert attention_mask.size() == targets.size() # bsz x (1 + s1 + 1 + s2)
+            return inputs_embeds, targets, attention_mask 
+
+
+    def forward(self, inputs):
+
+        if 'masks' in inputs:
+
+            image_paths = inputs['images']
+            img_embeds, _, patch_tokens = self.encode_image_from_tensor(image_paths)
+            class_name = inputs['class_names']
+
+            loss_pixel = 0
+            feats_text_tensor = encode_text_with_prompt_ensemble(self.visual_encoder, ['object' for _ in class_name], self.device)
+
+            anomaly_maps = []
+            for layer in range(len(patch_tokens)):
+                patch_tokens[layer] = patch_tokens[layer] / patch_tokens[layer].norm(dim=-1, keepdim=True)
+                # print(patch_tokens[layer].shape)
+                # anomaly_map = torch.bmm(patch_tokens[layer], feats_text_tensor.transpose(-2,-1))
+                anomaly_map = (100.0 * patch_tokens[layer] @ feats_text_tensor.transpose(-2,-1))
+                B, L, C = anomaly_map.shape
+                H = int(np.sqrt(L))
+                anomaly_map = F.interpolate(anomaly_map.permute(0, 2, 1).view(B, 2, H, H),
+                                            size=224, mode='bilinear', align_corners=True)
+                # anomaly_map_no_softmax = anomaly_map
+                anomaly_map = torch.softmax(anomaly_map, dim=1)
+                anomaly_maps.append(anomaly_map)
+                # anomaly_maps_ns.append(anomaly_map_no_softmax)           
+
+            gt = inputs['masks']
+            gt = torch.stack(gt, dim=0).to(self.device)
+            gt = gt.squeeze()
+            # print(gt.max(), gt.min())
+            gt[gt > 0.3], gt[gt <= 0.3] = 1, 0
+            
+
+            for num in range(len(anomaly_maps)):
+                f_loss = self.loss_focal(anomaly_maps[num], gt)
+                d_loss = self.loss_dice(anomaly_maps[num][:, 1, :, :], gt)
+                loss_pixel = loss_pixel + f_loss + d_loss
+
+            for num in range(len(anomaly_maps)):
+                anomaly_maps[num] = anomaly_maps[num][:,1,:,:]
+
+            anomaly_map_all = torch.mean(torch.stack(anomaly_maps, dim=0), dim=0).unsqueeze(1)
+        
+            if random.randint(0,1) == 0 and len(inputs['img_paths']) == len(image_paths):
+
+                normal_paths = []
+                for path in inputs['img_paths']:
+                    normal_path = path.replace('test', 'train')
+                    normal_path = find_first_file_in_directory("/".join(normal_path.split('/')[:-2])+'/good')
+                    normal_paths.append(normal_path)
+
+                print(normal_paths)
+                query_patch_tokens = self.encode_image_for_one_shot_from_tensor(image_paths)
+                normal_patch_tokens = self.encode_image_for_one_shot_with_aug(normal_paths)
+                sims = []
+                B = len(image_paths)
+
+                for i in range(len(query_patch_tokens)):
+                    query_patch_tokens_reshaped = query_patch_tokens[i].view(B,256,1,1280)
+                    normal_tokens_reshaped = normal_patch_tokens[i].reshape(B,1,-1,1280)
+                    cosine_similarity_matrix = F.cosine_similarity(query_patch_tokens_reshaped, normal_tokens_reshaped, dim=-1)
+                    sim_max, _ = torch.max(cosine_similarity_matrix, dim=-1)
+                    sims.append(sim_max)
+
+                sim = torch.mean(torch.stack(sims,dim=0), dim=0).reshape(B,1,16,16)
+                sim = F.interpolate(sim,size=224, mode='bilinear', align_corners=True)
+                anomaly_map_all = 1 - sim # (anomaly_map_all + 1 - sim) / 2
+
+            anomaly_map_prompts = self.prompt_learner(anomaly_map_all)
+
+            # img_embeds = img_embeds + anomaly_map_prompts
+
+            output_texts = inputs['texts']
+            input_ids, target_ids, attention_mask = process_batch_instance(self.llama_tokenizer, output_texts, self.max_tgt_len)
+            inputs_embeds, targets, attention_mask = self.prompt_wrap(img_embeds, input_ids, target_ids, attention_mask, anomaly_map_prompts)
+
+            outputs = self.llama_model(
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                return_dict=True,
+                labels=targets,
+            )
+            loss = outputs.loss
+
+            # loss_l2 = torch.norm(anomaly_map_prompts / 2 , p=2)
+            # loss_l2 = nn.MSELoss()(img_embeds_origin, img_embeds)
+            # calculate the token accuarcy
+            chosen_tokens = torch.max(outputs.logits, dim=-1)[1][:, 1:-1]    # [B, S-1]
+            # print(self.llama_tokenizer.decode(chosen_tokens[0], skip_special_tokens=True))
+            labels = targets[:, 2:]
+            gen_acc = (chosen_tokens.reshape(-1) == labels.reshape(-1)).to(torch.long)    # [B*S]
+            valid_mask = (labels != -100).reshape(-1)
+            # print(self.llama_tokenizer.decode(chosen_tokens.reshape(-1)[valid_mask], skip_special_tokens=True))
+            valid_tokens = gen_acc & valid_mask    # [B*S]
+            gen_acc = valid_tokens.sum().item() / valid_mask.sum().item()
+
+            return loss + loss_pixel, gen_acc
+        
+        else:
+            
+            image_paths = inputs['image_paths']
+            img_embeds, _, patch_tokens = self.encode_image_from_tensor(image_paths)
+
+            output_texts = inputs['output_texts']
+
+            c_name = 'object'
+            for name in CLASS_NAMES:
+                if name in output_texts:
+                    c_name = name
+                    break
+
+            feats_text_tensor = encode_text_with_prompt_ensemble(self.visual_encoder, ['object'] * len(image_paths), self.device)
+
+            anomaly_maps = []
+            for layer in range(len(patch_tokens)):
+                patch_tokens[layer] = patch_tokens[layer] / patch_tokens[layer].norm(dim=-1, keepdim=True)
+                # print(patch_tokens[layer].shape)
+                # anomaly_map = torch.bmm(patch_tokens[layer], feats_text_tensor.transpose(-2,-1))
+                anomaly_map = (100.0 * patch_tokens[layer] @ feats_text_tensor.transpose(-2,-1))
+                B, L, C = anomaly_map.shape
+                H = int(np.sqrt(L))
+                anomaly_map = F.interpolate(anomaly_map.permute(0, 2, 1).view(B, 2, H, H),
+                                            size=224, mode='bilinear', align_corners=True)
+                # anomaly_map_no_softmax = anomaly_map
+                anomaly_map = torch.softmax(anomaly_map, dim=1)
+                anomaly_maps.append(anomaly_map)
+
+            for num in range(len(anomaly_maps)):
+                anomaly_maps[num] = anomaly_maps[num][:,1,:,:]
+
+            anomaly_map_all = torch.mean(torch.stack(anomaly_maps, dim=0), dim=0).unsqueeze(1)
+
+            anomaly_map_prompts = self.prompt_learner(anomaly_map_all)
+
+            # img_embeds = img_embeds + anomaly_map_prompts
+
+            input_ids, target_ids, attention_mask = process_batch_instance(self.llama_tokenizer, output_texts, self.max_tgt_len)
+            inputs_embeds, targets, attention_mask = self.prompt_wrap(img_embeds, input_ids, target_ids, attention_mask, anomaly_map_prompts)
+
+            outputs = self.llama_model(
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                return_dict=True,
+                labels=targets,
+            )
+            loss = outputs.loss
+            # calculate the token accuarcy
+            chosen_tokens = torch.max(outputs.logits, dim=-1)[1][:, 1:-1]    # [B, S-1]
+            labels = targets[:, 2:]
+            gen_acc = (chosen_tokens.reshape(-1) == labels.reshape(-1)).to(torch.long)    # [B*S]
+            valid_mask = (labels != -100).reshape(-1)
+            valid_tokens = gen_acc & valid_mask    # [B*S]
+            gen_acc = valid_tokens.sum().item() / valid_mask.sum().item()
+            
+            return loss, gen_acc
+
+
+    def extract_multimodal_feature(self, inputs, web_demo):
+        features = []
+        if inputs['image_paths']:
+            
+            prompt = inputs['prompt']
+            c_name = 'object'
+            for name in CLASS_NAMES:
+                if name in prompt:
+                    c_name = name
+                    break
+                
+            if not web_demo:
+                image_embeds, _, patch_tokens = self.encode_image(inputs['image_paths'])
+                feats_text_tensor = encode_text_with_prompt_ensemble(self.visual_encoder, [c_name], self.device)
+            else:
+                image_embeds, _, patch_tokens = self.encode_image_for_web_demo(inputs['image_paths'])
+                feats_text_tensor = encode_text_with_prompt_ensemble(self.visual_encoder, ['object'], self.device)
+
+            anomaly_maps = []
+            for layer in range(len(patch_tokens)):
+                patch_tokens[layer] = patch_tokens[layer] / patch_tokens[layer].norm(dim=-1, keepdim=True)
+                # print(patch_tokens[layer].shape)
+                # anomaly_map = torch.bmm(patch_tokens[layer], feats_text_tensor.transpose(-2,-1))
+                anomaly_map = (100.0 * patch_tokens[layer] @ feats_text_tensor.transpose(-2,-1))
+                B, L, C = anomaly_map.shape
+                H = int(np.sqrt(L))
+                # anomaly_map = anomaly_map.to(torch.float16)
+                anomaly_map = F.interpolate(anomaly_map.permute(0, 2, 1).view(B, 2, H, H),
+                                            size=224, mode='bilinear', align_corners=True)
+                # anomaly_map = anomaly_map.to(torch.bfloat16)
+                anomaly_map = torch.softmax(anomaly_map, dim=1)
+                anomaly_maps.append(anomaly_map[:,1,:,:])
+
+            anomaly_map_ret = torch.mean(torch.stack(anomaly_maps, dim=0), dim=0).unsqueeze(1)
+            # anomaly_map_all = anomaly_map_ret.unsqueeze(1).repeat((1,3,1,1))
+            # anomaly_map_feature, _, _ = self.encode_image_from_tensor(anomaly_map_all)
+            # image_embeds = anomaly_map_feature + image_embeds
+            if inputs['normal_img_paths']:
+                query_patch_tokens = self.encode_image_for_one_shot(inputs['image_paths'])
+                if 'mvtec' in 'normal_img_paths':
+                    normal_patch_tokens = self.encode_image_for_one_shot_with_aug(inputs['normal_img_paths'])
+                else:
+                    normal_patch_tokens = self.encode_image_for_one_shot(inputs['normal_img_paths'])
+                sims = []
+
+                for i in range(len(query_patch_tokens)):
+                    query_patch_tokens_reshaped = query_patch_tokens[i].view(256,1,1280)
+                    normal_tokens_reshaped = normal_patch_tokens[i].reshape(1,-1,1280)
+                    cosine_similarity_matrix = F.cosine_similarity(query_patch_tokens_reshaped, normal_tokens_reshaped, dim=2)
+                    sim_max, _ = torch.max(cosine_similarity_matrix, dim=1)
+                    sims.append(sim_max)
+
+                sim = torch.mean(torch.stack(sims,dim=0), dim=0).reshape(1,1,16,16)
+                # anomaly_map = anomaly_map.to(torch.float16)
+                sim = F.interpolate(sim,size=224, mode='bilinear', align_corners=True)
+                # anomaly_map = anomaly_map.to(torch.bfloat16)
+                anomaly_map_ret = 1 - sim # (anomaly_map_ret + 1 - sim) / 2
+                
+
+            features.append(image_embeds)
+        if inputs['audio_paths']:
+            audio_embeds, _ = self.encode_audio(inputs['audio_paths'])
+            features.append(audio_embeds)
+        if inputs['video_paths']:
+            video_embeds, _ = self.encode_video(inputs['video_paths'])
+            features.append(video_embeds)
+        if inputs['thermal_paths']:
+            thermal_embeds, _ = self.encode_thermal(inputs['thermal_paths'])
+            features.append(thermal_embeds)
+
+        feature_embeds = torch.cat(features).sum(dim=0).unsqueeze(0)
+        return feature_embeds, anomaly_map_ret
+
+    def prepare_generation_embedding(self, inputs, web_demo):
+        prompt = inputs['prompt']
+        # if len(inputs['modality_embeds']) == 1:
+        #     feature_embeds = inputs['modality_embeds'][0]
+        # else:
+        feature_embeds, anomaly_map = self.extract_multimodal_feature(inputs, web_demo)
+        # print(anomaly_map.shape)
+        inputs['modality_embeds'].append(feature_embeds)
+
+        batch_size = feature_embeds.shape[0]
+        p_before = PROMPT_START
+        p_before_tokens = self.llama_tokenizer(p_before, 
+            return_tensors="pt", add_special_tokens=False).to(self.device)
+        p_before_embeds = self.llama_model.model.model.embed_tokens(p_before_tokens.input_ids).expand(batch_size, -1, -1) # bsz x s1 x embed_dim
+        
+        p_middle = '</Img> '
+        p_middle_tokens = self.llama_tokenizer(p_middle, 
+            return_tensors="pt", add_special_tokens=False).to(self.device)
+        # peft model need deeper call
+        p_middle_embeds = self.llama_model.model.model.embed_tokens(p_middle_tokens.input_ids).expand(batch_size, -1, -1) # bsz x s1 x embed_dim
+
+        # self.prompt_learner.eval()
+        anomaly_map_prompts = self.prompt_learner(anomaly_map)
+
+
+
+
+        text = prompt + '\n### Assistant:'
+        p_after_tokens = self.llama_tokenizer(text, add_special_tokens=False, return_tensors='pt').to(self.device)
+        p_after_embeds = self.llama_model.model.model.embed_tokens(p_after_tokens.input_ids).expand(batch_size, -1, -1) # bsz x s2 x embed_dim
+        bos = torch.ones([batch_size, 1],
+                         dtype=p_before_tokens.input_ids.dtype,
+                         device=p_before_tokens.input_ids.device) * self.llama_tokenizer.bos_token_id # bsz x 1
+        bos_embeds = self.llama_model.model.model.embed_tokens(bos) # bsz x 1 x embed_dim
+        inputs_embeds = torch.cat([bos_embeds, p_before_embeds, feature_embeds, p_middle_embeds, anomaly_map_prompts, p_after_embeds], dim=1) # bsz x (1+s1+1+s2) x embed_dim
+    
+        return inputs_embeds, anomaly_map
+
+    def generate(self, inputs, web_demo=False):
+        '''
+            inputs = {
+                'image_paths': optional,
+                'audio_paths': optional
+                'video_paths': optional
+                'thermal_paths': optional
+                'mode': generation mode,
+                'prompt': human input prompt,
+                'max_tgt_len': generation length,
+                'top_p': top_p,
+                'temperature': temperature
+                'modality_embeds': None or torch.tensor
+                'modality_cache': save the image cache
+            }
+        '''
+        # self.prompt_learner.eval()
+        # self.llama_model.eval()
+        # self.llama_proj.eval()
+        # self.image_decoder.eval()
+        # self.llama_tokenizer.eval()
+        input_embeds, pixel_output = self.prepare_generation_embedding(inputs, web_demo)
+        stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=[2277], encounters=1)])
+        outputs = self.llama_model.generate(
+            inputs_embeds=input_embeds,
+            max_new_tokens=inputs['max_tgt_len'],
+            top_p=inputs['top_p'],
+            temperature=inputs['temperature'],
+            do_sample=True,
+            use_cache=True,
+            stopping_criteria=stopping_criteria,
+        )
+        output_text = self.llama_tokenizer.decode(outputs[0][:-2], skip_special_tokens=True)
+        return output_text, pixel_output