Merge pull request divamgupta#295 from AEtheve/add-nsfw-filter

Add nsfw filter

Author: divamgupta

electron_app/package.json

@@ -32,6 +32,7 @@
  "electron-notarize": "^1.2.1",
  "electron-settings": "^4.0.2",
  "javascript-time-ago": "^2.3.13",
+ "nsfwjs": "^2.4.2",
  "konva": "^8.3.13",
  "v-click-outside": "^3.1.2",
  "vue": "^2.6.11",

electron_app/src/StableDiffusion.vue

@@ -5,6 +5,7 @@
 
 import { send_to_py } from "./py_vue_bridge.js"
 import {get_tokens} from './clip_tokeniser/clip_encoder.js'
+const nsfwjs = require('nsfwjs');
 
 function remove_non_ascii(str) {
  
@@ -27,6 +28,7 @@ export default {
  return {
  is_backend_loaded : false,
  is_input_avail : false,
+ nsfw_filter: true,
  model_loading_msg : "",
  model_loading_title : "",
  loading_percentage : -1 , 
@@ -50,11 +52,35 @@ export default {
  }
  if(msg_code == "nwim"){
  let impath = msg.substring(5).trim()
- if(this.attached_cbs){
- if(this.attached_cbs.on_img)
- this.attached_cbs.on_img(impath);
- }
 
+ const img = new Image();
+ img.src = "file://" + impath;
+
+ img.attached_cbs = this.attached_cbs;
+
+
+ img.onload = () => {
+ if (this.nsfw_filter == false) {
+ if (img.attached_cbs) {
+ if (img.attached_cbs.on_img)
+ img.attached_cbs.on_img(impath);
+ }
+ }
+ nsfwjs.load()
+ .then(model => model.classify(img))
+ .then(predictions => {
+ switch (predictions[0].className) {
+ case 'Hentai':
+ case 'Porn':
+ case 'Sexy':
+ impath = "nsfw";
+ }
+ if (img.attached_cbs) {
+ if (img.attached_cbs.on_img)
+ img.attached_cbs.on_img(impath);
+ }
+ });
+ }
  }
 
  if(msg_code == "mdvr"){

electron_app/src/assets/nsfw.png

@@ -18,7 +18,8 @@
  </div>
 
 
- <img @click="open_image_popup( path )" class="gal_img" :src="'file://' + path " style="max-height: 100% ; max-width: 100%;" >
+ <img v-if="path != 'nsfw'" @click="open_image_popup( path )" class="gal_img" :src="'file://' + path " style="max-height: 100% ; max-width: 100%;" >
+ <img v-else @click="open_image_popup( path )" class="gal_img" :src="require('@/assets/nsfw.png')" style="max-height: 100% ; max-width: 100%;" >
  <br>
  <div v-if="!hide_extra_save_button" @click="save_image(path)" class="l_button">Save Image</div>
  </div>