File size: 1,392 Bytes
0300664 8d3099f 0300664 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
---
license: mit
datasets:
- markury/AndroAtlas
language:
- en
base_model:
- google/paligemma-3b-mix-448
pipeline_tag: image-text-to-text
tags:
- captioning
---
```python
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()
``` |