From 1af9045ded96fff09aeb348c0952749e1dd528fb Mon Sep 17 00:00:00 2001
From: skier233 <39396856+skier233@users.noreply.github.com>
Date: Thu, 29 Aug 2024 21:54:14 -0400
Subject: [PATCH] [AI Tagger] Minor fix (#423)

---
 plugins/AITagger/ai_tagger.py     | 11 +++++++----
 plugins/AITagger/tag_mappings.csv |  3 ++-
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/plugins/AITagger/ai_tagger.py b/plugins/AITagger/ai_tagger.py
index 0d2f6268..cb86fd5b 100644
--- a/plugins/AITagger/ai_tagger.py
+++ b/plugins/AITagger/ai_tagger.py
@@ -172,10 +172,13 @@ async def __tag_images(images):
         finally:
             increment_progress()
             for temp_file in temp_files:
-                if os.path.isdir(temp_file):
-                    shutil.rmtree(temp_file)
-                else:
-                    os.remove(temp_file)
+                try:
+                    if os.path.isdir(temp_file):
+                        shutil.rmtree(temp_file)
+                    else:
+                        os.remove(temp_file)
+                except Exception as e:
+                    log.debug(f"Failed to remove temp file {temp_file}: {e}")
 
 # ----------------- Scene Processing -----------------
 
diff --git a/plugins/AITagger/tag_mappings.csv b/plugins/AITagger/tag_mappings.csv
index f65efdcf..593cfe70 100644
--- a/plugins/AITagger/tag_mappings.csv
+++ b/plugins/AITagger/tag_mappings.csv
@@ -38,6 +38,7 @@ Vibrating,Vibrating_AI,10,6,20s,0.5
 Ass,Ass_AI,-1,6,20s,0.5
 Asshole,Asshole_AI,-1,6,20s,0.5
 Anal Gape,Anal Gape_AI,10,6,20s,0.5
+Balls,Balls_AI,-1,6,20s,0.5
 Boobs,Boobs_AI,-1,6,20s,0.5
 Cum,Cum_AI,10,6,20s,0.5
 Dick,Dick_AI,-1,6,20s,0.5
@@ -53,4 +54,4 @@ Pussy,Pussy_AI,-1,6,20s,0.5
 Pussy Gape,Pussy Gape_AI,10,6,20s,0.5
 Spit,Spit_AI,10,6,20s,0.5
 Oiled,Oiled_AI,10,6,20s,0.5
-Wet (Water),Wet (Water)_AI,10,6,20s,0.5
+Wet (Water),Wet (Water)_AI,10,6,20s,0.5
\ No newline at end of file