backend/app/edit_models.py

"""
Edit Models Configuration - Centralized model management for Edit features.

This module provides:
- Detection of installed ComfyUI models for upscale/enhance/face restoration
- Default model configuration with fallbacks
- User-configurable model preferences
- Consistent model access across enhance.py, upscale.py, etc.

Model Categories:
- UPSCALE: 4x-UltraSharp, RealESRGAN, realesr-general, SwinIR
- FACE_RESTORE: GFPGANv1.4, CodeFormer
- BACKGROUND: u2net (rembg), SAM
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Dict, List, Optional, Any

from .providers import get_comfy_models_path


class ModelCategory(str, Enum):
    """Categories of edit models."""
    UPSCALE = "upscale"
    FACE_RESTORE = "face_restore"
    BACKGROUND = "background"
    # Additive: Avatar/Persona wizard specialized models
    AVATAR_GENERATION = "avatar_generation"


@dataclass
class ModelInfo:
    """Information about a single model."""
    id: str
    name: str
    category: ModelCategory
    filename: str
    subdir: str  # Directory under models/ (e.g., "upscale_models", "gfpgan")
    description: str = ""
    is_default: bool = False
    # Additive metadata (backwards-safe — all have defaults)
    license: str = ""
    commercial_use_ok: Optional[bool] = None
    homepage: str = ""
    download_url: str = ""
    sha256: str = ""
    requires: List[str] = field(default_factory=list)  # model IDs this depends on

    @property
    def path(self) -> Path:
        """Get the full path to this model."""
        models_path = get_comfy_models_path()
        return models_path / self.subdir / self.filename

    @property
    def installed(self) -> bool:
        """Check if this model is installed."""
        p = self.path
        return p.exists() and p.stat().st_size > 0


# =============================================================================
# MODEL DEFINITIONS
# =============================================================================

UPSCALE_MODELS: Dict[str, ModelInfo] = {
    "4x-UltraSharp": ModelInfo(
        id="4x-UltraSharp",
        name="4x UltraSharp",
        category=ModelCategory.UPSCALE,
        filename="4x-UltraSharp.pth",
        subdir="upscale_models",
        description="High-quality upscaling with sharp details",
        is_default=True,
    ),
    "RealESRGAN_x4plus": ModelInfo(
        id="RealESRGAN_x4plus",
        name="RealESRGAN x4+",
        category=ModelCategory.UPSCALE,
        filename="RealESRGAN_x4plus.pth",
        subdir="upscale_models",
        description="Photo enhancement with natural texture recovery",
    ),
    "realesr-general-x4v3": ModelInfo(
        id="realesr-general-x4v3",
        name="RealESRGAN General v3",
        category=ModelCategory.UPSCALE,
        filename="realesr-general-x4v3.pth",
        subdir="upscale_models",
        description="General purpose upscaling",
    ),
    "SwinIR_4x": ModelInfo(
        id="SwinIR_4x",
        name="SwinIR 4x",
        category=ModelCategory.UPSCALE,
        filename="SwinIR_4x.pth",
        subdir="upscale_models",
        description="Artifact removal and restoration",
    ),
}

FACE_RESTORE_MODELS: Dict[str, ModelInfo] = {
    "GFPGANv1.4": ModelInfo(
        id="GFPGANv1.4",
        name="GFPGAN v1.4",
        category=ModelCategory.FACE_RESTORE,
        filename="GFPGANv1.4.pth",
        subdir="gfpgan",  # Matches Makefile download path
        description="Face restoration and enhancement",
        is_default=True,
    ),
    "codeformer": ModelInfo(
        id="codeformer",
        name="CodeFormer",
        category=ModelCategory.FACE_RESTORE,
        filename="codeformer.pth",
        subdir="codeformer",  # Matches providers.py scan path
        description="AI face restoration with fidelity control",
    ),
}

BACKGROUND_MODELS: Dict[str, ModelInfo] = {
    "u2net": ModelInfo(
        id="u2net",
        name="U2-Net",
        category=ModelCategory.BACKGROUND,
        filename="u2net.onnx",
        subdir="rembg",
        description="Background removal segmentation",
        is_default=True,
    ),
    "sam_vit_h": ModelInfo(
        id="sam_vit_h",
        name="SAM ViT-H",
        category=ModelCategory.BACKGROUND,
        filename="sam_vit_h_4b8939.pth",
        subdir="sams",
        description="Segment Anything Model (high quality)",
    ),
}

# =============================================================================
# AVATAR GENERATION MODELS (Additive — Golden Rule 1.0)
# These models extend the registry for the PersonaWizard "Portrait Studio".
# They do NOT replace or affect existing text-to-image, edit, or enhance models.
# =============================================================================

AVATAR_GENERATION_MODELS: Dict[str, ModelInfo] = {
    # ── Face analysis / embeddings (commonly used by InstantID & face-swap) ──
    "insightface-antelopev2": ModelInfo(
        id="insightface-antelopev2",
        name="InsightFace AntelopeV2",
        category=ModelCategory.AVATAR_GENERATION,
        filename="antelopev2.zip",
        subdir="insightface/models",
        description="Face detection & embedding pack used by InstantID and face-swap workflows.",
        license="InsightFace model zoo (code MIT, models vary)",
        commercial_use_ok=True,
        homepage="https://github.com/deepinsight/insightface",
        download_url="https://huggingface.co/MonsterMMORPG/tools/resolve/main/antelopev2.zip",
        sha256="",
        requires=[],
        is_default=True,
    ),

    # ── Face swap model (optional) ──
    "insightface-inswapper-128": ModelInfo(
        id="insightface-inswapper-128",
        name="InsightFace InSwapper 128 (ONNX)",
        category=ModelCategory.AVATAR_GENERATION,
        filename="inswapper_128.onnx",
        subdir="insightface",
        description="Face swap ONNX model for consistent identity transfer onto generated images.",
        license="Model distribution varies (see source)",
        commercial_use_ok=True,
        homepage="https://github.com/deepinsight/insightface",
        download_url="https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx",
        sha256="",
        requires=["insightface-antelopev2"],
    ),

    # ── InstantID (identity-preserving adapter, Apache 2.0) ──
    "instantid-ip-adapter": ModelInfo(
        id="instantid-ip-adapter",
        name="InstantID IP-Adapter",
        category=ModelCategory.AVATAR_GENERATION,
        filename="ip-adapter.bin",
        subdir="instantid",
        description="InstantID adapter checkpoint for identity-preserving diffusion generation.",
        license="Apache 2.0",
        commercial_use_ok=True,
        homepage="https://huggingface.co/InstantX/InstantID",
        download_url="https://huggingface.co/InstantX/InstantID/resolve/main/ip-adapter.bin",
        sha256="",
        requires=["insightface-antelopev2"],
    ),

    # ── InstantID ControlNet ──
    "instantid-controlnet": ModelInfo(
        id="instantid-controlnet",
        name="InstantID ControlNet",
        category=ModelCategory.AVATAR_GENERATION,
        filename="diffusion_pytorch_model.safetensors",
        subdir="controlnet/InstantID",
        description="InstantID ControlNet for facial keypoint guidance during generation.",
        license="Apache 2.0",
        commercial_use_ok=True,
        homepage="https://huggingface.co/InstantX/InstantID",
        download_url="https://huggingface.co/InstantX/InstantID/resolve/main/ControlNetModel/diffusion_pytorch_model.safetensors",
        sha256="",
        requires=["instantid-ip-adapter"],
    ),

    # ── OpenPose ControlNet for SDXL (pose-guided body generation) ──
    "openpose-controlnet-sdxl": ModelInfo(
        id="openpose-controlnet-sdxl",
        name="OpenPose ControlNet (SDXL)",
        category=ModelCategory.AVATAR_GENERATION,
        filename="diffusion_pytorch_model.safetensors",
        subdir="controlnet/thibaud-openpose-sdxl-1.0",
        description="OpenPose ControlNet for SDXL enabling pose-guided body generation. Safetensors conversion by dimitribarbot.",
        license="CreativeML Open RAIL-M",
        commercial_use_ok=True,
        homepage="https://huggingface.co/thibaud/controlnet-openpose-sdxl-1.0",
        download_url="https://huggingface.co/dimitribarbot/controlnet-openpose-sdxl-1.0-safetensors/resolve/main/diffusion_pytorch_model.safetensors",
        sha256="",
        requires=["instantid-controlnet"],
    ),

    # ── PhotoMaker V2 (identity-preserving, Apache 2.0) ──
    "photomaker-v2": ModelInfo(
        id="photomaker-v2",
        name="PhotoMaker V2",
        category=ModelCategory.AVATAR_GENERATION,
        filename="photomaker-v2.bin",
        subdir="photomaker",
        description="PhotoMaker V2 identity-preserving encoder for SDXL-compatible workflows.",
        license="Apache 2.0",
        commercial_use_ok=True,
        homepage="https://huggingface.co/TencentARC/PhotoMaker-V2",
        download_url="https://huggingface.co/TencentARC/PhotoMaker-V2/resolve/main/photomaker-v2.bin",
        sha256="",
        requires=[],
    ),

    # ── PuLID (Flux adapter, advanced) ──
    "pulid-flux": ModelInfo(
        id="pulid-flux",
        name="PuLID for FLUX",
        category=ModelCategory.AVATAR_GENERATION,
        filename="pulid_flux_v0.9.0.safetensors",
        subdir="pulid",
        description="PuLID adapter for FLUX identity customization with minimal disruption.",
        license="Apache 2.0",
        commercial_use_ok=True,
        homepage="https://huggingface.co/guozinan/PuLID",
        download_url="https://huggingface.co/guozinan/PuLID/resolve/main/pulid_flux_v0.9.0.safetensors",
        sha256="",
        requires=["insightface-antelopev2"],
    ),

    # ── IP-Adapter FaceID Plus V2 (face-conditioned generation) ──
    "ip-adapter-faceid-plusv2": ModelInfo(
        id="ip-adapter-faceid-plusv2",
        name="IP-Adapter FaceID Plus V2 (SDXL)",
        category=ModelCategory.AVATAR_GENERATION,
        filename="ip-adapter-faceid-plusv2_sdxl.bin",
        subdir="ipadapter",
        description="Face-conditioned generation adapter for SDXL. Non-commercial license.",
        license="Non-Commercial (h94)",
        commercial_use_ok=False,
        homepage="https://huggingface.co/h94/IP-Adapter-FaceID",
        download_url="https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main/ip-adapter-faceid-plusv2_sdxl.bin",
        sha256="",
        requires=["insightface-antelopev2"],
    ),

    # ── StyleGAN2 FFHQ 256 (fast random faces, non-commercial) ──
    "stylegan2-ffhq-256": ModelInfo(
        id="stylegan2-ffhq-256",
        name="StyleGAN2 FFHQ 256",
        category=ModelCategory.AVATAR_GENERATION,
        filename="stylegan2-ffhq-256x256.pkl",
        subdir="avatar",
        description="Fast random-face generator (FFHQ 256). Non-commercial (NVIDIA).",
        license="NVIDIA Source Code License (Non-commercial)",
        commercial_use_ok=False,
        homepage="https://github.com/NVlabs/stylegan2",
        download_url="https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-256x256.pkl",
        sha256="",
        requires=[],
    ),

    # ── StyleGAN2 FFHQ 1024 (high-quality random faces, non-commercial) ──
    "stylegan2-ffhq-1024": ModelInfo(
        id="stylegan2-ffhq-1024",
        name="StyleGAN2 FFHQ 1024",
        category=ModelCategory.AVATAR_GENERATION,
        filename="stylegan2-ffhq-1024x1024.pkl",
        subdir="avatar",
        description="High-quality random-face generator (FFHQ 1024). Non-commercial (NVIDIA).",
        license="NVIDIA Source Code License (Non-commercial)",
        commercial_use_ok=False,
        homepage="https://github.com/NVlabs/stylegan2",
        download_url="https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-1024x1024.pkl",
        sha256="",
        requires=[],
    ),
}


# All models by category
ALL_MODELS: Dict[ModelCategory, Dict[str, ModelInfo]] = {
    ModelCategory.UPSCALE: UPSCALE_MODELS,
    ModelCategory.FACE_RESTORE: FACE_RESTORE_MODELS,
    ModelCategory.BACKGROUND: BACKGROUND_MODELS,
    # Additive: used by PersonaWizard "Portrait Studio" (future)
    ModelCategory.AVATAR_GENERATION: AVATAR_GENERATION_MODELS,
}


# =============================================================================
# ENHANCE MODE CONFIGURATION
# =============================================================================

@dataclass
class EnhanceModeConfig:
    """Configuration for an enhance mode."""
    mode: str
    name: str
    description: str
    workflow: str
    model_category: ModelCategory
    default_model_id: str
    param_name: str = "upscale_model"  # Workflow parameter name


# Default enhance mode configurations
ENHANCE_MODES: Dict[str, EnhanceModeConfig] = {
    "photo": EnhanceModeConfig(
        mode="photo",
        name="Photo Enhancement",
        description="Natural photo enhancement with texture recovery",
        workflow="upscale",
        model_category=ModelCategory.UPSCALE,
        default_model_id="4x-UltraSharp",
        param_name="upscale_model",
    ),
    "restore": EnhanceModeConfig(
        mode="restore",
        name="Restoration",
        description="Remove artifacts and mild blur",
        workflow="upscale",
        model_category=ModelCategory.UPSCALE,
        default_model_id="4x-UltraSharp",
        param_name="upscale_model",
    ),
    "faces": EnhanceModeConfig(
        mode="faces",
        name="Face Restoration",
        description="Restore and enhance faces using FaceDetailer",
        workflow="fix_faces_facedetailer",
        model_category=ModelCategory.FACE_RESTORE,
        default_model_id="GFPGANv1.4",
        param_name="model_name",
    ),
}


# =============================================================================
# USER PREFERENCES (Runtime State)
# =============================================================================

@dataclass
class EditModelPreferences:
    """User preferences for edit models."""
    # Selected model IDs by mode
    upscale_model: str = "4x-UltraSharp"
    photo_model: str = "4x-UltraSharp"
    restore_model: str = "4x-UltraSharp"
    faces_model: str = "GFPGANv1.4"

    def get_model_for_mode(self, mode: str) -> str:
        """Get the selected model ID for a given mode."""
        mapping = {
            "upscale": self.upscale_model,
            "photo": self.photo_model,
            "restore": self.restore_model,
            "faces": self.faces_model,
        }
        return mapping.get(mode, self.upscale_model)

    def set_model_for_mode(self, mode: str, model_id: str) -> None:
        """Set the selected model ID for a given mode."""
        if mode == "upscale":
            self.upscale_model = model_id
        elif mode == "photo":
            self.photo_model = model_id
        elif mode == "restore":
            self.restore_model = model_id
        elif mode == "faces":
            self.faces_model = model_id


# Global preferences instance (can be modified at runtime)
_preferences = EditModelPreferences()


def get_preferences() -> EditModelPreferences:
    """Get current model preferences."""
    return _preferences


def set_preferences(prefs: EditModelPreferences) -> None:
    """Set model preferences."""
    global _preferences
    _preferences = prefs


# =============================================================================
# MODEL DETECTION & SELECTION
# =============================================================================

def get_installed_models(category: ModelCategory) -> List[ModelInfo]:
    """Get list of installed models in a category."""
    models = ALL_MODELS.get(category, {})
    return [m for m in models.values() if m.installed]


def get_all_models(category: ModelCategory) -> List[ModelInfo]:
    """Get all models in a category (installed or not)."""
    return list(ALL_MODELS.get(category, {}).values())


def get_model_info(model_id: str) -> Optional[ModelInfo]:
    """Get model info by ID, searching all categories."""
    for category_models in ALL_MODELS.values():
        if model_id in category_models:
            return category_models[model_id]
    return None


def get_default_model(category: ModelCategory) -> Optional[ModelInfo]:
    """Get the default model for a category."""
    models = ALL_MODELS.get(category, {})
    for model in models.values():
        if model.is_default:
            return model
    # Return first if no default specified
    return next(iter(models.values()), None) if models else None


def get_available_model(category: ModelCategory, preferred_id: Optional[str] = None) -> Optional[ModelInfo]:
    """
    Get an available model, preferring the specified one.

    Priority:
    1. Preferred model if installed
    2. Default model if installed
    3. Any installed model
    4. None if nothing installed
    """
    # Try preferred model first
    if preferred_id:
        model = get_model_info(preferred_id)
        if model and model.installed:
            return model

    # Try default model
    default = get_default_model(category)
    if default and default.installed:
        return default

    # Try any installed model
    installed = get_installed_models(category)
    if installed:
        return installed[0]

    return None


# =============================================================================
# HIGH-LEVEL API FOR ENHANCE/UPSCALE
# =============================================================================

def get_upscale_model() -> tuple[Optional[str], Optional[str]]:
    """
    Get the upscale model to use.

    Returns:
        (model_filename, error_message) - filename if available, error if not
    """
    prefs = get_preferences()
    model = get_available_model(ModelCategory.UPSCALE, prefs.upscale_model)

    if model:
        return model.filename, None

    # No model available
    all_models = get_all_models(ModelCategory.UPSCALE)
    model_names = [m.name for m in all_models]
    return None, f"No upscale model installed. Please install one of: {', '.join(model_names)}"


def get_enhance_model(mode: str) -> tuple[Optional[str], Optional[str], Optional[EnhanceModeConfig]]:
    """
    Get the model to use for an enhance mode.

    Args:
        mode: "photo", "restore", or "faces"

    Returns:
        (model_filename, error_message, mode_config)
    """
    mode_config = ENHANCE_MODES.get(mode)
    if not mode_config:
        return None, f"Unknown enhance mode: {mode}", None

    prefs = get_preferences()
    preferred = prefs.get_model_for_mode(mode)
    model = get_available_model(mode_config.model_category, preferred)

    if model:
        return model.filename, None, mode_config

    # No model available
    all_models = get_all_models(mode_config.model_category)
    model_names = [m.name for m in all_models]
    return None, f"No {mode_config.name} model installed. Please install one of: {', '.join(model_names)}", mode_config


def get_face_restore_model() -> tuple[Optional[str], Optional[str]]:
    """
    Get the face restoration model to use.

    Returns:
        (model_filename, error_message)
    """
    prefs = get_preferences()
    model = get_available_model(ModelCategory.FACE_RESTORE, prefs.faces_model)

    if model:
        return model.filename, None

    all_models = get_all_models(ModelCategory.FACE_RESTORE)
    model_names = [m.name for m in all_models]
    return None, f"No face restoration model installed. Please install one of: {', '.join(model_names)}"


# =============================================================================
# STATUS & REPORTING
# =============================================================================

def get_edit_models_status() -> Dict[str, Any]:
    """
    Get comprehensive status of all edit models.

    Returns a dict suitable for API response with installed/available models.
    """
    # Check ComfyUI face restoration readiness
    try:
        from .face_restore import face_restore_ready
        comfyui_ok, comfyui_reason = face_restore_ready()
    except Exception:
        comfyui_ok, comfyui_reason = False, "Import error"

    status: Dict[str, Any] = {
        "upscale": {
            "installed": [],
            "available": [],
            "selected": get_preferences().upscale_model,
            "default": None,
        },
        "enhance": {
            "photo": {
                "installed": [],
                "available": [],
                "selected": get_preferences().photo_model,
            },
            "restore": {
                "installed": [],
                "available": [],
                "selected": get_preferences().restore_model,
            },
            "faces": {
                "installed": [],
                "available": [],
                "selected": get_preferences().faces_model,
                "comfyui_ready": comfyui_ok,
                "comfyui_status": comfyui_reason,
            },
        },
    }

    # Upscale models
    for model in get_all_models(ModelCategory.UPSCALE):
        model_info = {
            "id": model.id,
            "name": model.name,
            "description": model.description,
            "filename": model.filename,
            "installed": model.installed,
            "is_default": model.is_default,
        }
        status["upscale"]["available"].append(model_info)
        if model.installed:
            status["upscale"]["installed"].append(model.id)
        if model.is_default:
            status["upscale"]["default"] = model.id

    # Face restore models
    for model in get_all_models(ModelCategory.FACE_RESTORE):
        model_info = {
            "id": model.id,
            "name": model.name,
            "description": model.description,
            "filename": model.filename,
            "installed": model.installed,
            "is_default": model.is_default,
        }
        status["enhance"]["faces"]["available"].append(model_info)
        if model.installed:
            status["enhance"]["faces"]["installed"].append(model.id)

    # Copy upscale info to photo/restore (they use same models)
    status["enhance"]["photo"]["available"] = status["upscale"]["available"]
    status["enhance"]["photo"]["installed"] = status["upscale"]["installed"]
    status["enhance"]["restore"]["available"] = status["upscale"]["available"]
    status["enhance"]["restore"]["installed"] = status["upscale"]["installed"]

    return status


def set_model_preference(mode: str, model_id: str) -> tuple[bool, Optional[str]]:
    """
    Set the preferred model for a mode.

    Args:
        mode: "upscale", "photo", "restore", or "faces"
        model_id: The model ID to use

    Returns:
        (success, error_message)
    """
    # Validate model exists
    model = get_model_info(model_id)
    if not model:
        return False, f"Unknown model: {model_id}"

    # Validate model category matches mode
    mode_to_category = {
        "upscale": ModelCategory.UPSCALE,
        "photo": ModelCategory.UPSCALE,
        "restore": ModelCategory.UPSCALE,
        "faces": ModelCategory.FACE_RESTORE,
    }

    expected_category = mode_to_category.get(mode)
    if not expected_category:
        return False, f"Unknown mode: {mode}"

    if model.category != expected_category:
        return False, f"Model {model_id} is not compatible with mode {mode}"

    # Check if installed
    if not model.installed:
        return False, f"Model {model.name} is not installed. Please install it first."

    # Set preference
    prefs = get_preferences()
    prefs.set_model_for_mode(mode, model_id)

    return True, None


# =============================================================================
# AVATAR MODEL STATUS (Additive — does NOT affect edit/enhance behavior)
# =============================================================================

def get_avatar_models_status() -> Dict[str, Any]:
    """
    Return installed/available status for Avatar Generator models.

    Intended for the PersonaWizard UI to gate avatar generation modes
    (enable/disable based on which models are downloaded).
    This does NOT alter existing edit/enhance behavior.
    """
    models = get_all_models(ModelCategory.AVATAR_GENERATION)

    available = []
    installed_ids = []

    # Feature-to-model mapping: which models each feature needs
    FEATURE_MODELS: Dict[str, Dict[str, Any]] = {
        "photo_variations": {
            "label": "Same-Person Photo Variations",
            "description": "Generate new photos preserving the persona's identity",
            "required": ["insightface-antelopev2", "instantid-ip-adapter"],
            "recommended": ["instantid-controlnet"],
        },
        "outfit_generation": {
            "label": "Outfit / Wardrobe Generation",
            "description": "Generate outfit variations with identity preservation",
            "required": ["insightface-antelopev2", "instantid-ip-adapter", "instantid-controlnet"],
            "recommended": ["photomaker-v2", "pulid-flux"],
            "recommended_note": "PhotoMaker V2 for SDXL models, PuLID for FLUX models",
        },
        "face_swap": {
            "label": "Face Swap",
            "description": "Transfer identity onto generated or existing images",
            "required": ["insightface-antelopev2", "insightface-inswapper-128"],
            "recommended": [],
        },
        "pose_guided_body": {
            "label": "Pose-Guided Body Generation",
            "description": "Generate full-body images with OpenPose control for precise posing",
            "required": ["insightface-antelopev2", "instantid-ip-adapter", "instantid-controlnet", "openpose-controlnet-sdxl"],
            "recommended": [],
        },
        "random_faces": {
            "label": "Random Face Generator",
            "description": "Generate random realistic faces (non-commercial)",
            "required": ["stylegan2-ffhq-256"],
            "recommended": ["stylegan2-ffhq-1024"],
            "recommended_note": "1024px version for higher quality output",
        },
    }

    # Build model→features reverse index
    model_features: Dict[str, list] = {}
    for feat_id, feat in FEATURE_MODELS.items():
        for mid in feat["required"]:
            model_features.setdefault(mid, []).append({"feature": feat_id, "role": "required"})
        for mid in feat.get("recommended", []):
            model_features.setdefault(mid, []).append({"feature": feat_id, "role": "recommended"})

    for m in models:
        info = {
            "id": m.id,
            "name": m.name,
            "description": m.description,
            "filename": m.filename,
            "subdir": m.subdir,
            "installed": m.installed,
            "license": m.license,
            "commercial_use_ok": m.commercial_use_ok,
            "homepage": m.homepage,
            "download_url": m.download_url,
            "sha256": m.sha256,
            "requires": m.requires,
            "is_default": m.is_default,
            "used_by": model_features.get(m.id, []),
        }
        available.append(info)
        if m.installed:
            installed_ids.append(m.id)

    # Compute feature readiness
    installed_set = set(installed_ids)
    features = {}
    for feat_id, feat in FEATURE_MODELS.items():
        required_ok = all(mid in installed_set for mid in feat["required"])
        required_missing = [mid for mid in feat["required"] if mid not in installed_set]
        recommended_ok = all(mid in installed_set for mid in feat.get("recommended", []))
        features[feat_id] = {
            "label": feat["label"],
            "description": feat["description"],
            "ready": required_ok,
            "required_missing": required_missing,
            "recommended_installed": recommended_ok,
            "recommended_note": feat.get("recommended_note"),
        }

    return {
        "category": ModelCategory.AVATAR_GENERATION.value,
        "installed": installed_ids,
        "available": available,
        "defaults": [m.id for m in models if m.is_default],
        "features": features,
    }