prolapse/malensfw-paligemma-fp16

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-fp16"
    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()