Update app.py

app.py

@@ -40,7 +40,8 @@ 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()
 
@@ -58,6 +59,7 @@ 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):
@@ -70,6 +72,7 @@ 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
@@ -80,6 +83,7 @@ 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
@@ -90,6 +94,7 @@ 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
@@ -103,7 +108,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)
@@ -111,6 +116,7 @@ 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 self.character_indexes and x[1] > character_thresh]
@@ -134,113 +140,13 @@ 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 = ["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",
+    predefined_tags = ["2024",
                        "mosaic_censoring"]
 
     with gr.Blocks(title=TITLE) as demo:
@@ -248,6 +154,7 @@ def main():
         gr.Markdown(DESCRIPTION)
 
         with gr.Row():
+            
             with gr.Column():
 
                 submit = gr.Button(
@@ -277,20 +184,37 @@ 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):
+        def process_images(files, model_repo, general_thresh, character_thresh, filter_tags, conditional_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(
@@ -301,17 +225,34 @@ def main():
                 )
         
                 # Construct the prompt based on the presence of character_part
+                prompt = ""
                 if character_part:
-                    prompts.append(f"{character_part}, {general_part}")
+                    prompt = f"{character_part}, {general_part}"
                 else:
-                    prompts.append(general_part)
+                    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)
 
             # 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],
+            inputs=[image_files, model_repo, general_thresh, character_thresh, filter_tags, conditional_tags],
             outputs=output
         )