Check the following codes thoroughly. Understand the purpose and suggest fixes.
# src/preprocess_image.py

import json
import torch
from torch.utils.data import Dataset
from PIL import Image
import os

class LlavaDataset(Dataset):
    def init(self, json_file, image_dir, tokenizer, model, device, max_length=512):
        with open(json_file, 'r') as f:
            self.data = json.load(f)
        self.image_dir = image_dir
        self.tokenizer = tokenizer
        self.model = model
        self.device = device
        self.max_length = max_length

        # Print the structure of the first few items
        # print("Dataset structure:")
        # for i in range(min(5, len(self.data))):
        #     print(f"Item {i}:")
        #     print(json.dumps(self.data[i], indent=2))
        # print("End of dataset structure preview")

    def len(self):
        return len(self.data)

    def getitem(self, idx):
        item = self.data[idx]

        # Print the keys of the current item
        # print(f"Keys in item {idx}: {item.keys()}")

        # Check if 'image' key exists and print its value
        if 'image' in item:
            # print(f"Image value for item {idx}: {item['image']}")
            image_file = os.path.join(self.image_dir, item['image'])
            image = Image.open(image_file).convert('RGB')
            if not os.path.exists(image_file):
                print(f"Warning: Image file does not exist: {image_file}")
        else:
            print(f"Warning: 'image' key not found in item {idx}")
            image_file = None

        processed = self.model.module.clip_processor(images=image, return_tensors="pt", padding=True) if isinstance(self.model, torch.nn.DataParallel) else self.model.clip_processor(images=image, return_tensors="pt", padding=True)
        pixel_values = processed['pixel_values'].to(self.device)

        # Handle 5D tensor if produced by clip_processor
        if pixel_values.dim() == 5:
            pixel_values = pixel_values.squeeze(1)
            print(f"Reshaped 5D pixel_values to 4D. New shape: {pixel_values.shape}")

        # Construct prompt from conversation
        conversation = item.get('conversations', [])
        prompt = ""
        for turn in conversation:
            if turn['from'] == 'human':
                prompt += f"Human: {turn['value']}\n"
            elif turn['from'] == 'gpt':
                prompt += f"Assistant: {turn['value']}\n"

        # Replace '<image>' token with a special token that your model understands
        prompt = prompt.replace('<image>', '[IMAGE]')
        # print(prompt)

        # Tokenize the prompt
        encoded_prompt = self.tokenizer(prompt, truncation=True, max_length=self.max_length, padding='max_length', return_tensors='pt')

        return {
            'pixel_values': pixel_values,
            'input_ids': encoded_prompt['input_ids'].squeeze(),
            'attention_mask': encoded_prompt['attention_mask'].squeeze()
        }

def load_llava_dataset(json_file, image_dir, tokenizer, model, device, max_length=2048):
    return LlavaDataset(json_file, image_dir, tokenizer, model, device, max_length)

# src/train_img.py
import os
import torch
from torch import nn
from transformers import CLIPModel, CLIPProcessor, Trainer, TrainingArguments, AutoTokenizer
from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model
from peft.optimizers import create_loraplus_optimizer
import bitsandbytes as bnb
from preprocess_image import load_llava_dataset
from PIL import Image
import tempfile

# Set environment variable for PyTorch memory allocation
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

# Set custom directories
os.environ['TMPDIR'] = './.cache/temp'
os.environ['HF_HOME'] = './.cache'
tempfile.tempdir = './.cache/temp'

# Set the cache directory for datasets
from datasets import config
config.HF_DATASETS_CACHE = './.cache'
  
