import torch
import time
import threading
from aiohttp import web
import aiohttp
import asyncio
import traceback
import msgpack
import collections
import queue
import open_clip
from PIL import Image
from prometheus_client import Counter, Histogram, REGISTRY, generate_latest
import io
import json
import torchvision.transforms.transforms as transforms
import sys

with open(sys.argv[1], "r") as config_file:
    CONFIG = json.load(config_file)

batch = 2

device = torch.device("cpu")
model, _, preprocess = open_clip.create_model_and_transforms(CONFIG["model"], device=device, pretrained=CONFIG.get("model_path", dict(open_clip.list_pretrained())[CONFIG["model"]]), precision="fp16")
model.eval()
tokenizer = open_clip.get_tokenizer(CONFIG["model"])
print("Model loaded")

images = torch.stack([ preprocess(Image.open("/data/public/memes-or-something/0mg.jpg")).half() for _ in range(batch) ]).to(device)

class Wrapper(torch.nn.Module):
    def __init__(self, x):
        super().__init__()
        self.model = x

    def forward(self, images):
        return self.model.encode_images(images)

print(model)

torch.onnx.export(
    Wrapper(model),
    (images,),
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=14,
)