Extensibility and Plugins

The sd-webui-infinite-image-browsing project is designed with extensibility in mind, allowing users and developers to integrate support for new AI image generation software and extend functionality through custom parsers and plugins. This document outlines the mechanisms for extending the application.

Overview of Extensibility

The core of the application's extensibility revolves around two main points:

1. Custom Image Metadata Parsers: These are Python modules designed to understand and extract generation parameters from images created by different AI software (e.g., ComfyUI, Fooocus, NovelAI).
2. General Plugins: These offer a broader way to extend the application, allowing for custom logic, additional metadata processing, or integration with new data sources beyond just parsing embedded image information.

1. Custom Image Metadata Parsers

Image metadata parsers are crucial for Infinite Image Browsing to understand the various formats in which different AI generation tools embed information (like prompts, models, seeds, etc.) within image files.

How it Works

The application's parser system, managed by scripts/iib/parsers/index.py, iterates through a list of registered parser classes. Each parser has a test method to determine if it can handle a given image. If test returns True, its parse method is called to extract the detailed generation information.

Creating a Custom Parser

To create a new parser for unsupported AI software, follow these steps:

1. Create a New Python File:

   • Navigate to the scripts/iib/parsers/ directory.
   • Create a new Python file (e.g., my_software.py).

2. Define a Parser Class:

   • Inside my_software.py, define a class (e.g., MySoftwareParser). This class must implement two static methods: test and parse.
   • You'll need to import PIL.Image for image handling and ImageGenerationInfo, ImageGenerationParams from scripts/iib/parsers/model.py for structured output.

3. Implement the test Method:

   • This method takes an Image object and file_path as input.
   • Its purpose is to quickly determine if the given image was generated by your target AI software. It should return True if it can parse the image, False otherwise.
   • This method should be lightweight as it might be called for many images.

4. Implement the parse Method:

   • This method is called if test returns True. It takes the same inputs (Image object and file_path).
   • Its role is to extract all relevant generation information and return it as an ImageGenerationInfo object.
   • The ImageGenerationInfo object contains:
     • raw_info (str): The raw string of the generation information (e.g., the full prompt block).
     • params (ImageGenerationParams): A structured object containing:
       • meta (dict): Key-value pairs of various parameters (e.g., Model, Sampler, Steps, CFG scale, Size, Source Identifier). The final_width and final_height should always be included here.
       • pos_prompt (list): A list of positive prompt tags.
       • extra (dict): Any additional parsed data, like Lora/LyCORIS information.

Integration

New parser files placed in scripts/iib/parsers/ are automatically imported and added to the list of parsers in scripts/iib/parsers/index.py. The order of parsers in index.py matters; for example, the SdWebUIParser is usually the last one, acting as a general fallback if no other specialized parser matches.

Example Parsers: Refer to scripts/iib/parsers/comfyui.py, scripts/iib/parsers/fooocus.py, scripts/iib/parsers/novelai.py, and scripts/iib/parsers/stable_swarm_ui.py for concrete examples. The scripts/iib/parsers/sd_webui_stealth.py also demonstrates handling specific image formats (stealth PNG info) and can be enabled via the IIB_ENABLE_SD_WEBUI_STEALTH_PARSER environment variable (see .env.example).

2. General Plugins

General plugins provide a more flexible way to extend Infinite Image Browsing's functionality. They can be used for custom data processing, adding new "Source Identifier" specific logic, or integrating external services.

How it Works

The scripts/iib/plugin.py script scans a designated plugins directory (relative to the project's root) for subdirectories containing a main.py file. It then imports and instantiates a Main class from each main.py file. These instances are stored and can be accessed by other parts of the application.

Creating a General Plugin

To create a general plugin:

1. Create a Plugin Directory:

   • Create a new directory named plugins in the project's root.
   • Inside plugins/, create a subdirectory for your plugin (e.g., plugins/my_custom_plugin/).

2. Create main.py:

   • Inside your plugin's subdirectory, create a main.py file.
   • Define a Main class within this main.py file:

Example Use Case: Custom Tagging from Plugin Metadata

A key use case for general plugins is to specify additional metadata fields from your parser's output that should be automatically converted into searchable tags.

• In scripts/iib/db/update_image_data.py, the get_extra_meta_keys_from_plugins function retrieves the extra_convert_to_tag_meta_keys from loaded plugins.
• During image indexing (build_single_img_idx function), if an image's Source Identifier matches a plugin's source_identifier, the plugin's specified extra_convert_to_tag_meta_keys will be processed. The values corresponding to these keys in the image's metadata will be created as new tags (with the key as the tag type) and linked to the image.

3. Contributing New Software Support or Functionality

Developers are encouraged to contribute to sd-webui-infinite-image-browsing by adding support for new AI generation software or extending its functionality.

For New AI Software

1. Develop a Custom Parser: Create a new Python file in scripts/iib/parsers/ with a test and parse method as described in Section 1.
2. Ensure Source Identifier: Make sure your parse method populates the meta["Source Identifier"] field with a unique string representing your software. This is crucial for distinguishing images and can be used by general plugins.
3. Test Thoroughly: Test your parser with various images generated by the new software, including different models, parameters, and output formats (PNG, WebP, etc.).
4. Submit a Pull Request: Provide a clear description of the new software supported and include example images if possible.

For Broader Functionality / Custom Tagging

1. Develop a General Plugin: Create a new plugin in plugins/your_plugin_name/main.py as described in Section 2.
2. Define source_identifier: If your plugin is meant to enhance the processing of a specific AI software's images, ensure your source_identifier matches the Source Identifier reported by the corresponding parser.
3. Specify extra_convert_to_tag_meta_keys: List any additional metadata keys from your ImageGenerationParams.meta that you want to be automatically searchable as tags.
4. Implement Custom Logic: Add any other custom methods or properties to your Main class as needed.
5. Test: Verify that your plugin's logic executes correctly and integrates seamlessly with existing features.
6. Submit a Pull Request: Clearly explain the new functionality your plugin provides.

Relevant Files for Developers

• scripts/iib/parsers/index.py: The central hub for loading and orchestrating image metadata parsers.
• scripts/iib/parsers/model.py: Defines the ImageGenerationInfo and ImageGenerationParams data structures that all parsers must output.
• scripts/iib/parsers/*.py: Contains existing parser implementations (e.g., sd_webui.py, comfyui.py, fooocus.py, novelai.py, stable_swarm_ui.py, invoke_ai.py, sd_webui_stealth.py) as examples.
• scripts/iib/plugin.py: The script responsible for discovering and loading general plugins.
• scripts/iib/db/update_image_data.py: Contains the logic for processing image metadata and converting specified keys into tags, including those defined by plugins.
• .env.example: Provides examples of environment variables that can configure certain parsing behaviors (e.g., IIB_ENABLE_SD_WEBUI_STEALTH_PARSER).