Spaces:
Sleeping
Sleeping
| from transformers import CLIPImageProcessor, AutoModel | |
| import torch | |
| import json | |
| import torch.nn as nn | |
| from PIL import Image | |
| import gradio as gr | |
| import os | |
| from huggingface_hub import login, snapshot_download | |
| TITLE = "Danbooru Tagger" | |
| DESCRIPTION = """ | |
| ## Dataset | |
| - Source: Cleaned Danbooru | |
| ## Metrics | |
| - Validation Split: 10% of Dataset | |
| - Validation Results: | |
| ### General | |
| | Metric | Value | | |
| |-----------------|-------------| | |
| | Macro F1 | 0.4678 | | |
| | Macro Precision | 0.4605 | | |
| | Macro Recall | 0.5229 | | |
| | Micro F1 | 0.6661 | | |
| | Micro Precision | 0.6049 | | |
| | Micro Recall | 0.7411 | | |
| ### Character | |
| | Metric | Value | | |
| |-----------------|-------------| | |
| | Macro F1 | 0.8925 | | |
| | Macro Precision | 0.9099 | | |
| | Macro Recall | 0.8935 | | |
| | Micro F1 | 0.9232 | | |
| | Micro Precision | 0.9264 | | |
| | Micro Recall | 0.9199 | | |
| ### Artist | |
| | Metric | Value | | |
| |-----------------|-------------| | |
| | Macro F1 | 0.7904 | | |
| | Macro Precision | 0.8286 | | |
| | Macro Recall | 0.7904 | | |
| | Micro F1 | 0.5989 | | |
| | Micro Precision | 0.5975 | | |
| | Micro Recall | 0.6004 | | |
| """ | |
| kaomojis = [ | |
| "0_0", | |
| "(o)_(o)", | |
| "+_+", | |
| "+_-", | |
| "._.", | |
| "<o>_<o>", | |
| "<|>_<|>", | |
| "=_=", | |
| ">_<", | |
| "3_3", | |
| "6_9", | |
| ">_o", | |
| "@_@", | |
| "^_^", | |
| "o_o", | |
| "u_u", | |
| "x_x", | |
| "|_|", | |
| "||_||", | |
| ] | |
| device = torch.device('cpu') | |
| dtype = torch.float32 | |
| hf_token = os.getenv("HF_TOKEN") | |
| if hf_token: | |
| login(token=hf_token) | |
| else: | |
| raise ValueError("environment variable HF_TOKEN not found.") | |
| repo_id = "Johnny-Z/vit-e4" | |
| repo_dir = snapshot_download(repo_id) | |
| model = AutoModel.from_pretrained(repo_id, dtype=dtype, trust_remote_code=True, device_map=device) | |
| processor = CLIPImageProcessor.from_pretrained(repo_id) | |
| class MultiheadAttentionPoolingHead(nn.Module): | |
| def __init__(self, input_size): | |
| super().__init__() | |
| self.map_probe = nn.Parameter(torch.randn(1, 1, input_size)) | |
| self.map_layernorm0 = nn.LayerNorm(input_size, eps=1e-08) | |
| self.map_attention = torch.nn.MultiheadAttention(input_size, input_size // 64, batch_first=True) | |
| self.map_layernorm1 = nn.LayerNorm(input_size, eps=1e-08) | |
| self.map_ffn = nn.Sequential( | |
| nn.Linear(input_size, input_size * 4), | |
| nn.SiLU(), | |
| nn.Linear(input_size * 4, input_size) | |
| ) | |
| def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: | |
| batch_size = hidden_state.shape[0] | |
| probe = self.map_probe.repeat(batch_size, 1, 1) | |
| hidden_state = self.map_layernorm0(hidden_state) | |
| hidden_state = self.map_attention(probe, hidden_state, hidden_state)[0] | |
| hidden_state = self.map_layernorm1(hidden_state) | |
| residual = hidden_state | |
| hidden_state = residual + self.map_ffn(hidden_state) | |
| return hidden_state[:, 0] | |
| class MLP(nn.Module): | |
| def __init__(self, input_size, class_num): | |
| super().__init__() | |
| self.mlp_layer0 = nn.Sequential( | |
| nn.LayerNorm(input_size, eps=1e-08), | |
| nn.Linear(input_size, input_size // 2), | |
| nn.SiLU() | |
| ) | |
| self.mlp_layer1 = nn.Linear(input_size // 2, class_num) | |
| self.sigmoid = nn.Sigmoid() | |
| def forward(self, x): | |
| x = self.mlp_layer0(x) | |
| x = self.mlp_layer1(x) | |
| x = self.sigmoid(x) | |
| return x | |
| with open(os.path.join(repo_dir, 'general_tag_dict.json'), 'r', encoding='utf-8') as f: | |
| general_dict = json.load(f) | |
| with open(os.path.join(repo_dir, 'character_tag_dict.json'), 'r', encoding='utf-8') as f: | |
| character_dict = json.load(f) | |
| with open(os.path.join(repo_dir, 'artist_tag_dict.json'), 'r', encoding='utf-8') as f: | |
| artist_dict = json.load(f) | |
| with open(os.path.join(repo_dir, 'implications_list.json'), 'r', encoding='utf-8') as f: | |
| implications_list = json.load(f) | |
| with open(os.path.join(repo_dir, 'artist_threshold.json'), 'r', encoding='utf-8') as f: | |
| artist_thresholds = json.load(f) | |
| with open(os.path.join(repo_dir, 'character_threshold.json'), 'r', encoding='utf-8') as f: | |
| character_thresholds = json.load(f) | |
| with open(os.path.join(repo_dir, 'general_threshold.json'), 'r', encoding='utf-8') as f: | |
| general_thresholds = json.load(f) | |
| model_map = MultiheadAttentionPoolingHead(2048) | |
| model_map.load_state_dict(torch.load(os.path.join(repo_dir, "map_head.pth"), map_location=device, weights_only=True)) | |
| model_map.to(device).to(dtype).eval() | |
| general_class = 9775 | |
| mlp_general = MLP(2048, general_class) | |
| mlp_general.load_state_dict(torch.load(os.path.join(repo_dir, "cls_predictor_general.pth"), map_location=device, weights_only=True)) | |
| mlp_general.to(device).to(dtype).eval() | |
| character_class = 7568 | |
| mlp_character = MLP(2048, character_class) | |
| mlp_character.load_state_dict(torch.load(os.path.join(repo_dir, "cls_predictor_character.pth"), map_location=device, weights_only=True)) | |
| mlp_character.to(device).to(dtype).eval() | |
| artist_class = 13957 | |
| mlp_artist = MLP(2048, artist_class) | |
| mlp_artist.load_state_dict(torch.load(os.path.join(repo_dir, "cls_predictor_artist.pth"), map_location=device, weights_only=True)) | |
| mlp_artist.to(device).to(dtype).eval() | |
| def prediction_to_tag(prediction, tag_dict, class_num): | |
| prediction = prediction.view(class_num) | |
| predicted_ids = (prediction >= 0.2).nonzero(as_tuple=True)[0].cpu().numpy() + 1 | |
| general = {} | |
| character = {} | |
| artist = {} | |
| date = {} | |
| rating = {} | |
| for tag, value in tag_dict.items(): | |
| if value[2] in predicted_ids: | |
| tag_value = round(prediction[value[2] - 1].item(), 6) | |
| if value[1] == "general" and tag_value >= general_thresholds.get(tag, {}).get("Threshold", 0.75): | |
| general[tag] = tag_value | |
| elif value[1] == "character" and tag_value >= character_thresholds.get(tag, {}).get("Threshold", 0.75): | |
| character[tag] = tag_value | |
| elif value[1] == "artist" and tag_value >= artist_thresholds.get(tag, {}).get("Threshold", 0.75): | |
| artist[tag] = tag_value | |
| elif value[1] == "rating": | |
| rating[tag] = tag_value | |
| elif value[1] == "date": | |
| date[tag] = tag_value | |
| general = dict(sorted(general.items(), key=lambda item: item[1], reverse=True)) | |
| character = dict(sorted(character.items(), key=lambda item: item[1], reverse=True)) | |
| artist = dict(sorted(artist.items(), key=lambda item: item[1], reverse=True)) | |
| if date: | |
| date = {max(date, key=date.get): date[max(date, key=date.get)]} | |
| if rating: | |
| rating = {max(rating, key=rating.get): rating[max(rating, key=rating.get)]} | |
| return general, character, artist, date, rating | |
| def process_image(image): | |
| try: | |
| image = image.convert('RGBA') | |
| background = Image.new('RGBA', image.size, (255, 255, 255, 255)) | |
| image = Image.alpha_composite(background, image).convert('RGB') | |
| image_inputs = processor(images=[image], return_tensors="pt").to(device).to(dtype) | |
| except (OSError, IOError) as e: | |
| print(f"Error opening image: {e}") | |
| return | |
| with torch.no_grad(): | |
| embedding = model(image_inputs.pixel_values) | |
| embedding = model_map(embedding) | |
| general_prediction = mlp_general(embedding) | |
| general_ = prediction_to_tag(general_prediction, general_dict, general_class) | |
| general_tags = general_[0] | |
| rating = general_[4] | |
| character_prediction = mlp_character(embedding) | |
| character_ = prediction_to_tag(character_prediction, character_dict, character_class) | |
| character_tags = character_[1] | |
| artist_prediction = mlp_artist(embedding) | |
| artist_ = prediction_to_tag(artist_prediction, artist_dict, artist_class) | |
| artist_tags = artist_[2] | |
| date = artist_[3] | |
| combined_tags = {**general_tags} | |
| tags_list = [tag for tag in combined_tags] | |
| remove_list = [] | |
| for tag in tags_list: | |
| if tag in implications_list: | |
| for implication in implications_list[tag]: | |
| remove_list.append(implication) | |
| tags_list = [tag for tag in tags_list if tag not in remove_list] | |
| tags_list = [tag.replace("_", " ") if tag not in kaomojis else tag for tag in tags_list] | |
| tags_str = ", ".join(tags_list).replace("(", r"\(").replace(")", r"\)") | |
| return tags_str, artist_tags, character_tags, general_tags, rating, date | |
| def main(): | |
| with gr.Blocks(title=TITLE) as demo: | |
| with gr.Column(): | |
| gr.Markdown( | |
| value=f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(variant="panel"): | |
| submit = gr.Button(value="Submit", variant="primary", size="lg") | |
| image = gr.Image(type="pil", image_mode="RGBA", label="Input") | |
| with gr.Row(): | |
| clear = gr.ClearButton( | |
| components=[ | |
| image, | |
| ], | |
| variant="secondary", | |
| size="lg", | |
| ) | |
| gr.Markdown(value=DESCRIPTION) | |
| with gr.Column(variant="panel"): | |
| tags_str = gr.Textbox(label="Output", lines=4) | |
| with gr.Row(): | |
| rating = gr.Label(label="Rating") | |
| date = gr.Label(label="Year") | |
| artist_tags = gr.Label(label="Artist") | |
| character_tags = gr.Label(label="Character") | |
| general_tags = gr.Label(label="General") | |
| clear.add( | |
| [ | |
| tags_str, | |
| artist_tags, | |
| general_tags, | |
| character_tags, | |
| rating, | |
| date, | |
| ] | |
| ) | |
| submit.click( | |
| process_image, | |
| inputs=[ | |
| image | |
| ], | |
| outputs=[tags_str, artist_tags, character_tags, general_tags, rating, date], | |
| ) | |
| demo.queue(max_size=10) | |
| demo.launch() | |
| if __name__ == "__main__": | |
| main() | |