TekeshiX/JointTaggerProject-Inference

1import json2 3import gradio as gr4from PIL import Image5import safetensors.torch6import spaces7import timm8from timm.models import VisionTransformer9import torch10from torchvision.transforms import transforms11from torchvision.transforms import InterpolationMode12import torchvision.transforms.functional as TF13 14torch.set_grad_enabled(False)15 16class Fit(torch.nn.Module):17    def __init__(18        self,19        bounds: tuple[int, int] | int,20        interpolation = InterpolationMode.LANCZOS,21        grow: bool = True,22        pad: float | None = None23    ):24        super().__init__()25 26        self.bounds = (bounds, bounds) if isinstance(bounds, int) else bounds27        self.interpolation = interpolation28        self.grow = grow29        self.pad = pad30 31    def forward(self, img: Image) -> Image:32        wimg, himg = img.size33        hbound, wbound = self.bounds34 35        hscale = hbound / himg36        wscale = wbound / wimg37 38        if not self.grow:39            hscale = min(hscale, 1.0)40            wscale = min(wscale, 1.0)41 42        scale = min(hscale, wscale)43        if scale == 1.0:44            return img45 46        hnew = min(round(himg * scale), hbound)47        wnew = min(round(wimg * scale), wbound)48 49        img = TF.resize(img, (hnew, wnew), self.interpolation)50 51        if self.pad is None:52            return img53 54        hpad = hbound - hnew55        wpad = wbound - wnew56 57        tpad = hpad // 258        bpad = hpad - tpad59 60        lpad = wpad // 261        rpad = wpad - lpad62 63        return TF.pad(img, (lpad, tpad, rpad, bpad), self.pad)64 65    def __repr__(self) -> str:66        return (67            f"{self.__class__.__name__}(" +68            f"bounds={self.bounds}, " +69            f"interpolation={self.interpolation.value}, " +70            f"grow={self.grow}, " +71            f"pad={self.pad})"72        )73 74class CompositeAlpha(torch.nn.Module):75    def __init__(76        self,77        background: tuple[float, float, float] | float,78    ):79        super().__init__()80 81        self.background = (background, background, background) if isinstance(background, float) else background82        self.background = torch.tensor(self.background).unsqueeze(1).unsqueeze(2)83 84    def forward(self, img: torch.Tensor) -> torch.Tensor:85        if img.shape[-3] == 3:86            return img87 88        alpha = img[..., 3, None, :, :]89 90        img[..., :3, :, :] *= alpha91 92        background = self.background.expand(-1, img.shape[-2], img.shape[-1])93        if background.ndim == 1:94            background = background[:, None, None]95        elif background.ndim == 2:96            background = background[None, :, :]97 98        img[..., :3, :, :] += (1.0 - alpha) * background99        return img[..., :3, :, :]100 101    def __repr__(self) -> str:102        return (103            f"{self.__class__.__name__}(" +104            f"background={self.background})"105        )106 107transform = transforms.Compose([108    Fit((384, 384)),109    transforms.ToTensor(),110    CompositeAlpha(0.5),111    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),112    transforms.CenterCrop((384, 384)),113])114 115model = timm.create_model(116    "vit_so400m_patch14_siglip_384.webli",117    pretrained=False,118    num_classes=9083,119) # type: VisionTransformer120 121safetensors.torch.load_model(model, "JTP_PILOT-e4-vit_so400m_patch14_siglip_384.safetensors")122model.eval()123 124with open("tagger_tags.json", "r") as file:125    tags = json.load(file) # type: dict126allowed_tags = list(tags.keys())127 128for idx, tag in enumerate(allowed_tags):129    allowed_tags[idx] = tag.replace("_", " ")130 131sorted_tag_score = {}132 133@spaces.GPU(duration=5)134def run_classifier(image, threshold):135    global sorted_tag_score136    img = image.convert('RGBA')137    tensor = transform(img).unsqueeze(0)138 139    with torch.no_grad():140        logits = model(tensor)141        probits = torch.nn.functional.sigmoid(logits[0])142        values, indices = probits.topk(250)143 144    tag_score = dict()145    for i in range(indices.size(0)):146        tag_score[allowed_tags[indices[i]]] = values[i].item()147    sorted_tag_score = dict(sorted(tag_score.items(), key=lambda item: item[1], reverse=True))148 149    return create_tags(threshold)150 151def create_tags(threshold):152    global sorted_tag_score153    filtered_tag_score = {key: value for key, value in sorted_tag_score.items() if value > threshold}154    text_no_impl = ", ".join(filtered_tag_score.keys())155    return text_no_impl, filtered_tag_score156    157 158with gr.Blocks(css=".output-class { display: none; }") as demo:159    gr.Markdown("""160    ## Joint Tagger Project: PILOT Demo161    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.162 163    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.164    """)165    with gr.Row():166        with gr.Column():167            image_input = gr.Image(label="Source", sources=['upload'], type='pil', height=512, show_label=False)168            threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")169        with gr.Column():170            tag_string = gr.Textbox(label="Tag String")171            label_box = gr.Label(label="Tag Predictions", num_top_classes=250, show_label=False)172 173    image_input.upload(174        fn=run_classifier,175        inputs=[image_input, threshold_slider],176        outputs=[tag_string, label_box]177    )178 179    threshold_slider.input(180        fn=create_tags,181        inputs=[threshold_slider],182        outputs=[tag_string, label_box]183    )184 185if __name__ == "__main__":186    demo.launch()