cdnuts/JointTaggerProject-Inference-GPU
0
1import json2import os3import zipfile4from io import BytesIO5from tempfile import NamedTemporaryFile6import tempfile7 8import gradio as gr9import pandas as pd10from PIL import Image11import safetensors.torch12import spaces13import timm14from timm.models import VisionTransformer15import torch16from torchvision.transforms import transforms17from torchvision.transforms import InterpolationMode18import torchvision.transforms.functional as TF19from torch.utils.data import Dataset, DataLoader20 21 22torch.set_grad_enabled(False)23 24class Fit(torch.nn.Module):25 def __init__(26 self,27 bounds: tuple[int, int] | int,28 interpolation = InterpolationMode.LANCZOS,29 grow: bool = True,30 pad: float | None = None31 ):32 super().__init__()33 34 self.bounds = (bounds, bounds) if isinstance(bounds, int) else bounds35 self.interpolation = interpolation36 self.grow = grow37 self.pad = pad38 39 def forward(self, img: Image) -> Image:40 wimg, himg = img.size41 hbound, wbound = self.bounds42 43 hscale = hbound / himg44 wscale = wbound / wimg45 46 if not self.grow:47 hscale = min(hscale, 1.0)48 wscale = min(wscale, 1.0)49 50 scale = min(hscale, wscale)51 if scale == 1.0:52 return img53 54 hnew = min(round(himg * scale), hbound)55 wnew = min(round(wimg * scale), wbound)56 57 img = TF.resize(img, (hnew, wnew), self.interpolation)58 59 if self.pad is None:60 return img61 62 hpad = hbound - hnew63 wpad = wbound - wnew64 65 tpad = hpad // 266 bpad = hpad - tpad67 68 lpad = wpad // 269 rpad = wpad - lpad70 71 return TF.pad(img, (lpad, tpad, rpad, bpad), self.pad)72 73 def __repr__(self) -> str:74 return (75 f"{self.__class__.__name__}(" +76 f"bounds={self.bounds}, " +77 f"interpolation={self.interpolation.value}, " +78 f"grow={self.grow}, " +79 f"pad={self.pad})"80 )81 82class CompositeAlpha(torch.nn.Module):83 def __init__(84 self,85 background: tuple[float, float, float] | float,86 ):87 super().__init__()88 89 self.background = (background, background, background) if isinstance(background, float) else background90 self.background = torch.tensor(self.background).unsqueeze(1).unsqueeze(2)91 92 def forward(self, img: torch.Tensor) -> torch.Tensor:93 if img.shape[-3] == 3:94 return img95 96 alpha = img[..., 3, None, :, :]97 98 img[..., :3, :, :] *= alpha99 100 background = self.background.expand(-1, img.shape[-2], img.shape[-1])101 if background.ndim == 1:102 background = background[:, None, None]103 elif background.ndim == 2:104 background = background[None, :, :]105 106 img[..., :3, :, :] += (1.0 - alpha) * background107 return img[..., :3, :, :]108 109 def __repr__(self) -> str:110 return (111 f"{self.__class__.__name__}(" +112 f"background={self.background})"113 )114 115transform = transforms.Compose([116 Fit((384, 384)),117 transforms.ToTensor(),118 CompositeAlpha(0.5),119 transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),120 transforms.CenterCrop((384, 384)),121])122 123model = timm.create_model(124 "vit_so400m_patch14_siglip_384.webli",125 pretrained=False,126 num_classes=9083,127) # type: VisionTransformer128 129safetensors.torch.load_model(model, "JTP_PILOT-e4-vit_so400m_patch14_siglip_384.safetensors")130device = torch.device("cuda" if torch.cuda.is_available() else "cpu")131model.to(device)132model.eval()133 134with open("tagger_tags.json", "r") as file:135 tags = json.load(file) # type: dict136allowed_tags = list(tags.keys())137 138for idx, tag in enumerate(allowed_tags):139 allowed_tags[idx] = tag.replace("_", " ")140 141sorted_tag_score = {}142 143@spaces.GPU(duration=9)144def run_classifier(image, threshold):145 global sorted_tag_score146 img = image.convert('RGB')147 tensor = transform(img).unsqueeze(0).to(device)148 with torch.no_grad():149 logits = model(tensor)150 probabilities = torch.nn.functional.sigmoid(logits[0])151 indices = torch.topk(probabilities, 250).indices152 values = probabilities[indices]153 154 tag_score = dict()155 for i in range(indices.size(0)):156 tag_score[allowed_tags[indices[i]]] = values[i].item()157 sorted_tag_score = dict(sorted(tag_score.items(), key=lambda item: item[1], reverse=True))158 159 return create_tags(threshold)160 161def create_tags(threshold):162 global sorted_tag_score163 filtered_tag_score = {key: value for key, value in sorted_tag_score.items() if value > threshold}164 text_no_impl = ", ".join(filtered_tag_score.keys())165 return text_no_impl, filtered_tag_score166 167 168class ImageDataset(Dataset):169 def __init__(self, image_files, transform):170 self.image_files = image_files171 self.transform = transform172 173 def __len__(self):174 return len(self.image_files)175 176 def __getitem__(self, idx):177 img_path = self.image_files[idx]178 img = Image.open(img_path).convert('RGB')179 return self.transform(img), os.path.basename(img_path)180 181@spaces.GPU(duration=299)182def process_images(images, threshold):183 dataset = ImageDataset(images, transform)184 185 dataloader = DataLoader(dataset, batch_size=64, num_workers=0, pin_memory=True, drop_last=False)186 187 all_results = []188 189 with torch.no_grad():190 for batch, filenames in dataloader:191 192 batch = batch.to(device) 193 with torch.no_grad():194 logits = model(batch)195 probabilities = torch.nn.functional.sigmoid(logits)196 197 for i, prob in enumerate(probabilities):198 indices = torch.where(prob > threshold)[0]199 values = prob[indices]200 201 temp = []202 tag_score = dict()203 for j in range(indices.size(0)):204 temp.append([allowed_tags[indices[j]], values[j].item()])205 tag_score[allowed_tags[indices[j]]] = values[j].item()206 207 tags = ", ".join([t[0] for t in temp])208 all_results.append((filenames[i], tags, tag_score))209 210 return all_results211 212def is_valid_image(file_path):213 try:214 with Image.open(file_path) as img:215 img.verify()216 return True217 except:218 return False219 220def process_zip(zip_file, threshold):221 if zip_file is None:222 return None, None223 224 with tempfile.TemporaryDirectory() as temp_dir:225 with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:226 zip_ref.extractall(temp_dir)227 228 all_files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir)]229 image_files = [f for f in all_files if is_valid_image(f)]230 results = process_images(image_files, threshold)231 232 temp_file = NamedTemporaryFile(delete=False, suffix=".zip")233 with zipfile.ZipFile(temp_file, "w") as zip_ref:234 for image_name, text_no_impl, _ in results:235 with zip_ref.open(''.join(image_name.split('.')[:-1]) + ".txt", 'w') as file:236 file.write(text_no_impl.encode())237 temp_file.seek(0)238 df = pd.DataFrame([(os.path.basename(f), t) for f, t, _ in results], columns=['Image', 'Tags'])239 240 return temp_file.name, df241 242@spaces.GPU(duration=120) # Reduced GPU duration for less wait time...243def process_images_light(images, threshold):244 dataset = ImageDataset(images, transform)245 246 dataloader = DataLoader(dataset, batch_size=32, num_workers=0, pin_memory=True, drop_last=False)247 248 all_results = []249 250 with torch.no_grad():251 for batch, filenames in dataloader:252 253 batch = batch.to(device) 254 with torch.no_grad():255 logits = model(batch)256 probabilities = torch.nn.functional.sigmoid(logits)257 258 for i, prob in enumerate(probabilities):259 indices = torch.where(prob > threshold)[0]260 values = prob[indices]261 262 temp = []263 tag_score = dict()264 for j in range(indices.size(0)):265 temp.append([allowed_tags[indices[j]], values[j].item()])266 tag_score[allowed_tags[indices[j]]] = values[j].item()267 268 tags = ", ".join([t[0] for t in temp])269 all_results.append((filenames[i], tags, tag_score))270 271 return all_results272 273def process_zip_light(zip_file, threshold):274 if zip_file is None:275 return None, None276 277 with tempfile.TemporaryDirectory() as temp_dir:278 with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:279 zip_ref.extractall(temp_dir)280 281 all_files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir)]282 image_files = [f for f in all_files if is_valid_image(f)]283 results = process_images_light(image_files, threshold)284 285 temp_file = NamedTemporaryFile(delete=False, suffix=".zip")286 with zipfile.ZipFile(temp_file, "w") as zip_ref:287 for image_name, text_no_impl, _ in results:288 with zip_ref.open(''.join(image_name.split('.')[:-1]) + ".txt", 'w') as file:289 file.write(text_no_impl.encode())290 temp_file.seek(0)291 df = pd.DataFrame([(os.path.basename(f), t) for f, t, _ in results], columns=['Image', 'Tags'])292 293 return temp_file.name, df294 295with gr.Blocks(css=".output-class { display: none; }") as demo:296 gr.Markdown("""297 ## Joint Tagger Project: PILOT Demo298 This tagger is designed for use on furry images (though may very well work on out-of-distribution images, potentially with funny results). A threshold of 0.2 is recommended. Lower thresholds often turn up more valid tags, but can also result in some amount of hallucinated tags.299 300 This tagger is the result of joint efforts between members of the RedRocket team. Special thanks to Minotoro at frosting.ai for providing the compute power for this project.301 302 Usage Note for batch tagging:303 304 the normal version is limited to 300s and uses batch size 64305 306 the light version is limited to 120s with batch size 32307 308 if your image count is low use the light version for lower gpu wait time (most of the time you instantly get a gpu anyway)309 """)310 311 with gr.Tabs():312 with gr.TabItem("Single Image"):313 with gr.Row():314 with gr.Column():315 image_input = gr.Image(label="Source", sources=['upload'], type='pil', height=512, show_label=False)316 threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")317 with gr.Column():318 tag_string = gr.Textbox(label="Tag String")319 label_box = gr.Label(label="Tag Predictions", num_top_classes=250, show_label=False)320 321 image_input.upload(322 fn=run_classifier,323 inputs=[image_input, threshold_slider],324 outputs=[tag_string, label_box]325 )326 327 threshold_slider.input(328 fn=create_tags,329 inputs=[threshold_slider],330 outputs=[tag_string, label_box]331 )332 333 with gr.TabItem("Multiple Images"):334 with gr.Row():335 with gr.Column():336 zip_input = gr.File(label="Upload ZIP file", file_types=['.zip'])337 multi_threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")338 process_button = gr.Button("Process Images")339 with gr.Column():340 zip_output = gr.File(label="Download Tagged Text Files (ZIP)")341 dataframe_output = gr.Dataframe(label="Image Tags Summary")342 343 process_button.click(344 fn=process_zip,345 inputs=[zip_input, multi_threshold_slider],346 outputs=[zip_output, dataframe_output]347 )348 with gr.TabItem("Multiple Images (Light)"):349 with gr.Row():350 with gr.Column():351 zip_input_light = gr.File(label="Upload ZIP file", file_types=['.zip'])352 multi_threshold_slider_light = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")353 process_button_light = gr.Button("Process Images (Light)")354 with gr.Column():355 zip_output_light = gr.File(label="Download Tagged Text Files (ZIP)")356 dataframe_output_light = gr.Dataframe(label="Image Tags Summary")357 358 process_button_light.click(359 fn=process_zip_light,360 inputs=[zip_input_light, multi_threshold_slider_light],361 outputs=[zip_output_light, dataframe_output_light]362 )363 364if __name__ == "__main__":365 demo.queue().launch()