app.py

import gradio as gr
import open_clip
import torch

# Load model and tokenizer
DEVICE='cpu'
model, preprocess = open_clip.create_model_from_pretrained('hf-hub:woweenie/open-clip-vit-h-nsfw-finetune', device=DEVICE)
tokenizer = open_clip.get_tokenizer('hf-hub:woweenie/open-clip-vit-h-nsfw-finetune')

# Define labels
type_labels = ['2.5d render', '3d render', 'photograph', 'anime drawing', 'drawing', 'illustration', 'painting', 'pre-raphaelite painting', 'concept artwork', 'screenshot']
scene_labels = ['in an airport', 'in the bath', 'on a bed', 'in bed', 'in a bedroom', 'at the beach', 'on a boat', 'in a tent', 'in a car', 'on a chair', 'in the city', 'in a dressing room', 'on the floor', 'at the gym', 'in a hotel room', 'in a kitchen', 'in a living room', 'in an office', 'by a harbor', 'on a bench', 'in a park', 'by a piano', 'on a forest road', 'in a forest', 'in a garden', 'at a lake', 'on the grass', 'on the ground', 'on a paved surface', 'outdoors, on a rock', 'outdoors, on a rug', 'outdoors, on a towel', 'in a photo studio', 'at the pool', 'at a river', 'on a road', 'by the sea', 'showering', 'in the shower', 'on a stool', 'on a rug', 'on a rock', 'on a sofa', 'on a table', 'at a table', 'in a store', 'on snow', 'by a waterfall', 'with a water feature', 'on a windowsill']
expression_labels = ['scared', 'annoyed', 'aroused', 'bored', 'confident', 'distracted', 'dominating', 'embarrassed', 'scared', 'laughing', 'shy', 'orgasm']
clothing_labels = ['a bikini that is too small', 'bikini bottoms', 'a bikini top', 'a bikini', 'a bodysuit', 'a bra', 'a crop top', 'a dress', 'garters', 'glasses', 'goggles', 'gym shorts', 'a halter top', 'a hat', 'a handbra', 'a hoodie', 'a jacket', 'jeans', 'a jumper', 'a gown', 'a lace-up top', 'leggings', 'lingerie', 'a long sleeved top', 'a off-shoulder top', 'a nightgown', 'a coat', 'overalls', 'pink pajamas', 'pajamas', 'panties', 'pantyhose', 'a t-shirt', 'a robe', 'a bathrobe', 'a piece of fabric', 'a scarf', 'a shirt', 'shorts', 'a skirt', 'a sleeveless top', 'a slip', 'sneakers', 'tube socks', 'a sports bra', 'sunglasses', 'sweatpants', 'a one piece swimsuit', 'a t-shirt', 'a tank top', 'a tied shirt', 'a top', 'long pants', 'a wetsuit', 'a backpack', 'high hem', 'see-through', 'short', 'tight']
clothing_labels = ['wearing ' + label for label in clothing_labels]

