metadata
license: mit
datasets:
- markury/AndroAtlas
language:
- en
base_model:
- google/paligemma-3b-mix-448
pipeline_tag: image-text-to-text
tags:
- captioning
import torch
from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor
from PIL import Image
def download_model(model_id):
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id)
processor = PaliGemmaProcessor.from_pretrained(model_id)
return model, processor
def infer(model, processor, image_path, text, max_new_tokens=128):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model.eval()
image = Image.open(image_path)
inputs = processor(text=text, images=image, return_tensors="pt").to(device)
with torch.inference_mode():
generated_ids = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=False
)
result = processor.batch_decode(generated_ids, skip_special_tokens=True)
return result[0][len(text):].lstrip("\n")
def main():
model_id = "prolapse/malensfw-paligemma-fp8_e5m2"
model, processor = download_model(model_id)
image_path = "/path/to/image.png"
prompt = "describe this photo"
result = infer(model, processor, image_path, prompt)
print(result)
if __name__ == "__main__":
main()