[AI Tagger] Minor fix (#423)

stashapp · Aug 30, 2024 · 1af9045 · 1af9045
1 parent 97999c1
commit 1af9045
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git 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
@@ -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