def process_image_separate_tags(image):
    # Preprocess image
    image = preprocess(image).unsqueeze(0).to(DEVICE)

    # Tokenize labels
    type_text = tokenizer(type_labels).to(DEVICE)
    scene_text = tokenizer(scene_labels).to(DEVICE) 
    expression_text = tokenizer(expression_labels).to(DEVICE)
    clothing_text = tokenizer(clothing_labels).to(DEVICE)

    with torch.no_grad():
        # Encode image and text
        image_features = model.encode_image(image)
        type_text_features = model.encode_text(type_text)
        scene_text_features = model.encode_text(scene_text)
        expression_text_features = model.encode_text(expression_text)
        clothing_text_features = model.encode_text(clothing_text)

        # Normalize features
        image_features /= image_features.norm(dim=-1, keepdim=True)
        type_text_features /= type_text_features.norm(dim=-1, keepdim=True)
        scene_text_features /= scene_text_features.norm(dim=-1, keepdim=True)
        expression_text_features /= expression_text_features.norm(dim=-1, keepdim=True)
        clothing_text_features /= clothing_text_features.norm(dim=-1, keepdim=True)

        # Calculate cosine similarities and apply softmax
        # Using temperature parameter to control the "sharpness" of the distribution
        temperature = 0.1  # Lower values make the distribution more peaked
        type_text_probs = torch.softmax(image_features @ type_text_features.T / temperature, dim=-1)
        scene_text_probs = torch.softmax(image_features @ scene_text_features.T / temperature, dim=-1)
        expression_text_probs = torch.softmax(image_features @ expression_text_features.T / temperature, dim=-1)
        clothing_text_probs = torch.softmax(image_features @ clothing_text_features.T / temperature, dim=-1)

    # Convert to dictionaries
    type_results = {label: float(type_text_probs[0][i]) for i, label in enumerate(type_labels)}
    scene_results = {label: float(scene_text_probs[0][i]) for i, label in enumerate(scene_labels)}
    expression_results = {label: float(expression_text_probs[0][i]) for i, label in enumerate(expression_labels)}
    clothing_results = {label: float(clothing_text_probs[0][i]) for i, label in enumerate(clothing_labels)}

    return type_results, scene_results, expression_results, clothing_results

def process_image_combined_tags(image):
    # Preprocess image
    image = preprocess(image).unsqueeze(0).to(DEVICE)

    # Tokenize labels
    all_text = tokenizer(type_labels + scene_labels + expression_labels + clothing_labels).to(DEVICE)

    with torch.no_grad():
        # Encode image and text
        image_features = model.encode_image(image)
        all_text_features = model.encode_text(all_text)

        # Normalize features
        image_features /= image_features.norm(dim=-1, keepdim=True)
        all_text_features /= all_text_features.norm(dim=-1, keepdim=True)

        # Calculate cosine similarities and apply softmax
        # Using temperature parameter to control the "sharpness" of the distribution
        temperature = 0.1  # Lower values make the distribution more peaked
        cosine_similarities = image_features @ all_text_features.T
        all_text_probs = torch.softmax(cosine_similarities / temperature, dim=-1)

    # Convert to dictionaries
    all_results = {label: float(all_text_probs[0][i]) for i, label in enumerate(type_labels + scene_labels + expression_labels + clothing_labels)}

    return all_results

# Create Gradio interface
iface = gr.Blocks(title="NSFW Tagging with Finetuned CLIP")

with iface:
    gr.Markdown("# NSFW Tagging with Finetuned CLIP")
    gr.Markdown("Upload an image to analyze its content across multiple NSFW categories.")
    gr.Markdown("Uses [woweenie/open-clip-vit-h-nsfw-finetune](https://huggingface.co/woweenie/open-clip-vit-h-nsfw-finetune) finetuned on NSFW images.")
    gr.Markdown("Disclaimer: This model is not perfect and may make mistakes. Use at your own risk.")

    with gr.Tabs():
        with gr.Tab("Categorical Predictions"):
            with gr.Row():
                image_input2 = gr.Image(type="pil", label="Upload Image")
            with gr.Row():
                type_output = gr.Label(label="Predicted Type", num_top_classes=3)
                scene_output = gr.Label(label="Predicted Scene", num_top_classes=10) 
                expression_output = gr.Label(label="Predicted Expression", num_top_classes=3)
                clothing_output = gr.Label(label="Predicted Clothing", num_top_classes=5)
            predict_btn2 = gr.Button("Analyze")
            predict_btn2.click(
                fn=process_image_separate_tags,
                inputs=image_input2,
                outputs=[type_output, scene_output, expression_output, clothing_output]
            )

        with gr.Tab("Combined Predictions"):
            with gr.Row():
                image_input1 = gr.Image(type="pil", label="Upload Image")
            with gr.Row():
                combined_output = gr.Label(label="Predicted Tags", num_top_classes=10)
            predict_btn1 = gr.Button("Analyze")
            predict_btn1.click(
                fn=process_image_combined_tags,
                inputs=image_input1,
                outputs=combined_output
            )



if __name__ == "__main__":
    iface.launch()