import base64 import io import json import logging import time from pathlib import Path from typing import Any import requests import timm import torch import torchvision.transforms as transforms from PIL import Image import safetensors.torch class TaggingHead(torch.nn.Module): def __init__(self, input_dim, num_classes): super().__init__() self.input_dim = input_dim self.num_classes = num_classes self.head = torch.nn.Sequential(torch.nn.Linear(input_dim, num_classes)) def forward(self, x): logits = self.head(x) probs = torch.nn.functional.sigmoid(logits) return probs def get_tags(tags_file: Path) -> tuple[dict[str, int], int, int]: with tags_file.open("r", encoding="utf-8") as f: tag_info = json.load(f) tag_map = tag_info["tag_map"] tag_split = tag_info["tag_split"] gen_tag_count = tag_split["gen_tag_count"] character_tag_count = tag_split["character_tag_count"] return tag_map, gen_tag_count, character_tag_count def get_character_ip_mapping(mapping_file: Path): with mapping_file.open("r", encoding="utf-8") as f: mapping = json.load(f) return mapping def get_encoder(): base_model_repo = "hf_hub:SmilingWolf/wd-eva02-large-tagger-v3" encoder = timm.create_model(base_model_repo, pretrained=False) encoder.reset_classifier(0) return encoder def get_decoder(): decoder = TaggingHead(1024, 13461) return decoder def get_model(): encoder = get_encoder() decoder = get_decoder() model = torch.nn.Sequential(encoder, decoder) return model def load_model(weights_file, device): model = get_model() # Load from safetensors or pickle weights_path = Path(weights_file) if weights_path.suffix == '.safetensors': print(f"Loading safetensors model: {weights_path}") states_dict = safetensors.torch.load_file(str(weights_path)) else: print(f"Loading pickle model: {weights_path}") states_dict = torch.load(str(weights_path), map_location=device, weights_only=True) model.load_state_dict(states_dict) model.to(device) model.eval() return model def pure_pil_alpha_to_color_v2( image: Image.Image, color: tuple[int, int, int] = (255, 255, 255) ) -> Image.Image: image.load() background = Image.new("RGB", image.size, color) background.paste(image, mask=image.split()[3]) return background def pil_to_rgb(image: Image.Image) -> Image.Image: if image.mode == "RGBA": image = pure_pil_alpha_to_color_v2(image) elif image.mode == "P": image = pure_pil_alpha_to_color_v2(image.convert("RGBA")) else: image = image.convert("RGB") return image class EndpointHandler: def __init__(self, path: str): repo_path = Path(path) assert repo_path.is_dir(), f"Model directory not found: {repo_path}" # Check for safetensors first, then fallback to .pth weights_file = repo_path / "pixai-tagger_v0.9-bf16.safetensors" if not weights_file.exists(): weights_file = repo_path / "model_v0.9.pth" tags_file = repo_path / "tags_v0.9_13k.json" mapping_file = repo_path / "char_ip_map.json" if not weights_file.exists(): raise FileNotFoundError(f"Model file not found: {weights_file}") if not tags_file.exists(): raise FileNotFoundError(f"Tags file not found: {tags_file}") if not mapping_file.exists(): raise FileNotFoundError(f"Mapping file not found: {mapping_file}") self.device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Loading model on {self.device}...") self.model = load_model(weights_file, self.device) self.transform = transforms.Compose( [ transforms.Resize((448, 448)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ] ) self.fetch_image_timeout = 5.0 self.default_general_threshold = 0.3 self.default_character_threshold = 0.85 tag_map, self.gen_tag_count, self.character_tag_count = get_tags(tags_file) self.index_to_tag_map = {v: k for k, v in tag_map.items()} self.character_ip_mapping = get_character_ip_mapping(mapping_file) print("Model loaded successfully!") def __call__(self, data: dict[str, Any]) -> dict[str, Any]: inputs = data.pop("inputs", data) fetch_start_time = time.time() if isinstance(inputs, Image.Image): image = inputs elif image_url := inputs.pop("url", None): with requests.get( image_url, stream=True, timeout=self.fetch_image_timeout ) as res: res.raise_for_status() image = Image.open(res.raw) elif image_base64_encoded := inputs.pop("image", None): image = Image.open(io.BytesIO(base64.b64decode(image_base64_encoded))) else: raise ValueError(f"No image or url provided: {data}") image = pil_to_rgb(image) fetch_time = time.time() - fetch_start_time parameters = data.pop("parameters", {}) general_threshold = parameters.pop( "general_threshold", self.default_general_threshold ) character_threshold = parameters.pop( "character_threshold", self.default_character_threshold ) inference_start_time = time.time() with torch.inference_mode(): image_tensor = self.transform(image).unsqueeze(0).pin_memory() image_tensor = image_tensor.to(self.device, non_blocking=True) probs = self.model(image_tensor)[0] general_mask = probs[: self.gen_tag_count] > general_threshold character_mask = probs[self.gen_tag_count :] > character_threshold general_indices = general_mask.nonzero(as_tuple=True)[0] character_indices = ( character_mask.nonzero(as_tuple=True)[0] + self.gen_tag_count ) combined_indices = torch.cat((general_indices, character_indices)).cpu() inference_time = time.time() - inference_start_time post_process_start_time = time.time() cur_gen_tags = [] cur_char_tags = [] for i in combined_indices: idx = i.item() tag = self.index_to_tag_map[idx] if idx < self.gen_tag_count: cur_gen_tags.append(tag) else: cur_char_tags.append(tag) ip_tags = [] for tag in cur_char_tags: if tag in self.character_ip_mapping: ip_tags.extend(self.character_ip_mapping[tag]) ip_tags = sorted(set(ip_tags)) post_process_time = time.time() - post_process_start_time logging.info( f"Timing - Fetch: {fetch_time:.3f}s, Inference: {inference_time:.3f}s, Post-process: {post_process_time:.3f}s, Total: {fetch_time + inference_time + post_process_time:.3f}s" ) return { "feature": cur_gen_tags, "character": cur_char_tags, "ip": ip_tags, }