Merge pull request #852 from zanllp/feat/auto-tagging

feat: implement auto-tagging

Author: zanllp

javascript/index.js

@@ -13,7 +13,7 @@ Promise.resolve().then(async () => {
     <link rel="icon" href="/favicon.ico" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
     <title>Infinite Image Browsing</title>
-    <script type="module" crossorigin src="/infinite_image_browsing/fe-static/assets/index-a7581aec.js"></script>
+    <script type="module" crossorigin src="/infinite_image_browsing/fe-static/assets/index-6c47ee34.js"></script>
     <link rel="stylesheet" href="/infinite_image_browsing/fe-static/assets/index-a44325c7.css">
   </head>
 

scripts/iib/api.py

@@ -313,6 +313,10 @@ class AppFeSettingReq(BaseModel):
     async def app_fe_setting(req: AppFeSettingReq):
         conn = DataBase.get_conn()
         GlobalSetting.save_setting(conn, req.name, req.value)
+        # 如果更新的是自动标签规则，重新加载
+        if req.name == "auto_tag_rules":
+            from scripts.iib.auto_tag import AutoTagMatcher
+            AutoTagMatcher.reload_rules(conn)
 
     class AppFeSettingDelReq(BaseModel):
         name: str

scripts/iib/auto_tag.py

@@ -0,0 +1,101 @@
+from typing import List, Dict, Any, Optional
+import re
+from scripts.iib.db.datamodel import GlobalSetting, Tag, ImageTag, DataBase
+from scripts.iib.logger import logger
+from scripts.iib.parsers.model import ImageGenerationParams
+
+class AutoTagMatcher:
+    _instance: Optional['AutoTagMatcher'] = None
+    
+    def __init__(self, conn):
+        self.conn = conn
+        self.rules = self.load_rules()
+
+    @classmethod
+    def get_instance(cls, conn) -> 'AutoTagMatcher':
+        """获取全局单例实例"""
+        if cls._instance is None:
+            cls._instance = cls(conn)
+        else:
+            # 更新连接
+            cls._instance.conn = conn
+        return cls._instance
+    
+    @classmethod
+    def reload_rules(cls, conn) -> None:
+        """重新加载规则"""
+        if cls._instance is not None:
+            cls._instance.conn = conn
+            cls._instance.rules = cls._instance.load_rules()
+
+    def load_rules(self) -> List[Dict[str, Any]]:
+        try:
+            setting = GlobalSetting.get_setting(self.conn, "auto_tag_rules")
+            print("Loaded auto tag rules:", setting)
+            if setting:
+                return setting
+        except Exception as e:
+            logger.error(f"Error loading auto tag rules: {e}")
+        return []
+
+    def match(self, params: ImageGenerationParams, rule: Dict[str, Any]) -> bool:
+        # params is ImageGenerationParams
+        # rule structure: { "tag": "TagName", "filters": [...] }
+        # filter structure: { "field": "pos_prompt", "operator": "contains", "value": "foo" }
+        
+        filters = rule.get("filters", [])
+        if not filters:
+            return False
+
+        for filter in filters:
+            field = filter.get("field")
+            operator = filter.get("operator")
+            value = filter.get("value")
+            
+            target_value = ""
+            if field == "pos_prompt":
+                target_value = ",".join(params.pos_prompt) if isinstance(params.pos_prompt, list) else str(params.pos_prompt)
+            elif field == "neg_prompt":
+                 # Assuming neg_prompt might be in meta or extra, but params usually has pos_prompt. 
+                 # Let's check where negative prompt is usually stored. 
+                 # In api.py: params = parse_generation_parameters(content)
+                 # parse_generation_parameters returns dict with "meta", "pos_prompt".
+                 # "meta" usually contains "Negative prompt".
+                 target_value = params.meta.get("Negative prompt", "")
+            else:
+                # Try to find in meta
+                target_value = params.meta.get(field, "")
+                if not target_value and field in params.extra:
+                     target_value = params.extra.get(field, "")
+            
+            target_value = str(target_value)
+            
+            if operator == "contains":
+                if value.lower() not in target_value.lower():
+                    return False
+            elif operator == "equals":
+                if value.lower() != target_value.lower():
+                    return False
+            elif operator == "regex":
+                try:
+                    if not re.search(value, target_value, re.IGNORECASE):
+                        return False
+                except:
+                    return False
+            # Add more operators if needed
+            
+        return True
+
+    def apply(self, img_id: int, params: Any):
+        if not self.rules:
+            return
+            
+        for rule in self.rules:
+            try:
+                if self.match(params, rule):
+                    tag_name = rule.get("tag")
+                    if tag_name:
+                        tag = Tag.get_or_create(self.conn, tag_name, "custom")
+                        ImageTag(img_id, tag.id).save_or_ignore(self.conn)
+            except Exception as e:
+                logger.error(f"Error applying auto tag rule {rule}: {e}")

scripts/iib/db/datamodel.py

@@ -908,7 +908,7 @@ def get_cached_media_files(cls, conn, folder_path):
         else:
             return []
 
-
+# Global settings storage, also use as key-value store
 class GlobalSetting:
     @classmethod
     def create_table(cls, conn):

scripts/iib/db/update_image_data.py

@@ -17,6 +17,7 @@
 from scripts.iib.logger import logger
 from scripts.iib.parsers.index import parse_image_info
 from scripts.iib.plugin import plugin_inst_map
+from scripts.iib.auto_tag import AutoTagMatcher
 
 # 定义一个函数来获取图片文件的EXIF数据
 def get_exif_data(file_path):
@@ -222,4 +223,6 @@ def build_single_img_idx(conn, file_path, is_rebuild, safe_save_img_tag):
         safe_save_img_tag(ImageTag(img.id, tag.id))
     for k in pos:
         tag = Tag.get_or_create(conn, k, "pos")
-        safe_save_img_tag(ImageTag(img.id, tag.id))
+        safe_save_img_tag(ImageTag(img.id, tag.id))
+    
+    AutoTagMatcher.get_instance(conn).apply(img.id, parsed_params)

scripts/iib/parsers/comfyui.py

@@ -32,7 +32,7 @@ def parse(clz, img, file_path):
                 params = get_comfyui_exif_data(img)
                 info = comfyui_exif_data_to_str(params)
         except Exception as e:
-            logger.error("parse comfyui image failed. prompt:")
+            logger.error("parse comfyui image failed. prompt:", exc_info=e)
             logger.error(img.info.get("prompt"))
             return ImageGenerationInfo(
                 params=ImageGenerationParams(

scripts/iib/tool.py

@@ -460,11 +460,20 @@ def get_comfyui_exif_data(img: Image):
     meta["Model"] = data[KSampler_entry["model"][0]]["inputs"].get("ckpt_name")
     meta["Source Identifier"] = "ComfyUI"
     def get_text_from_clip(idx: str) :
-        inputs = data[idx]["inputs"]
-        text = inputs["text"] if "text" in inputs else inputs["t5xxl"]
-        if isinstance(text, list): # type:CLIPTextEncode (NSP) mode:Wildcards
-            text = data[text[0]]["inputs"]["text"]
-        return text.strip()
+        try:
+            inputs = data[idx]["inputs"]
+            if "text" in inputs:
+                text = inputs["text"]
+            elif "t5xxl" in inputs:
+                text = inputs["t5xxl"]
+            else:
+                return ""
+            if isinstance(text, list): # type:CLIPTextEncode (NSP) mode:Wildcards
+                text = data[text[0]]["inputs"]["text"]
+            return text.strip()
+        except Exception as e:
+            print(e)
+            return ""   
 
     in_node = data[str(KSampler_entry["positive"][0])]
     if in_node["class_type"] != "FluxGuidance":