Update app.py

app.py

@@ -40,8 +40,7 @@ LABEL_FILENAME = "selected_tags.csv"
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument("--score-slider-step", type=float, default=0.05)
-    parser.add_argument("--score-general-threshold", type=float, 
- default=0.25)
+    parser.add_argument("--score-general-threshold", type=float, default=0.25)
     parser.add_argument("--score-character-threshold", type=float, default=1.0)
     return parser.parse_args()
 
@@ -59,7 +58,6 @@ class Predictor:
     def download_model(self, model_repo):
         csv_path = huggingface_hub.hf_hub_download(model_repo, LABEL_FILENAME, use_auth_token=HF_TOKEN)
         model_path = huggingface_hub.hf_hub_download(model_repo, MODEL_FILENAME, use_auth_token=HF_TOKEN)
- 
         return csv_path, model_path
 
     def load_model(self, model_repo):
@@ -72,7 +70,6 @@ class Predictor:
 
         model = rt.InferenceSession(model_path)
         _, height, width, _ = model.get_inputs()[0].shape
-      
         self.model_target_size = height
         self.last_loaded_repo = model_repo
         self.model = model
@@ -83,7 +80,6 @@ class Predictor:
     
         # Ensure the input image has an alpha channel for compositing
         if image.mode != "RGBA":
-   
             image = image.convert("RGBA")
     
         # Composite the input image onto the canvas
@@ -94,7 +90,6 @@ class Predictor:
 
         # Resize the image to a square of size (model_target_size x model_target_size)
         max_dim = max(image.size)
-  
         padded_image = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
         pad_left = (max_dim - image.width) // 2
         pad_top = (max_dim - image.height) // 2
@@ -108,7 +103,7 @@ class Predictor:
 
     def predict(self, images, model_repo, general_thresh, character_thresh):
         self.load_model(model_repo)
-        results =
+        results = []
 
         for image in images:
             image = self.prepare_image(image)
@@ -116,10 +111,9 @@ class Predictor:
             label_name = self.model.get_outputs()[0].name
             preds = self.model.run([label_name], {input_name: image})[0]
 
-   
             labels = list(zip(self.tag_names, preds[0].astype(float)))
             general_res = [x[0] for i, x in enumerate(labels) if i in self.general_indexes and x[1] > general_thresh]
-            character_res = [x[0] for i, x in enumerate(labels) if i in the character_indexes and x[1] > character_thresh]
+            character_res = [x[0] for i, x in enumerate(labels) if i in self.character_indexes and x[1] > character_thresh]
             results.append((general_res, character_res))
 
         return results
@@ -140,13 +134,113 @@ def main():
         CONV_MODEL_DSV2_REPO,
         CONV2_MODEL_DSV2_REPO,
         VIT_MODEL_DSV2_REPO,
-    
         # ---
         SWINV2_MODEL_IS_DSV1_REPO,
         EVA02_LARGE_MODEL_IS_DSV1_REPO,
     ]
 
