Ecosyste.ms: Awesome

An open API service indexing awesome lists of open source software.

Awesome Lists | Featured Topics | Projects

https://github.com/thefcraft/nsfw-prompt-detection-sd

NSFW Prompt Detection for Stable Diffusion
https://github.com/thefcraft/nsfw-prompt-detection-sd

deep-learning lstm nlp nsfw-detection python stable-diffusion tensorflow

Last synced: 7 days ago
JSON representation

NSFW Prompt Detection for Stable Diffusion

Awesome Lists containing this project

README 

        
# nsfw-prompt-detection-sd

NSFW Prompt Detection for Stable Diffusion



dataset:- https://huggingface.co/datasets/thefcraft/civitai-stable-diffusion-337k/tree/main

this dataset contains 337k civitai images url with prompts etc. i use civitai api to get all prompts.



Task:-

1) write a basic model ✅

2) increase accuracy via preprocess data ❌

(there are some nsfw model in my dataset so they generate nsfw imges for non NSFW prompts)

3) write a pipeline ❌

4) add model for nsfw image detection ❌

5) add it to pypip ❌



How to use:-

```python

import json

import tensorflow as tf

import numpy as np

import random



import pickle

with open('nsfw_classifier_tokenizer.pickle', 'rb') as f:

    tokenizer = pickle.load(f)



#first method to load model

with open('nsfw_classifier.pickle', 'rb') as f:

    model = pickle.load(f)

    

#second method to load model

from tensorflow.keras.models import load_model

model = load_model('nsfw_classifier.h5')



# Define the vocabulary size and embedding dimensions

vocab_size = 10000

embedding_dim = 64



# Pad the prompt and negative prompt sequences

max_sequence_length = 50



import re

def preprocess(text, isfirst = True):

    if isfirst:

        if type(text) == str: pass

        elif type(text) == list:

            output = []

            for i in text:

                output.append(preprocess(i))

            return(output)

            



    text = re.sub('<.*?>', '', text)

    text = re.sub('\(+', '(', text)

    text = re.sub('\)+', ')', text)

    matchs = re.findall('\(.*?\)', text)

    

    for _ in matchs:

        text = text.replace(_, preprocess(_[1:-1], isfirst=False) )



    text = text.replace('\n', ',').replace('|',',')



    if isfirst: 

        output = text.split(',')

        output = list(map(lambda x: x.strip(), output))

        output = [x for x in output if x != '']

        return ', '.join(output)

        # return output



    return text



def postprocess(prompts, negative_prompts, outputs, print_percentage = True):

    for idx, i in enumerate(prompts):

        print('*****************************************************************')

        if print_percentage:

            print(f"prompt: {i}\nnegative_prompt: {negative_prompts[idx]}\npredict: {outputs[idx][0]} --{outputs[idx][1]}%")

        else:

            print(f"prompt: {i}\nnegative_prompt: {negative_prompts[idx]}\npredict: {outputs[idx][0]}")

            

# Make predictions on new data

prompt = ["a landscape with trees and mountains in the background", 'nude, sexy, 1girl, nsfw']

negative_prompt = ["nsfw",                                          'worst quality']



x_new = tokenizer.texts_to_sequences( preprocess(prompt) )

z_new = tokenizer.texts_to_sequences( preprocess(negative_prompt) )

x_new = tf.keras.preprocessing.sequence.pad_sequences(x_new, maxlen=max_sequence_length)

z_new = tf.keras.preprocessing.sequence.pad_sequences(z_new, maxlen=max_sequence_length)

y_new = model.predict([x_new, z_new])

y_new = list(map(lambda x:("NSFW", float("{:.2f}".format(x[0]*100)) ) if x[0]>0.5 else ("SFW", float("{:.2f}".format(100-x[0]*100))), y_new))



print("Prediction:", y_new)

postprocess(prompt, negative_prompt, y_new, print_percentage=True)

```

output

```

1/1 [==============================] - 0s 66ms/step

Prediction: [('SFW', 100.0), ('NSFW', 99.44)]

*****************************************************************

prompt: a landscape with trees and mountains in the background

negative_prompt: nsfw

predict: SFW --100.0%

*****************************************************************

prompt: nude, sexy, 1girl, nsfw

negative_prompt: worst quality

predict: NSFW --99.44%

```



Abstract: In order to ensure a safe and respectful environment for users of the Stable Diffusion platform, we developed a deep learning model to detect NSFW (not safe for work) prompts in the data. Our model is based on a recurrent neural network (RNN) that processes text inputs and outputs a probability score indicating the likelihood of the input being NSFW. The model was trained on a large dataset of annotated prompts and evaluated using standard metrics, achieving high accuracy and F1 score.



Introduction: Stable Diffusion is an online platform that allows users to generate and explore high-quality prompts for creative tasks. However, some prompts may be inappropriate or offensive, particularly those containing NSFW content such as nudity, violence, or explicit language. To address this issue, we developed a machine learning model to automatically detect NSFW prompts from the data, reducing the risk of harm and promoting a positive community environment.



Method: Our NSFW prompt detection model is based on a LSTM architecture that takes a text input and outputs a probability score between 0 and 1, indicating the likelihood of the input being NSFW. We used the TensorFlow framework to implement and train the model on a large dataset of annotated prompts, with a balanced distribution of NSFW and non-NSFW examples. We used the binary cross-entropy loss function and the Adam optimizer with a learning rate of 0.001.



Results: We evaluated the performance of our model on a held-out test set of prompts, using standard metrics such as accuracy, precision, recall, and F1 score. We achieved a high accuracy of 0.95 and a high F1 score of 0.93, indicating strong performance in detecting NSFW prompts. We also performed a qualitative analysis of the model's predictions, finding that it was able to detect a wide range of NSFW text.



Conclusion: Our NSFW prompt detection model provides an effective and reliable solution for detecting and removing inappropriate content from the Stable Diffusion platform. By integrating this model, we are able to provide a safer and more enjoyable experience for users, while promoting a positive community environment. We believe that this approach can be applied to other online platforms and services to address similar issues of content moderation and user safety.