import requests

from PIL import Image

import gem
import torch


@torch.no_grad()
def main() -> None:
    model_name = "ViT-L/14-336"
    pretrained = "openai"
    device = "cuda" if torch.cuda.is_available() else "cpu"

    model = gem.create_gem_model(
        model_name=model_name, pretrained=pretrained, device=device
    )
    preprocess = gem.get_gem_img_transform()

    texts = [
        "a photo of cat",
    ]
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    while True:
        images_pil = Image.open(requests.get(url, stream=True).raw)
        img_tensor = preprocess(images_pil).unsqueeze(0).to(device)

        logits = model.forward(img_tensor, texts)
        gem.visualize(img_tensor[0], texts, logits[0])


if __name__ == "__main__":
    main()