Views
No views yet
1import torch
2from transformers import AutoModelForImageTextToText, AutoProcessor, BitsAndBytesConfig
3from peft import PeftModel
4
5# Configuration
6BASE_MODEL_ID = "google/medgemma-4b-it"
7ADAPTER_ID = "pamessina/medgemma-4b-it-cure"
8
9# 4-bit Quantization Config (Must match training settings)
10quantization_config = BitsAndBytesConfig(
11 load_in_4bit=True,
12 bnb_4bit_use_double_quant=True,
13 bnb_4bit_quant_type="nf4",
14 bnb_4bit_compute_dtype=torch.bfloat16,
15 bnb_4bit_quant_storage=torch.bfloat16,
16)
17
18# Load Model
19base_model = AutoModelForImageTextToText.from_pretrained(
20 BASE_MODEL_ID,
21 quantization_config=quantization_config,
22 torch_dtype=torch.bfloat16,
23 device_map="auto",
24)
25
26# Load the LoRA adapter onto the base model
27model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
28
29# Load the processor
30processor = AutoProcessor.from_pretrained(BASE_MODEL_ID)
31processor.tokenizer.padding_side = "left"
32
33model.eval()
34print("CURE model and processor are ready for inference.")
35get_pil_with_augmentations_transforms Albumentations setup, please refer to the GitHub Repository or the Colab link above.1@InProceedings{Messina_2026_CVPR,
2 author = {Messina, Pablo and Villa, Andr\'es and Alcazar, Juan Leon and Sanchez, Karen and Hinojosa, Carlos and Parra, Denis and Soto, Alvaro and Ghanem, Bernard},
3 title = {CURE: Curriculum-guided Multi-task Training for Reliable Anatomy Grounded Report Generation},
4 booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
5 month = {June},
6 year = {2026},
7 pages = {36279-36289}
8}
9