remove sex from non-sex prompts

crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/prompt_enrichment.rs

@@ -7,6 +7,8 @@ pub struct EnrichedPrompts {
   pub maybe_negative_prompt: Option<String>,
 }
 
+const DEFAULT_NEGATIVE_PROMPT : &str = "nudity, nsfw, naked, sex";
+
 /// Enrich the user-provided prompts.
 /// In the future, we can add prompt engineering to these.
 pub fn enrich_prompt(request: &EnqueueImageGenRequest) -> Option<EnrichedPrompts> {
@@ -15,10 +17,25 @@ pub fn enrich_prompt(request: &EnqueueImageGenRequest) -> Option<EnrichedPrompts
     Some(prompt) => prompt.to_string(),
   };
 
-  let mut maybe_negative_prompt = request.maybe_n_prompt.clone();
+  let mut maybe_negative_prompt = request.maybe_n_prompt
+      .as_ref()
+      .map(|prompt| prompt.to_string());
 
   let classification = classify_prompt(&positive_prompt);
 
+  if !classification.prompt_references_sex {
+    // If the prompt doesn't have sex terms in it, try to make sure it doesn't get generated.
+    match maybe_negative_prompt.as_deref() {
+      None => {
+        maybe_negative_prompt = Some(DEFAULT_NEGATIVE_PROMPT.to_string());
+      }
+      Some(negative_prompt) => {
+        let new_negative_prompt = format!("{}, {}", negative_prompt, DEFAULT_NEGATIVE_PROMPT);
+        maybe_negative_prompt = Some(new_negative_prompt);
+      }
+    }
+  }
+
   if classification.is_abusive() {
     // NB: Save the original prompt as the negative prompt so that we can study it.
     maybe_negative_prompt = Some(positive_prompt.clone());

crates/service/web/storyteller_web/src/util/classify_prompt.rs

@@ -142,6 +142,7 @@ mod tests {
       "Girl naked, masturbating, ahegao, squatting",
       "Girl naked, masturbating, ahegao, squatting, ultra realistic",
       "girl in a strapless dress in a feald of flower",
+      "two girls kissing in bed",
     ];
 
     for prompt in prompts {
@@ -159,6 +160,7 @@ mod tests {
 
   #[test]
   fn permissible_nsfw() {
+    // These are actual user prompts used as test cases
     // We're allowing NSFW for now, but this may change
     // These may not be underage or illegal
     let prompts = [
@@ -166,19 +168,29 @@ mod tests {
       "2B. white swimsuit, butt, black eyeliner, white droplets on ass, sunscreen bottle in hand",
       "Big boobs",
       "COCK COCK AND DICK BALLS HEELLPPPP",
+      "Girl nude, pink hair, pink pussy ",
       "Hot naked anime maids",
       "LOOK AT MY MASSIVE DICK",
-      "hot gaping pussy and ass legs wide open",
-      "hot cat girl pussy getting fucked by a dick",
-      "hot white girl gaping pussy and ass close up in the camera",
-      "minecraft creeper having sex in the pussy with another minecraft creeper",
       "Naked, cute, hot, nice boobs",
+      "black women, cum inflation, cute, sex, big boobs, big ass",
+      "hot cat girl pussy getting fucked by a dick",
+      "hot gaping pussy and ass legs wide open",
       "hot girl with a tight vagina",
       "hot hentai fuck",
       "hot naked girl",
+      "hot white girl gaping pussy and ass close up in the camera",
+      "hyper boobs, boobjob, cum, smaller male, black girl, white man, paizuri, sex, black female, inflation, expansion, cum inflation, dialouge",
+      "hyper boobs, boobjob, cum, smaller male, black girl, white man, paizuri, sex,",
+      "inflation, water inflation, oral, hose in mouth, water, orgasm, english dialouge, solo female, belly expansion, expansion, cum, inflation, water inflation, water inflation, water inflation, water inflation, human girl, orgasm, human girl inflation",
       "intercource between man and women",
+      "minecraft creeper having sex in the pussy with another minecraft creeper",
       "naked girl sitting in chair",
       "naked women fingring",
+      "sex",
+      "titfuck, huge boobs",
+
+      // Spanish
+      "Una mujer japonesa embarazada pero más femenina con cabello blanco, ojos azules, un traje rojo de artes marciales de cuerpo completo, una cicatriz en la mejilla, con un efecto de nieve en el fondo de ciudad, en estilo de dibujo manga",
     ];
 
     for prompt in prompts {
@@ -198,6 +210,7 @@ mod tests {
     // These are actual user prompts that use borderline language, but are safe
     let prompts = [
       "3 dancing people in a concert",
+      "A 12 ft tall terrifying shadow analog horror demon",
       "A boy, yellow short hair, blue eyes a star in it, hikaru kamiki",
       "A cute small fox in the woods",
       "A girl from Robotech universe, SDF1 mecha superdimensional fortess",
@@ -211,8 +224,7 @@ mod tests {
       "Overweight Taco Bell manager Female. Hair and a bun with braces, smiling. Black uniform with apron.",
       "Scary, fear, horror, creepypasta, spooky",
       "Tifa and aerith. Bikini. Laughing. Squirting sunscreen on each other",
-      "Uma mulher chorando", // NB: "A woman crying"
-      "Una mujer con piel de dragon, y cuernos", // NB: "A woman with dragon skin and horns"
+      "White and blue fox furry in a river",
       "an image of kojiro hiuga kicking tsubasa osora in the stomach",
       "boy with black and blue hair in a red shirt and jeans happy",
       "dead girl in a white dress with blood comeing frome her head in an abendend hospitel",
@@ -221,9 +233,14 @@ mod tests {
       "girl teen in a dress with horns and a tail in a feald of flowers",
       "girl with bloned hair wering a light blue dress whell drawing with a boy with black hair wering a red shirt siting nex to her",
       "girl with horns and a tail that has durdyblond hair wereing a pink sweter and riped jeans",
+      "morbidly obese blonde woman",
       "pichu witch pockimon",
-      "rio con personas", // NB: "river with people"
       "verey fansey dresses",
+
+      // Non-english
+      "Uma mulher chorando", // NB: "A woman crying"
+      "Una mujer con piel de dragon, y cuernos", // NB: "A woman with dragon skin and horns"
+      "rio con personas", // NB: "river with people"
     ];
 
     for prompt in prompts {