Update app.py

app.py

@@ -220,25 +220,19 @@ def main():
             with gr.Column():
                 output = gr.Textbox(label="Output", lines=10)
 
-        def process_images(files, model_repo, general_thresh, character_thresh, filter_tags, custom_tags):
+        def process_images(files, model_repo, general_thresh, character_thresh, filter_tags, custom_tags_input):
             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 custom tags
-            custom_tag_list = [tag.strip().lower() for tag in custom_tags.split(",")]
-            
-            # Define the fallback tags for each custom tag
-            fallback_tags = {
-                "shy, happy": "smile",
-                "happy": "joy",
-                "cute": "adorable",
-                "energetic": "vibrant",
-                "blush": "flushed",
-                # Add more custom to fallback mappings as needed
-            }
+        
+            # Parse custom tags and their fallback pairs
+            fallback_tags = {}
+            for pair in custom_tags_input.split(","):
+                if ":" in pair:
+                    tag, fallback = pair.split(":")
+                    fallback_tags[tag.strip().lower()] = fallback.strip().lower()
         
             # Generate formatted output
             prompts = []
@@ -251,22 +245,20 @@ def main():
                     tag.replace('_', ' ') for tag in general_tags if tag.lower() not in filter_set
                 )
                 
-                # Check for missing custom tags and apply corresponding fallback tags
-                for custom_tag in custom_tag_list:
-                    if custom_tag not in (general_tags + character_tags):
-                        # Append the corresponding fallback tag for the missing custom tag
-                        if custom_tag in fallback_tags:
-                            general_part += ", " + fallback_tags[custom_tag]
-                
-                # Construct the prompt based on the presence of character_part
-                if character_part:
-                    prompts.append(f"{character_part}, {general_part}")
-                else:
-                    prompts.append(general_part)
+                # Check if custom tags are missing and apply fallback tags
+                all_tags = set(general_tags + character_tags)
+                for tag, fallback in fallback_tags.items():
+                    if tag not in all_tags:
+                        all_tags.add(fallback)
+        
+                # Construct the final prompt
+                final_tags = ", ".join(tag.replace('_', ' ') for tag in all_tags if tag.lower() not in filter_set)
+                prompts.append(final_tags)
         
             # 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, custom_tags],