An Open State-of-the-Art Vision-Language Foundation Model

				
					import torch
from peft import LoraConfig
from transformers import AutoProcessor, BitsAndBytesConfig, Idefics2ForConditionalGeneration

# Specify the device to use for model training and inference (CUDA for GPU)
DEVICE = "cuda:0"
# Flags to indicate the use of Low-Rank Adaptation (LoRA) or Quantized LoRA (QLORA)
USE_LORA = False
USE_QLORA = True

# Load the pre-trained Idefics2 model processor from the Hugging Face hub
processor = AutoProcessor.from_pretrained(
    "HuggingFaceM4/idefics2-8b",
    do_image_splitting=False # Disable image splitting for simplicity
)

# Configure the model based on the flags for LoRA and QLORA
if USE_QLORA or USE_LORA:
    # Define the configuration for LoRA
    lora_config = LoraConfig(
        r=8, # Rank of the LoRA updates
        lora_alpha=8, # Scaling factor for the updates
        lora_dropout=0.1, # Dropout rate for LoRA
        target_modules='.*(text_model|modality_projection|perceiver_resampler).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$',
        use_dora=False if USE_QLORA else True, # Disable DORA if using QLORA
        init_lora_weights="gaussian" # Initialization method for LoRA weights
    )
    
    if USE_QLORA:
        # Define the configuration for QLORA, including 4-bit quantization
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True, # Enable 4-bit quantization
            bnb_4bit_quant_type="nf4", # Use NF4 quantization type
            bnb_4bit_compute_dtype=torch.float16 # Use FP16 for computation
        )
    
    # Load the Idefics2 model with the specified LoRA or QLORA configurations
    model = Idefics2ForConditionalGeneration.from_pretrained(
        "HuggingFaceM4/idefics2-8b",
        torch_dtype=torch.float16, # Use FP16 for model weights
        quantization_config=bnb_config if USE_QLORA else None, # Apply QLORA config if enabled
    )
    
    # Add the LoRA configuration to the model and enable adapters
    model.add_adapter(lora_config)
    model.enable_adapters()
else:
    # Load the Idefics2 model for full fine-tuning without LoRA or QLORA
    model = Idefics2ForConditionalGeneration.from_pretrained(
        "HuggingFaceM4/idefics2-8b",
        torch_dtype=torch.float16, # Use FP16 for model weights
        _attn_implementation="flash_attention_2", # Enable Flash Attention (only available on A100 or H100 GPUs)
    ).to(DEVICE) # Move the model to the specified device