ComfyUI-Lora-Manager/py/recipes/parsers/sui_image_params.py

"""Parser for SuiImage (Stable Diffusion WebUI) metadata format."""

import json
import logging
from typing import Dict, Any, Optional, List
from ..base import RecipeMetadataParser
from ...services.metadata_service import get_default_metadata_provider

logger = logging.getLogger(__name__)


class SuiImageParamsParser(RecipeMetadataParser):
    """Parser for SuiImage metadata JSON format.

    This format is used by some Stable Diffusion WebUI variants.
    Structure:
    {
        "sui_image_params": {
            "prompt": "...",
            "negativeprompt": "...",
            "model": "...",
            "seed": ...,
            "steps": ...,
            ...
        },
        "sui_models": [
            {"name": "...", "param": "model", "hash": "..."},
            ...
        ],
        "sui_extra_data": {...}
    }
    """

    def is_metadata_matching(self, user_comment: str) -> bool:
        """Check if the user comment matches the SuiImage metadata format"""
        try:
            data = json.loads(user_comment)
            return isinstance(data, dict) and 'sui_image_params' in data
        except (json.JSONDecodeError, TypeError):
            return False

    async def parse_metadata(self, user_comment: str, recipe_scanner=None, civitai_client=None) -> Dict[str, Any]:
        """Parse metadata from SuiImage metadata format"""
        try:
            metadata_provider = await get_default_metadata_provider()

            data = json.loads(user_comment)
            params = data.get('sui_image_params', {})
            models = data.get('sui_models', [])

            # Extract prompt and negative prompt
            prompt = params.get('prompt', '')
            negative_prompt = params.get('negativeprompt', '') or params.get('negative_prompt', '')

            # Extract generation parameters
            gen_params = {}
            if prompt:
                gen_params['prompt'] = prompt
            if negative_prompt:
                gen_params['negative_prompt'] = negative_prompt

            # Map standard parameters
            param_mapping = {
                'steps': 'steps',
                'seed': 'seed',
                'cfgscale': 'cfg_scale',
                'cfg_scale': 'cfg_scale',
                'width': 'width',
                'height': 'height',
                'sampler': 'sampler',
                'scheduler': 'scheduler',
                'model': 'model',
                'vae': 'vae',
            }

            for src_key, dest_key in param_mapping.items():
                if src_key in params and params[src_key] is not None:
                    gen_params[dest_key] = params[src_key]

            # Add size info if available
            if 'width' in gen_params and 'height' in gen_params:
                gen_params['size'] = f"{gen_params['width']}x{gen_params['height']}"

            # Process models - extract checkpoint and loras
            loras: List[Dict[str, Any]] = []
            checkpoint: Optional[Dict[str, Any]] = None

            for model in models:
                model_name = model.get('name', '')
                param_type = model.get('param', '')
                model_hash = model.get('hash', '')

                # Remove .safetensors extension for cleaner name
                clean_name = model_name.replace('.safetensors', '') if model_name else ''

                # Check if this is a LoRA by looking at the name or param type
                is_lora = 'lora' in model_name.lower() or param_type.lower().startswith('lora')

                if is_lora:
                    lora_entry = {
                        'id': 0,
                        'modelId': 0,
                        'name': clean_name,
                        'version': '',
                        'type': 'lora',
                        'weight': 1.0,
                        'existsLocally': False,
                        'localPath': None,
                        'file_name': model_name,
                        'hash': model_hash.replace('0x', '') if model_hash.startswith('0x') else model_hash,
                        'thumbnailUrl': '/loras_static/images/no-preview.png',
                        'baseModel': '',
                        'size': 0,
                        'downloadUrl': '',
                        'isDeleted': False
                    }

                    # Try to get additional info from metadata provider
                    if metadata_provider and model_hash:
                        try:
                            civitai_info = await metadata_provider.get_model_by_hash(
                                model_hash.replace('0x', '') if model_hash.startswith('0x') else model_hash
                            )
                            if civitai_info:
                                lora_entry = await self.populate_lora_from_civitai(
                                    lora_entry, civitai_info, recipe_scanner
                                )
                        except Exception as e:
                            logger.debug(f"Error fetching info for LoRA {clean_name}: {e}")

                    if lora_entry:
                        loras.append(lora_entry)
                elif param_type == 'model' or 'lora' not in model_name.lower():
                    # This is likely a checkpoint
                    checkpoint_entry = {
                        'id': 0,
                        'modelId': 0,
                        'name': clean_name,
                        'version': '',
                        'type': 'checkpoint',
                        'hash': model_hash.replace('0x', '') if model_hash.startswith('0x') else model_hash,
                        'existsLocally': False,
                        'localPath': None,
                        'file_name': model_name,
                        'thumbnailUrl': '/loras_static/images/no-preview.png',
                        'baseModel': '',
                        'size': 0,
                        'downloadUrl': '',
                        'isDeleted': False
                    }

                    # Try to get additional info from metadata provider
                    if metadata_provider and model_hash:
                        try:
                            civitai_info = await metadata_provider.get_model_by_hash(
                                model_hash.replace('0x', '') if model_hash.startswith('0x') else model_hash
                            )
                            if civitai_info:
                                checkpoint_entry = await self.populate_checkpoint_from_civitai(
                                    checkpoint_entry, civitai_info
                                )
                        except Exception as e:
                            logger.debug(f"Error fetching info for checkpoint {clean_name}: {e}")

                    checkpoint = checkpoint_entry

            # Determine base model from loras or checkpoint
            base_model = None
            if loras:
                base_models = [lora.get('baseModel') for lora in loras if lora.get('baseModel')]
                if base_models:
                    from collections import Counter
                    base_model_counts = Counter(base_models)
                    base_model = base_model_counts.most_common(1)[0][0]
            elif checkpoint and checkpoint.get('baseModel'):
                base_model = checkpoint['baseModel']

            return {
                'base_model': base_model,
                'loras': loras,
                'checkpoint': checkpoint,
                'gen_params': gen_params,
                'from_sui_image_params': True
            }

        except Exception as e:
            logger.error(f"Error parsing SuiImage metadata: {e}", exc_info=True)
            return {"error": str(e), "loras": []}