class ImageModel(nn.Module):
    def __init__(self, clip_model_name):
        super().__init__()
        self.clip_model = CLIPModel.from_pretrained(clip_model_name, torch_dtype=torch.bfloat16, attn_implementation="sdpa")
        self.clip_processor = CLIPProcessor.from_pretrained(clip_model_name)
        
        # Freeze CLIP model
        for param in self.clip_model.parameters():
            param.requires_grad = False
        
        clip_hidden_size = 768 #self.clip_model.config.vision_config.hidden_size
        llm_hidden_size = 2048  # Adjust this to match your LLM's hidden size
        # print(clip_hidden_size,"- clip hidden size")
        # Add a trainable adapter layer
        self.adapter = nn.Sequential(
            bnb.nn.Linear4bit(clip_hidden_size, llm_hidden_size, compress_statistics=True),
            nn.GELU(),
            bnb.nn.Linear4bit(llm_hidden_size, llm_hidden_size, compress_statistics=True)
        )

    def forward(self, pixel_values):
        # Print input shape for debugging
        print(f"Input shape: {pixel_values.shape}")

        # Handle 5D input (if produced by clip_processor)
        if pixel_values.dim() == 5:
            b, extra_dim, c, h, w = pixel_values.shape
            pixel_values = pixel_values.squeeze(1)
            print(f"Reshaped 5D input to 4D. New shape: {pixel_values.shape}")

        # Ensure pixel_values are in the correct shape (batch_size, channels, height, width)
        if pixel_values.dim() != 4:
            raise ValueError(f"Expected 4D input, got {pixel_values.dim()}D input")

        # Ensure all tensors are on the same device
        device = pixel_values.device
        self.clip_model = self.clip_model.to(device)
        self.adapter = self.adapter.to(device)

        with torch.no_grad():
            image_features = self.clip_model.get_image_features(pixel_values, output_hidden_states=True)
        
        print(f"CLIP output shape: {image_features.shape}")
        adapter_output = self.adapter(image_features)
        print(f"Adapter output shape: {adapter_output.shape}")
        return adapter_output
    
    def gradient_checkpointing_enable(self):
        # Enable gradient checkpointing for the adapter
        self.adapter.gradient_checkpointing_enable()
    
def main():
     # Check if CUDA is available
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Load the image model
    model = ImageModel(clip_model_name="openai/clip-vit-large-patch14-336")
    model = model.to(device)

    # Prepare only the adapter for k-bit training
    model.adapter = prepare_model_for_kbit_training(model.adapter)
    
    # Enable gradient checkpointing
    # model.gradient_checkpointing_enable()

    # Define LoRA config for the adapter layer
    peft_config = LoraConfig(
        r=8,
        lora_alpha=32,
        lora_dropout=0.1,
        bias="none",
        target_modules=["0", "2"]  # Target the Linear layers in the adapter
    )

    # Get the PEFT model (only for the adapter)
    model.adapter = get_peft_model(model.adapter, peft_config)

    # Use DataParallel if multiple GPUs are available
    if torch.cuda.device_count() > 1:
        print(f"Using {torch.cuda.device_count()} GPUs")
        model = torch.nn.DataParallel(model)

    # Load tokenizer (use the same tokenizer as your text model)
    tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
    
    # Add [IMAGE] token to the tokenizer
    tokenizer.add_special_tokens({'additional_special_tokens': ['[IMAGE]']})

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

    # Load the dataset
    train_dataset = load_llava_dataset(
        json_file='data/llava_instruct_150k/llava_instruct_150k.json',
        image_dir='data/coco/train2017',  # Adjust this path to where your images are stored
        tokenizer=tokenizer,
        model = model,
        device = device
    )
    # train_dataset.save_to_disk('data/processed/llava_instruct_150k_train_phi2')
    # # train_dataset.load_from_disk('data/processed/llava_instruct_150k_train_phi2')

    # Define training arguments
    training_args = TrainingArguments(
        output_dir="./results",
        num_train_epochs=1,
        per_device_train_batch_size=16,
        gradient_accumulation_steps=8,
        warmup_steps=100,
        weight_decay=0.01,
        logging_dir='./logs',
        logging_steps=10,
        save_steps=500,
        # fp16=True,
        bf16=True,
        optim="paged_adamw_32bit",
        # gradient_checkpointing=True,
        report_to="tensorboard",
    )

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

    # Define a custom collate function to handle image processing
    def collate_fn(batch):
        batch = [item for item in batch if item is not None]
        
        if not batch:
            print("Warning: Empty batch encountered")
            return None
        
        pixel_values = torch.stack([item['pixel_values'] 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])

        print(f"Collated batch size: {len(batch)}")
        print(f"pixel_values shape: {pixel_values.shape}")
        print(f"input_ids shape: {input_ids.shape}")
        print(f"attention_mask shape: {attention_mask.shape}")
        return {
            'pixel_values': pixel_values,
            'input_ids': input_ids,
            'attention_mask': attention_mask,
        }

    # Initialize the Trainer
    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()