Spaces:
Sleeping
Sleeping
| import torch | |
| import clip | |
| import os | |
| import time | |
| from PIL import Image | |
| from typing import List | |
| from .mlp import MLP | |
| from .utils import download_from_url | |
| MLP_MODEL_URL = "https://huggingface.co/Eugeoter/waifu-scorer/waifu-scorer-v1-large.pth" | |
| class WaifuScorer: | |
| def __init__(self, model_path: str = None, device: str = 'cuda', verbose=False): | |
| self.verbose = verbose | |
| if self.verbose: | |
| tic = time.time() | |
| print(f"loading pretrained model from `{model_path}`") | |
| if model_path is None or not os.path.isfile(model_path): | |
| model_path = download_from_url(MLP_MODEL_URL) | |
| if device == 'cuda' and not torch.cuda.is_available(): | |
| device = 'cpu' | |
| print("CUDA is not available, using CPU instead") | |
| self.mlp = load_model(model_path, input_size=768, device=device) | |
| self.model2, self.preprocess = load_clip_models("ViT-L/14", device=device) | |
| self.device = self.mlp.device | |
| self.dtype = self.mlp.dtype | |
| self.mlp.eval() | |
| if self.verbose: | |
| toc = time.time() | |
| print(f"model loaded: time_cost={toc-tic:.2f} | device={self.device} | dtype={self.dtype}") | |
| def __call__(self, images: List[Image.Image]) -> List[float]: | |
| if isinstance(images, Image.Image): | |
| images = [images] | |
| n = len(images) | |
| if n == 1: | |
| images = images*2 # batch norm | |
| images = encode_images(images, self.model2, self.preprocess, device=self.device).to(device=self.device, dtype=self.dtype) | |
| predictions = self.mlp(images) | |
| scores = predictions.clamp(0, 10).cpu().numpy().reshape(-1).tolist() | |
| if n == 1: | |
| scores = scores[0] | |
| return scores | |
| def load_clip_models(name: str = "ViT-L/14", device='cuda'): | |
| model2, preprocess = clip.load(name, device=device) # RN50x64 | |
| return model2, preprocess | |
| def load_model(model_path: str = None, input_size=768, device: str = 'cuda', dtype=None): | |
| model = MLP(input_size=input_size) | |
| if model_path: | |
| s = torch.load(model_path, map_location=device) | |
| model.load_state_dict(s) | |
| model.to(device) | |
| if dtype: | |
| model = model.to(dtype=dtype) | |
| return model | |
| def normalized(a: torch.Tensor, order=2, dim=-1): | |
| l2 = a.norm(order, dim, keepdim=True) | |
| l2[l2 == 0] = 1 | |
| return a / l2 | |
| def encode_images(images: List[Image.Image], model2, preprocess, device='cuda') -> torch.Tensor: | |
| if isinstance(images, Image.Image): | |
| images = [images] | |
| image_tensors = [preprocess(img).unsqueeze(0) for img in images] | |
| image_batch = torch.cat(image_tensors).to(device) | |
| image_features = model2.encode_image(image_batch) | |
| im_emb_arr = normalized(image_features).cpu().float() | |
| return im_emb_arr | |