33 lines
1.4 KiB
Python
33 lines
1.4 KiB
Python
|
from transformers import CLIPVisionModel, CLIPVisionConfig, CLIPImageProcessor
|
||
|
|
from comfy.sd import load_torch_file
|
||
|
|
import os
|
||
|
|
|
||
|
|
class ClipVisionModel():
|
||
|
|
def __init__(self):
|
||
|
|
json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config.json")
|
||
|
|
config = CLIPVisionConfig.from_json_file(json_config)
|
||
|
|
self.model = CLIPVisionModel(config)
|
||
|
|
self.processor = CLIPImageProcessor(crop_size=224,
|
||
|
|
do_center_crop=True,
|
||
|
|
do_convert_rgb=True,
|
||
|
|
do_normalize=True,
|
||
|
|
do_resize=True,
|
||
|
|
image_mean=[ 0.48145466,0.4578275,0.40821073],
|
||
|
|
image_std=[0.26862954,0.26130258,0.27577711],
|
||
|
|
resample=3, #bicubic
|
||
|
|
size=224)
|
||
|
|
|
||
|
|
def load_sd(self, sd):
|
||
|
|
self.model.load_state_dict(sd, strict=False)
|
||
|
|
|
||
|
|
def encode_image(self, image):
|
||
|
|
inputs = self.processor(images=[image[0]], return_tensors="pt")
|
||
|
|
outputs = self.model(**inputs)
|
||
|
|
return outputs
|
||
|
|
|
||
|
|
def load(ckpt_path):
|
||
|
|
clip_data = load_torch_file(ckpt_path)
|
||
|
|
clip = ClipVisionModel()
|
||
|
|
clip.load_sd(clip_data)
|
||
|
|
return clip
|