How to use qlora in the code below:
import os
import torch
from torch import nn
from transformers import CLIPModel, CLIPProcessor, Trainer, TrainingArguments, AutoTokenizer, AutoModelForCausalLM
from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model
from peft.optimizers import create_loraplus_optimizer
import bitsandbytes as bnb
from torch.amp import autocast
from preprocess_image import load_llava_dataset
import tempfile

class ImageModel(nn.Module):
    def __init__(self, clip_model_name, lm_model_name, tokenizer):
        super().__init__()
        # Load CLIP model
        self.clip_model = CLIPModel.from_pretrained(
            clip_model_name,
            torch_dtype=torch.bfloat16,
        )
        # Load language model
        self.lm_model = AutoModelForCausalLM.from_pretrained(
            lm_model_name,
            torch_dtype=torch.bfloat16,
        )
        self.tokenizer = tokenizer

        # Freeze CLIP model
        for param in self.clip_model.parameters():
            param.requires_grad = False

        # Optionally, freeze the language model to train only the adapter
        # for param in self.lm_model.parameters():
        #     param.requires_grad = False
        # max_length = self.lm_model.config.max_position_embeddings
        # new_max_length = max_length + 1  # Increase by 1 for the image token
        # self.lm_model.resize_position_embeddings(new_max_length)

        # Define a projection layer to map image features to language model's hidden size
        self.image_proj = nn.Linear(
            self.clip_model.config.projection_dim,
            self.lm_model.config.hidden_size,
        )
        # self.image_proj = self.image_proj.to(dtype=torch.bfloat16)

        # For the loss function
        self.loss_fn = nn.CrossEntropyLoss()

    def forward(self, image, input_ids, attention_mask, labels=None):
        device = input_ids.device
        with autocast(device_type='cuda', dtype=torch.bfloat16):
            # Process images
            with torch.no_grad():
                image_features = self.clip_model.get_image_features(
                    pixel_values=image
                )

            # Project image features to match LM's hidden size
            projected_image_features = self.image_proj(image_features)

            # Prepare inputs for language model
            # Assuming you want to prepend the image features to the text embeddings
            batch_size = input_ids.size(0)
            seq_length = input_ids.size(1)

            # Expand image features to match sequence length if needed
            # Here, we'll concatenate the image features as a new token at the beginning
            # Get embeddings from the language model
            inputs_embeds = self.lm_model.get_input_embeddings()(input_ids)

            # Concatenate image features to the embeddings
            # Reshape projected_image_features to (batch_size, 1, hidden_size)
            image_embeds = projected_image_features.unsqueeze(1)
            inputs_embeds = torch.cat([image_embeds, inputs_embeds], dim=1)

            # Adjust attention_mask accordingly
            extended_attention_mask = torch.cat(
                [torch.ones(batch_size, 1, device=device), attention_mask],
                dim=1
            )

            # Pass through language model
            outputs = self.lm_model(
                inputs_embeds=inputs_embeds,
                attention_mask=extended_attention_mask,
                labels=labels
            )

        return outputs

def main():
    # Device setup
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Tokenizer and special tokens
    tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
    tokenizer.add_special_tokens({'additional_special_tokens': ['[IMAGE]']})

    if tokenizer.pad_token is None:
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})

    # num_labels = tokenizer.vocab_size

    # Load the image model
    model = ImageModel(
        clip_model_name="openai/clip-vit-large-patch14-336",
        lm_model_name="microsoft/phi-2",
        tokenizer=tokenizer
    )
    # model = prepare_model_for_kbit_training(model)
    model = model.to(device)
    # Inspect model's modules to find correct target module names
    # print("Model's named modules:")
    # for name, module in model.named_modules():
    #     print(name)
    print(model.lm_model)

    # LoRA configuration and setup
    peft_config = LoraConfig(
        r=8,
        lora_alpha=32,
        lora_dropout=0.1,
        bias="none",
        target_modules=['Wqkv','out_proj', 'fc1', 'fc2'],#"lm_model.transformer.h.*.mlp.*", "lm_model.transformer.h.*.attn.*"], # ["lm_model.*"],
        modules_to_save=["lm_model.lm_head","lm_model.model.embed_tokens"],
        # target_modules=["adapter.0","adapater.2"],  # Target the adapter layers
    )
    model = get_peft_model(model, peft_config)

    # Data loading
    train_dataset = load_llava_dataset(
        json_file='data/llava_instruct_150k/llava_instruct_150k.json',
        image_dir='data/coco/train2017',
        tokenizer=tokenizer,
    )

    # Training arguments
    training_args = TrainingArguments(
        output_dir="./results",
        num_train_epochs=1,
        per_device_train_batch_size=1,  # Adjust based on your GPU memory
        gradient_accumulation_steps=8,
        warmup_steps=100,
        weight_decay=0.01,
        logging_dir='./logs',
        logging_steps=10,
        save_steps=500,
        bf16=True,
        optim="paged_adamw_32bit",
        report_to="tensorboard",
        remove_unused_columns=False
    )

    optimizer = create_loraplus_optimizer(
        model=model,
        optimizer_cls=bnb.optim.PagedAdamW32bit,
        lr=5e-5,
        loraplus_lr_ratio=16,
    )
    scheduler = None

    # Custom collate function
    def collate_fn(batch):
        batch = [item for item in batch if item is not None]
        if not batch:
            return None
        images = torch.stack([item['image'] for item in batch])
        input_ids = torch.stack([item['input_ids'] for item in batch])
        attention_mask = torch.stack([item['attention_mask'] for item in batch])

        # Adjust labels to account for the image token
        labels = input_ids.clone()
        pad_token_id = tokenizer.pad_token_id
        labels = torch.cat([torch.full((labels.size(0), 1), pad_token_id, dtype=torch.long), labels], dim=1)

        return {
            'image': images,
            'input_ids': input_ids,
            'attention_mask': attention_mask,
            'labels': labels,
        }

    # Trainer initialization
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        optimizers=(optimizer, scheduler),
        data_collator=collate_fn,
    )

    # Start training
    trainer.train()

    # Save the final model
    trainer.save_model("./final_image_model")

if __name__ == "__main__":
    main()