Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import torch | |
| import torchvision.transforms as TF | |
| from safetensors.torch import load_file | |
| from huggingface_hub import hf_hub_download | |
| from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection | |
| from model import PredictorModel | |
| class CityAestheticsPipeline: | |
| """ | |
| Demo model pipeline for [image=>score] prediction | |
| Accepts a single model path on initialization. | |
| Resulting object can be called directly with a PIL image as the input | |
| Returns a single float value with the predicted score [0.0;1.0]. | |
| """ | |
| clip_ver = "openai/clip-vit-large-patch14-336" | |
| def __init__(self, model_path, device="cpu", clip_dtype=torch.float32): | |
| self.device = device | |
| self.clip_dtype = clip_dtype | |
| self._init_clip() | |
| self.model = self._load_model(model_path) | |
| print("CityAesthetics: Pipeline init ok") # debug | |
| def __call__(self, raw): | |
| emb = self.get_clip_emb(raw) | |
| return self.get_model_pred(self.model, emb) | |
| def get_model_pred(self, model, emb): | |
| with torch.no_grad(): | |
| pred = model(emb) | |
| return float(pred.detach().cpu().squeeze(0)) | |
| def get_clip_emb(self, raw): | |
| img = self.proc( | |
| images = raw, | |
| return_tensors = "pt" | |
| )["pixel_values"].to(self.clip_dtype).to(self.device) | |
| with torch.no_grad(): | |
| emb = self.clip(pixel_values=img) | |
| return emb["image_embeds"].detach().to(torch.float32) | |
| def _init_clip(self): | |
| self.proc = CLIPImageProcessor.from_pretrained(self.clip_ver) | |
| self.clip = CLIPVisionModelWithProjection.from_pretrained( | |
| self.clip_ver, | |
| device_map = self.device, | |
| torch_dtype = self.clip_dtype, | |
| ) | |
| def _load_model(self, path): | |
| sd = load_file(path) | |
| assert tuple(sd["up.0.weight"].shape) == (1024, 768) # only allow CLIP ver | |
| model = PredictorModel(outputs=1) | |
| model.eval() | |
| model.load_state_dict(sd) | |
| model.to(self.device) | |
| return model | |
| class CityAestheticsMultiModelPipeline(CityAestheticsPipeline): | |
| """ | |
| Demo multi-model pipeline for [image=>score] prediction | |
| Accepts a list of model paths on initialization. | |
| Resulting object can be called directly with a PIL image as the input. | |
| Returns a dict with the model name as key and the score [0.0;1.0] as a value. | |
| """ | |
| def __init__(self, model_paths, device="cpu", clip_dtype=torch.float32): | |
| self.device = device | |
| self.clip_dtype = clip_dtype | |
| self._init_clip() | |
| self.models = {} | |
| for path in model_paths: | |
| name = os.path.splitext(os.path.basename(path))[0] | |
| self.models[name] = self._load_model(path) | |
| print("CityAesthetics: Pipeline init ok") # debug | |
| def __call__(self, raw): | |
| emb = self.get_clip_emb(raw) | |
| out = {} | |
| for name, model in self.models.items(): | |
| pred = model(emb) | |
| out[name] = self.get_model_pred(model, emb) | |
| return out | |
| class CityClassifierPipeline: | |
| """ | |
| Demo model pipeline for [image=>label] prediction | |
| Accepts a single model path and (optionally) a JSON file on initialization. | |
| Resulting object can be called directly with a PIL image as the input | |
| Returns a single float value with the predicted score [0.0;1.0]. | |
| """ | |
| clip_ver = "openai/clip-vit-large-patch14-336" | |
| def __init__(self, model_path, config_path=None, device="cpu", clip_dtype=torch.float32): | |
| self.device = device | |
| self.clip_dtype = clip_dtype | |
| self._init_clip() | |
| self.labels, model_args = self._load_config(config_path) | |
| self.model = self._load_model(model_path, model_args) | |
| print("CityClassifier: Pipeline init ok") # debug | |
| def __call__(self, raw, default=True, tiling=True, tile_strat="mean"): | |
| emb = self.get_clip_emb(raw, tiling=tiling) | |
| pred = self.get_model_pred(self.model, emb) | |
| return self.format_pred( | |
| pred, | |
| labels = self.labels, | |
| drop = [] if default else [0], | |
| ts = tile_strat if tiling else "raw", | |
| ) | |
| def format_pred(self, pred, labels, drop=[], ts="mean"): | |
| # recombine strategy | |
| if ts == "mean" : vp = lambda x: float(torch.mean(x)) | |
| elif ts == "median": vp = lambda x: float(torch.median(x)) | |
| elif ts == "max" : vp = lambda x: float(torch.max(x)) | |
| elif ts == "min" : vp = lambda x: float(torch.min(x)) | |
| elif ts == "raw" : vp = lambda x: float(x) | |
| else: raise NotImplementedError(f"CityClassifier: Invalid combine strategy '{ts}'!") | |
| # combine pred w/ labels | |
| out = {} | |
| for k in range(len(pred)): | |
| if k in drop: continue | |
| key = labels.get(str(k), str(k)) | |
| out[key] = vp(pred[k]) | |
| return out | |
| def get_model_pred(self, model, emb): | |
| with torch.no_grad(): | |
| pred = model(emb) | |
| pred = pred.detach().cpu() | |
| return [pred[:, x] for x in range(pred.shape[1])] # split | |
| def get_clip_emb(self, raw, tiling=False): | |
| if tiling and min(raw.size) > self.size*2: | |
| if max(raw.size)>1536: | |
| raw = TF.functional.resize(raw, 1536) | |
| raw = TF.functional.five_crop(raw, self.size*2) | |
| img = self.proc( | |
| images = raw, | |
| return_tensors = "pt" | |
| )["pixel_values"].to(self.clip_dtype).to(self.device) | |
| with torch.no_grad(): | |
| emb = self.clip(pixel_values=img) | |
| return emb["image_embeds"].detach().to(torch.float32) | |
| def _init_clip(self): | |
| self.proc = CLIPImageProcessor.from_pretrained(self.clip_ver) | |
| self.size = self.proc.size.get("shortest_edge", 256) | |
| self.clip = CLIPVisionModelWithProjection.from_pretrained( | |
| self.clip_ver, | |
| device_map = self.device, | |
| torch_dtype = self.clip_dtype, | |
| ) | |
| def _load_model(self, path, args=None): | |
| sd = load_file(path) | |
| assert tuple(sd["up.0.weight"].shape) == (1024, 768) # only allow CLIP ver | |
| args = args or { # infer from model | |
| "outputs" : int(sd["down.5.bias"].shape[0]) | |
| } | |
| model = PredictorModel(**args) | |
| model.eval() | |
| model.load_state_dict(sd) | |
| model.to(self.device) | |
| return model | |
| def _load_config(self, path): | |
| if not path or not os.path.isfile(path): | |
| return ({},None) | |
| with open(path) as f: | |
| data = json.loads(f.read()) | |
| return ( | |
| data.get("labels", {}), | |
| data.get("model_params", {}), | |
| ) | |
| class CityClassifierMultiModelPipeline(CityClassifierPipeline): | |
| """ | |
| Demo model pipeline for [image=>label] prediction | |
| Accepts a list of model paths on initialization. | |
| A matching list of JSON files can also be passed in the same order. | |
| Resulting object can be called directly with a PIL image as the input | |
| Returns a single float value with the predicted score [0.0;1.0]. | |
| """ | |
| def __init__(self, model_paths, config_paths=[], device="cpu", clip_dtype=torch.float32): | |
| self.device = device | |
| self.clip_dtype = clip_dtype | |
| self._init_clip() | |
| self.models = {} | |
| self.labels = {} | |
| assert len(model_paths) == len(config_paths) or not config_paths, "CityClassifier: Model and config paths must match!" | |
| for k in range(len(model_paths)): | |
| name = os.path.splitext(os.path.basename(model_paths[k]))[0] # TODO: read from config | |
| self.labels[name], model_args = self._load_config(config_paths[k] if config_paths else None) | |
| self.models[name] = self._load_model(model_paths[k], model_args) | |
| print("CityClassifier: Pipeline init ok") # debug | |
| def __call__(self, raw, default=True, tiling=True, tile_strat="mean"): | |
| emb = self.get_clip_emb(raw, tiling=tiling) | |
| out = {} | |
| for name, model in self.models.items(): | |
| pred = self.get_model_pred(model, emb) | |
| out[name] = self.format_pred( | |
| pred, | |
| labels = self.labels[name], | |
| drop = [] if default else [0], | |
| ts = tile_strat if tiling else "raw", | |
| ) | |
| if len(out.values()) == 1: return list(out.values())[0] # GRADIO HOTFIX | |
| return list(out.values()) | |
| def get_model_path(name, repo, token=True, extension="safetensors", local=False): | |
| """ | |
| Returns local model path or falls back to HF hub if required. | |
| """ | |
| fname = f"{name}.{extension}" | |
| # local path: [models/AesPred-Anime-v1.8.safetensors] | |
| path = os.path.join(os.path.dirname(os.path.realpath(__file__)),"models") | |
| if os.path.isfile(os.path.join(path, fname)): | |
| print(f"Using local model for '{fname}'") | |
| return os.path.join(path, fname) | |
| if local: raise OSError(f"Can't find local model '{fname}'!") | |
| # huggingface hub fallback | |
| print(f"Using HF Hub model for '{fname}'") | |
| return str(hf_hub_download( | |
| token = token, | |
| repo_id = repo, | |
| filename = fname, | |
| )) | |