import torch
from transformers import CLIPProcessor, CLIPVisionModel
from modules import devices
import os
from annotator.annotator_path import clip_vision_path


remote_model_path = "https://huggingface.co./openai/clip-vit-large-patch14/resolve/main/pytorch_model.bin"
clip_path = clip_vision_path
print(f'ControlNet ClipVision location: {clip_path}')

clip_proc = None
clip_vision_model = None


def apply_clip(img):
    global clip_proc, clip_vision_model
    
    if clip_vision_model is None:
        modelpath = os.path.join(clip_path, 'pytorch_model.bin')
        if not os.path.exists(modelpath):
            from basicsr.utils.download_util import load_file_from_url
            load_file_from_url(remote_model_path, model_dir=clip_path)

        clip_proc = CLIPProcessor.from_pretrained(clip_path)
        clip_vision_model = CLIPVisionModel.from_pretrained(clip_path)

    with torch.no_grad():
        clip_vision_model = clip_vision_model.to(devices.get_device_for("controlnet"))
        style_for_clip = clip_proc(images=img, return_tensors="pt")['pixel_values']
        style_feat = clip_vision_model(style_for_clip.to(devices.get_device_for("controlnet")))['last_hidden_state']

    return style_feat


def unload_clip_model():
    global clip_proc, clip_vision_model
    if clip_vision_model is not None:
        clip_vision_model.cpu()