model = AutoModelForVision2Seq.from_pretrained(
  args.pretrained_ckpt,
  torch_dtype=compute_dtype,
  quantization_config=BitsAndBytesConfig(
      load_in_4bit=bits == 4,
      load_in_8bit=bits == 8,
      llm_int8_threshold=6.0,
      int8_quant_skip_modules=int8_quant_skip_modules,
      llm_int8_has_fp16_weight=False,
      bnb_4bit_compute_dtype=compute_dtype,
      bnb_4bit_use_double_quant=True,
      bnb_4bit_quant_type='nf4'  # {'fp4', 'nf4'}
  ) if bits < 16 else None,
  attn_implementation=args.attn_implementation,
)

if (args.use_lora and not resume_from_checkpoint and not ft_checkpoint_dir):
  target_modules = get_target_modules(model.model.text_model, args, bits)
  peft_config = LoraConfig(
      target_modules=target_modules,
      inference_mode=args.inference_mode,
      r=args.lora_r,
      lora_alpha=args.lora_alpha,
      lora_dropout=args.lora_dropout,
      use_dora=args.use_dora
  )
  model.model.text_model = get_peft_model(model.model.text_model, peft_config)

  if args.vit_train:
      target_modules = get_target_modules(model.model.vision_model, args, args.vit_bits, vit=True)
      peft_config = LoraConfig(
          target_modules=target_modules,
          inference_mode=args.inference_mode,
          r=args.vit_lora_r,
          lora_alpha=args.vit_lora_alpha,
          lora_dropout=args.lora_dropout,
          use_dora=args.use_dora_vit
      )
      model.model.vision_model = get_peft_model(model.model.vision_model, peft_config)

  if args.lora_abstractor:
      target_modules = get_target_modules(model.model.connector, args, args.bits)
      peft_config = LoraConfig(
          target_modules=target_modules,
          inference_mode=args.inference_mode,
          r=args.lora_r,
          lora_alpha=args.lora_alpha,
          lora_dropout=args.lora_dropout,
          use_dora=args.use_dora
      )
      model.model.connector = get_peft_model(model.model.connector, peft_config)