-    predefined_tags = ["2024",
+    predefined_tags = ["loli",
+                       "oppai_loli",
+                       "2024",
+                       "2023",
+                       "2025",
+                       "2022",
+                       "2021",
+                       "onee-shota",
+                       "incest",
+                       "furry",
+                       "twitter_strip_game_(meme)",
+                       "like_and_retweet",
+                       "furry_female",
+                       "realistic",
+                       "egg_vibrator",
+                       "tongue_piercing",
+                       "handheld_game_console",
+                       "game_controller",
+                       "nintendo_switch",
+                       "talking",
+                       "swastika",
+                       "character_name",
+                       "vibrator",
+                       "black-framed_eyewear",
+                       "heterochromia",
+                       "controller",
+                       "remote_control_vibrator",
+                       "vibrator_under_clothes",
+                       "thank_you",
+                       "vibrator_cord",
+                       "shota",
+                       "male_focus",
+                       "signature",
+                       "web_address",
+                       "censored_nipples",
+                       "rhodes_island_logo_(arknights)",
+                       "gothic_lolita",
+                       "glasses",
+                       "reference_inset",
+                       "twitter_logo",
+                       "mother_and_daughter",
+                       "holding_controller",
+                       "holding_game_controller",
+                       "baby",
+                       "heart_censor",
+                       "pixiv_username",
+                       "korean_text",
+                       "pixiv_logo",
+                       "greyscale_with_colored_background",
+                       "water_bottle",
+                       "body_writing",
+                       "used_condom",
+                       "multiple_condoms",
+                       "condom_belt", 
+                       "holding_phone",
+                       "multiple_views",
+                       "phone",
+                       "cellphone",
+                       "zoom_layer",
+                       "smartphone",
+                       "lolita_hairband",
+                       "lactation",
+                       "otoko_no_ko",
+                       "minigirl",
+                       "babydoll",
+                       "domino_mask",
+                       "pixiv_id",
+                       "qr_code",
+                       "monochrome",
+                       "trick_or_treat",
+                       "happy_birthday",
+                       "lolita_fashion",
+                       "arrow_(symbol)",
+                       "happy_new_year", 
+                       "dated", 
+                       "thought_bubble",
+                       "greyscale",
+                       "speech_bubble",
+                       "mask",
+                       "bottle",
+                       "holding_bottle",
+                       "milk",
+                       "milk_bottle",
+                       "english_text",
+                       "copyright_name",
+                       "twitter_username",
+                       "fanbox_username",
+                       "patreon_username",
+                       "patreon_logo",
+                       "cover", 
+                       "signature",
+                       "content_rating",
+                       "cover_page", 
+                       "doujin_cover", 
+                       "sex",
+                       "artist_name",
+                       "watermark",
+                       "censored",
+                       "bar_censor",
+                       "blank_censor",
+                       "blur_censor",
+                       "light_censor",
                        "mosaic_censoring"]
 
     with gr.Blocks(title=TITLE) as demo:
@@ -154,7 +248,6 @@ def main():
         gr.Markdown(DESCRIPTION)
 
         with gr.Row():
-            
             with gr.Column():
 
                 submit = gr.Button(
@@ -184,37 +277,20 @@ def main():
                         placeholder="Add tags to filter out (e.g., winter, red, from above)",
                         lines=9
                     )
-                    
-                    conditional_tags = gr.Textbox(
-                        label="Conditional Tag Rules",
-                        placeholder="Enter tag rules (e.g., sun: hot,day)",
-                        lines=3,
-                    )
 
 
             with gr.Column():
                 output = gr.Textbox(label="Output", lines=10)
 
-        def process_images(files, model_repo, general_thresh, character_thresh, filter_tags, conditional_tags):
+        def process_images(files, model_repo, general_thresh, character_thresh, filter_tags):
             images = [Image.open(file.name) for file in files]
             results = predictor.predict(images, model_repo, general_thresh, character_thresh)
 
             # Parse filter tags
             filter_set = set(tag.strip().lower() for tag in filter_tags.split(","))
-            
-            # Parse conditional tag rules
-            tag_rules = {}
-            if conditional_tags:
-                for rule in conditional_tags.splitlines():
-                    if ":" in rule:
-                        trigger_tag, tags_to_add = rule.split(":", 1)
-                        tag_rules[trigger_tag.strip().lower()] = [
-                            tag.strip() for tag in tags_to_add.split(",")
-                        ]
-
 
             # Generate formatted output
-            prompts =
+            prompts = []
             for i, (general_tags, character_tags) in enumerate(results):
                 # Replace underscores with spaces for both character and general tags
                 character_part = ", ".join(
@@ -225,34 +301,17 @@ def main():
                 )
         
                 # Construct the prompt based on the presence of character_part
-                prompt = ""
                 if character_part:
-                    prompt = f"{character_part}, {general_part}"
+                    prompts.append(f"{character_part}, {general_part}")
                 else:
-                    prompt = general_part
-                    
-                # Apply conditional tag rules
-                found_trigger = False
-                for trigger_tag, tags_to_add in tag_rules.items():
-                    if trigger_tag in prompt.lower():
-                        prompt += ", " + ", ".join(tags_to_add)
-                        found_trigger = True
-                        break
-                
-                if not found_trigger:
-                    for trigger_tag, tags_to_add in tag_rules.items():
-                         if trigger_tag not in prompt.lower():
-                            prompt += ", " + ", ".join(tags_to_add)
-                            break # Only apply the first rule that matches
-
-                prompts.append(prompt)
+                    prompts.append(general_part)
 
             # Join all prompts with blank lines
             return "\n\n".join(prompts)
 
         submit.click(
             process_images,
-            inputs=[image_files, model_repo, general_thresh, character_thresh, filter_tags, conditional_tags],
+            inputs=[image_files, model_repo, general_thresh, character_thresh, filter_tags],
             outputs=output
         )