Views
No views yet
1import torch
2from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor
3from PIL import Image
4
5def download_model(model_id):
6 model = PaliGemmaForConditionalGeneration.from_pretrained(model_id)
7 processor = PaliGemmaProcessor.from_pretrained(model_id)
8 return model, processor
9
10def infer(model, processor, image_path, text, max_new_tokens=128):
11 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12 model = model.to(device)
13 model.eval()
14
15 image = Image.open(image_path)
16 inputs = processor(text=text, images=image, return_tensors="pt").to(device)
17
18 with torch.inference_mode():
19 generated_ids = model.generate(
20 **inputs,
21 max_new_tokens=max_new_tokens,
22 do_sample=False
23 )
24
25 result = processor.batch_decode(generated_ids, skip_special_tokens=True)
26 return result[0][len(text):].lstrip("\n")
27
28def main():
29 model_id = "prolapse/malensfw-paligemma-bf16"
30 model, processor = download_model(model_id)
31 image_path = "/path/to/image.png"
32 prompt = "describe this photo"
33 result = infer(model, processor, image_path, prompt)
34 print(result)
35
36if __name__ == "__main__":
37 main()