From c7e22ccd757c7f5ba5e459bcf86416fd803c3afa Mon Sep 17 00:00:00 2001
From: younesbelkada <younesbelkada@gmail.com>
Date: Tue, 4 Apr 2023 07:59:03 +0000
Subject: [PATCH] v1

---
 src/peft/mapping.py     |  8 +++++++-
 src/peft/tuners/lora.py | 11 +++++++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/peft/mapping.py b/src/peft/mapping.py
index dbb9f36..1e98edb 100644
--- a/src/peft/mapping.py
+++ b/src/peft/mapping.py
@@ -19,6 +19,7 @@ from .peft_model import (
     PeftModelForSeq2SeqLM,
     PeftModelForSequenceClassification,
     PeftModelForTokenClassification,
+    PeftModelForVision2Seq,
 )
 from .tuners import LoraConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig
 from .utils import PromptLearningConfig
@@ -29,6 +30,7 @@ MODEL_TYPE_TO_PEFT_MODEL_MAPPING = {
     "SEQ_2_SEQ_LM": PeftModelForSeq2SeqLM,
     "CAUSAL_LM": PeftModelForCausalLM,
     "TOKEN_CLS": PeftModelForTokenClassification,
+    "VISION_2_SEQ": PeftModelForVision2Seq,
 }
 
 PEFT_TYPE_TO_CONFIG_MAPPING = {
@@ -44,6 +46,7 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
     "bart": ["q_proj", "v_proj"],
     "gpt2": ["c_attn"],
     "bloom": ["query_key_value"],
+    "blip2": ["q", "v", "q_proj", "v_proj"],
     "opt": ["q_proj", "v_proj"],
     "gptj": ["q_proj", "v_proj"],
     "gpt_neox": ["query_key_value"],
@@ -134,9 +137,12 @@ def get_peft_model(model, peft_config):
         model ([`transformers.PreTrainedModel`]): Model to be wrapped.
         peft_config ([`PeftConfig`]): Configuration object containing the parameters of the Peft model.
     """
-
     model_config = model.config.to_dict()
     peft_config.base_model_name_or_path = model.__dict__.get("name_or_path", None)
+
+    if peft_config.task_type == "VISION_2_SEQ" and not isinstance(peft_config, LoraConfig):
+        raise ValueError("Vision2Seq task type is only supported with LORA")
+
     if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys():
         peft_config = _prepare_lora_config(peft_config, model_config)
         return PeftModel(model, peft_config)
diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
index 51cd56f..d4d17ac 100644
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -394,6 +394,7 @@ class Linear(nn.Linear, LoraLayer):
         self.lora_B.eval()
 
     def forward(self, x: torch.Tensor):
+
         if self.disable_adapters:
             if self.r > 0 and self.merged:
                 self.weight.data -= (
@@ -401,14 +402,20 @@ class Linear(nn.Linear, LoraLayer):
                 )
                 self.merged = False
 
-            return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
+            result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
+
+            return result
         elif self.r > 0 and not self.merged:
             result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
             if self.r > 0:
+                x = x.to(self.lora_A.weight.dtype)
+
                 result += self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling
             return result
         else:
-            return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
+            result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
+
+            return result
 
 
 class MergedLinear(nn.Linear, LoraLayer):