invoke-ai--invokeai/invokeai/app/api/routers/recall_parameters.py

"""Router for updating recallable parameters on the frontend."""

import json
from typing import Any, Literal, Optional

from fastapi import Body, HTTPException, Path, Query
from fastapi.routing import APIRouter
from pydantic import BaseModel, ConfigDict, Field

from invokeai.app.api.auth_dependencies import CurrentUserOrDefault
from invokeai.app.api.dependencies import ApiDependencies
from invokeai.app.api.routers.image_move_maintenance import assert_image_move_maintenance_inactive
from invokeai.backend.image_util.controlnet_processor import process_controlnet_image
from invokeai.backend.model_manager.taxonomy import ModelType

recall_parameters_router = APIRouter(prefix="/v1/recall", tags=["recall"])


class LoRARecallParameter(BaseModel):
    """LoRA configuration for recall"""

    model_name: str = Field(description="The name of the LoRA model")
    weight: float = Field(default=0.75, ge=-10, le=10, description="The weight for the LoRA")
    is_enabled: bool = Field(default=True, description="Whether the LoRA is enabled")


class ControlNetRecallParameter(BaseModel):
    """ControlNet configuration for recall"""

    model_name: str = Field(description="The name of the ControlNet/T2I Adapter/Control LoRA model")
    image_name: Optional[str] = Field(default=None, description="The filename of the control image in outputs/images")
    weight: float = Field(default=1.0, ge=-1, le=2, description="The weight for the control adapter")
    begin_step_percent: Optional[float] = Field(
        default=None, ge=0, le=1, description="When the control adapter is first applied (% of total steps)"
    )
    end_step_percent: Optional[float] = Field(
        default=None, ge=0, le=1, description="When the control adapter is last applied (% of total steps)"
    )
    control_mode: Optional[Literal["balanced", "more_prompt", "more_control"]] = Field(
        default=None, description="The control mode (ControlNet only)"
    )


class IPAdapterRecallParameter(BaseModel):
    """IP Adapter configuration for recall"""

    model_name: str = Field(description="The name of the IP Adapter model")
    image_name: Optional[str] = Field(default=None, description="The filename of the reference image in outputs/images")
    weight: float = Field(default=1.0, ge=-1, le=2, description="The weight for the IP Adapter")
    begin_step_percent: Optional[float] = Field(
        default=None, ge=0, le=1, description="When the IP Adapter is first applied (% of total steps)"
    )
    end_step_percent: Optional[float] = Field(
        default=None, ge=0, le=1, description="When the IP Adapter is last applied (% of total steps)"
    )
    method: Optional[Literal["full", "style", "composition"]] = Field(default=None, description="The IP Adapter method")
    image_influence: Optional[Literal["lowest", "low", "medium", "high", "highest"]] = Field(
        default=None, description="FLUX Redux image influence (if model is flux_redux)"
    )


class ReferenceImageRecallParameter(BaseModel):
    """Global reference-image configuration for recall.

    Used for reference images that feed directly into the main model rather
    than through a separate IP-Adapter / ControlNet model — for example
    FLUX.2 Klein, FLUX Kontext, and Qwen Image Edit. The receiving frontend
    picks the correct config type (``flux2_reference_image`` /
    ``qwen_image_reference_image`` / ``flux_kontext_reference_image``) based
    on the currently-selected main model.
    """

    image_name: str = Field(description="The filename of the reference image in outputs/images")


class RecallParameter(BaseModel):
    """Request model for updating recallable parameters."""

    model_config = ConfigDict(extra="forbid")

    # Prompts
    positive_prompt: Optional[str] = Field(None, description="Positive prompt text")
    negative_prompt: Optional[str] = Field(None, description="Negative prompt text")

    # Model configuration
    model: Optional[str] = Field(None, description="Main model name/identifier")
    refiner_model: Optional[str] = Field(None, description="Refiner model name/identifier")
    vae_model: Optional[str] = Field(None, description="VAE model name/identifier")
    scheduler: Optional[str] = Field(None, description="Scheduler name")

    # Generation parameters
    steps: Optional[int] = Field(None, ge=1, description="Number of generation steps")
    refiner_steps: Optional[int] = Field(None, ge=0, description="Number of refiner steps")
    cfg_scale: Optional[float] = Field(None, description="CFG scale for guidance")
    cfg_rescale_multiplier: Optional[float] = Field(None, description="CFG rescale multiplier")
    refiner_cfg_scale: Optional[float] = Field(None, description="Refiner CFG scale")
    guidance: Optional[float] = Field(None, description="Guidance scale")

    # Image parameters
    width: Optional[int] = Field(None, ge=64, description="Image width in pixels")
    height: Optional[int] = Field(None, ge=64, description="Image height in pixels")
    seed: Optional[int] = Field(None, ge=0, description="Random seed")

    # Advanced parameters
    denoise_strength: Optional[float] = Field(None, ge=0, le=1, description="Denoising strength")
    refiner_denoise_start: Optional[float] = Field(None, ge=0, le=1, description="Refiner denoising start")
    clip_skip: Optional[int] = Field(None, ge=0, description="CLIP skip layers")
    seamless_x: Optional[bool] = Field(None, description="Enable seamless X tiling")
    seamless_y: Optional[bool] = Field(None, description="Enable seamless Y tiling")

    # Refiner aesthetics
    refiner_positive_aesthetic_score: Optional[float] = Field(None, description="Refiner positive aesthetic score")
    refiner_negative_aesthetic_score: Optional[float] = Field(None, description="Refiner negative aesthetic score")

    # LoRAs, ControlNets, and IP Adapters
    loras: Optional[list[LoRARecallParameter]] = Field(None, description="List of LoRAs with their weights")
    control_layers: Optional[list[ControlNetRecallParameter]] = Field(
        None, description="List of control adapters (ControlNet, T2I Adapter, Control LoRA) with their settings"
    )
    ip_adapters: Optional[list[IPAdapterRecallParameter]] = Field(
        None, description="List of IP Adapters with their settings"
    )
    reference_images: Optional[list[ReferenceImageRecallParameter]] = Field(
        None,
        description=(
            "List of model-free reference images for architectures that consume reference "
            "images directly (FLUX.2 Klein, FLUX Kontext, Qwen Image Edit). The frontend "
            "picks the correct config type based on the currently-selected main model."
        ),
    )


def resolve_model_name_to_key(model_name: str, model_type: ModelType = ModelType.Main) -> Optional[str]:
    """
    Look up a model by name and return its key.

    Args:
        model_name: The name of the model to look up
        model_type: The type of model to search for (default: Main)

    Returns:
        The key of the first matching model, or None if not found.
    """
    logger = ApiDependencies.invoker.services.logger
    try:
        models = ApiDependencies.invoker.services.model_manager.store.search_by_attr(
            model_name=model_name, model_type=model_type
        )

        if models:
            logger.info(f"Resolved {model_type.value} model name '{model_name}' to key '{models[0].key}'")
            return models[0].key

        logger.warning(f"Could not find {model_type.value} model with name '{model_name}'")
        return None
    except Exception as e:
        logger.error(f"Exception during {model_type.value} model lookup: {e}", exc_info=True)
        return None


def load_image_file(image_name: str) -> Optional[dict[str, Any]]:
    """
    Load an image from the outputs/images directory.

    Args:
        image_name: The filename of the image in outputs/images

    Returns:
        A dictionary with image_name, width, and height, or None if the image cannot be found
    """
    logger = ApiDependencies.invoker.services.logger
    try:
        images_service = ApiDependencies.invoker.services.images
        # Use images service which handles subfolder resolution via DB record
        path = images_service.get_path(image_name)

        if not images_service.validate_path(path):
            logger.warning(f"Image file not found: {image_name}")
            return None

        pil_image = images_service.get_pil_image(image_name)
        width, height = pil_image.size
        logger.info(f"Found image file: {image_name} ({width}x{height})")
        return {"image_name": image_name, "width": width, "height": height}
    except Exception as e:
        logger.warning(f"Error loading image file {image_name}: {e}")
        return None


def resolve_lora_models(loras: list[LoRARecallParameter]) -> list[dict[str, Any]]:
    """
    Resolve LoRA model names to keys and build configuration list.

    Args:
        loras: List of LoRA recall parameters

    Returns:
        List of resolved LoRA configurations with model keys
    """
    logger = ApiDependencies.invoker.services.logger
    resolved_loras = []

    for lora in loras:
        model_key = resolve_model_name_to_key(lora.model_name, ModelType.LoRA)
        if model_key:
            resolved_loras.append({"model_key": model_key, "weight": lora.weight, "is_enabled": lora.is_enabled})
        else:
            logger.warning(f"Skipping LoRA '{lora.model_name}' - model not found")

    return resolved_loras


def resolve_control_models(control_layers: list[ControlNetRecallParameter]) -> list[dict[str, Any]]:
    """
    Resolve control adapter model names to keys and build configuration list.

    Tries to resolve as ControlNet, T2I Adapter, or Control LoRA in that order.

    Args:
        control_layers: List of control adapter recall parameters

    Returns:
        List of resolved control adapter configurations with model keys
    """
    logger = ApiDependencies.invoker.services.logger
    services = ApiDependencies.invoker.services
    resolved_controls = []

    for control in control_layers:
        model_key = None

        # Try ControlNet first
        model_key = resolve_model_name_to_key(control.model_name, ModelType.ControlNet)
        if not model_key:
            # Try T2I Adapter
            model_key = resolve_model_name_to_key(control.model_name, ModelType.T2IAdapter)
        if not model_key:
            # Try Control LoRA (also uses LoRA type)
            model_key = resolve_model_name_to_key(control.model_name, ModelType.LoRA)

        if model_key:
            config: dict[str, Any] = {"model_key": model_key, "weight": control.weight}
            if control.image_name is not None:
                image_data = load_image_file(control.image_name)
                if image_data:
                    config["image"] = image_data

                    # Try to process the image using the model's default processor
                    processed_image_data = process_controlnet_image(control.image_name, model_key, services)
                    if processed_image_data:
                        config["processed_image"] = processed_image_data
                        logger.info(f"Added processed image for control adapter {control.model_name}")
                else:
                    logger.warning(f"Could not load image for control adapter: {control.image_name}")
            if control.begin_step_percent is not None:
                config["begin_step_percent"] = control.begin_step_percent
            if control.end_step_percent is not None:
                config["end_step_percent"] = control.end_step_percent
            if control.control_mode is not None:
                config["control_mode"] = control.control_mode

            resolved_controls.append(config)
        else:
            logger.warning(f"Skipping control adapter '{control.model_name}' - model not found")

    return resolved_controls


def resolve_ip_adapter_models(ip_adapters: list[IPAdapterRecallParameter]) -> list[dict[str, Any]]:
    """
    Resolve IP Adapter model names to keys and build configuration list.

    Args:
        ip_adapters: List of IP Adapter recall parameters

    Returns:
        List of resolved IP Adapter configurations with model keys
    """
    logger = ApiDependencies.invoker.services.logger
    resolved_adapters = []

    for adapter in ip_adapters:
        # Try resolving as IP Adapter; if not found, try FLUX Redux
        model_key = resolve_model_name_to_key(adapter.model_name, ModelType.IPAdapter)
        if not model_key:
            model_key = resolve_model_name_to_key(adapter.model_name, ModelType.FluxRedux)
        if model_key:
            config: dict[str, Any] = {
                "model_key": model_key,
                # Always include weight; ignored by FLUX Redux on the frontend
                "weight": adapter.weight,
            }
            if adapter.image_name is not None:
                image_data = load_image_file(adapter.image_name)
                if image_data:
                    config["image"] = image_data
                else:
                    logger.warning(f"Could not load image for IP Adapter: {adapter.image_name}")
            if adapter.begin_step_percent is not None:
                config["begin_step_percent"] = adapter.begin_step_percent
            if adapter.end_step_percent is not None:
                config["end_step_percent"] = adapter.end_step_percent
            if adapter.method is not None:
                config["method"] = adapter.method
            # Include FLUX Redux image influence when provided
            if adapter.image_influence is not None:
                config["image_influence"] = adapter.image_influence

            resolved_adapters.append(config)
        else:
            logger.warning(f"Skipping IP Adapter '{adapter.model_name}' - model not found")

    return resolved_adapters


def resolve_reference_images(
    reference_images: list[ReferenceImageRecallParameter],
) -> list[dict[str, Any]]:
    """
    Validate model-free reference images and build the configuration list.

    Unlike IP Adapters and ControlNets, these reference images are consumed
    directly by the main model (FLUX.2 Klein, FLUX Kontext, Qwen Image Edit),
    so there is no adapter-model name to resolve. We simply verify that each
    referenced file exists in ``outputs/images`` and pass the image metadata
    through to the frontend.

    Args:
        reference_images: List of reference-image recall parameters

    Returns:
        List of reference-image configurations with resolved image metadata.
        Entries whose image file cannot be loaded are dropped with a warning.
    """
    logger = ApiDependencies.invoker.services.logger
    resolved: list[dict[str, Any]] = []

    for ref in reference_images:
        image_data = load_image_file(ref.image_name)
        if image_data is None:
            logger.warning(f"Skipping reference image '{ref.image_name}' - file not found")
            continue
        resolved.append({"image": image_data})

    return resolved


def _assert_recall_image_access(parameters: "RecallParameter", current_user: CurrentUserOrDefault) -> None:
    """Validate that the caller can read every image referenced in the recall parameters.

    Control layers, IP adapters, and reference images may reference image_name fields.
    Without this check an attacker who knows another user's image UUID could use the recall
    endpoint to extract image dimensions and — for ControlNet preprocessors — mint
    a derived processed image they can then fetch.
    """
    from invokeai.app.services.board_records.board_records_common import BoardVisibility

    image_names: list[str] = []
    if parameters.control_layers:
        for layer in parameters.control_layers:
            if layer.image_name is not None:
                image_names.append(layer.image_name)
    if parameters.ip_adapters:
        for adapter in parameters.ip_adapters:
            if adapter.image_name is not None:
                image_names.append(adapter.image_name)
    if parameters.reference_images:
        for ref in parameters.reference_images:
            if ref.image_name is not None:
                image_names.append(ref.image_name)

    if not image_names:
        return

    # Admin can access all images
    if current_user.is_admin:
        return

    for image_name in image_names:
        owner = ApiDependencies.invoker.services.image_records.get_user_id(image_name)
        if owner is not None and owner == current_user.user_id:
            continue

        # Check board visibility
        board_id = ApiDependencies.invoker.services.board_image_records.get_board_for_image(image_name)
        if board_id is not None:
            try:
                board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id)
                if board.board_visibility in (BoardVisibility.Shared, BoardVisibility.Public):
                    continue
            except Exception:
                pass

        raise HTTPException(status_code=403, detail=f"Not authorized to access image {image_name}")


@recall_parameters_router.post(
    "/{queue_id}",
    operation_id="update_recall_parameters",
    response_model=dict[str, Any],
)
async def update_recall_parameters(
    current_user: CurrentUserOrDefault,
    queue_id: str = Path(..., description="The queue id to perform this operation on"),
    parameters: RecallParameter = Body(..., description="Recall parameters to update"),
    strict: bool = Query(
        default=False,
        description="When true, parameters not included in the request are reset to their defaults (cleared).",
    ),
    append: bool = Query(
        default=False,
        description=(
            "When true, recalled reference images (ip_adapters and reference_images) are "
            "appended to the frontend's existing reference-image list instead of replacing it. "
            "Mutually exclusive with strict."
        ),
    ),
) -> dict[str, Any]:
    """
    Update recallable parameters that can be recalled on the frontend.

    This endpoint allows updating parameters such as prompt, model, steps, and other
    generation settings. These parameters are stored in client state and can be
    accessed by the frontend to populate UI elements.

    Args:
        queue_id: The queue ID to associate these parameters with
        parameters: The RecallParameter object containing the parameters to update
        strict: When true, parameters not included in the request body are reset
            to their defaults (cleared on the frontend).  Defaults to false,
            which preserves the existing behaviour of only updating the
            parameters that are explicitly provided.
        append: When true, recalled reference images (``ip_adapters`` and
            ``reference_images``) are appended to whatever reference images the
            frontend already has, instead of replacing the whole list.  Mutually
            exclusive with ``strict`` (which clears omitted parameters).

    Returns:
        A dictionary containing the updated parameters and status

    Example:
        POST /api/v1/recall/{queue_id}?strict=true
        {
            "positive_prompt": "a beautiful landscape",
            "model": "sd-1.5",
            "steps": 20
        }
        # In strict mode, all other parameters (reference_images, loras, etc.)
        # are cleared.  In non-strict mode (default) they would be left as-is.
    """
    logger = ApiDependencies.invoker.services.logger

    if strict and append:
        raise HTTPException(
            status_code=400,
            detail="The 'strict' and 'append' query parameters are mutually exclusive",
        )

    # Validate image access before processing — prevents information leakage
    # (dimensions) and derived-image minting via ControlNet preprocessors.
    _assert_recall_image_access(parameters, current_user)
    assert_image_move_maintenance_inactive()

    try:
        # In strict mode, include all parameters so the frontend clears anything
        # not explicitly provided.  List-typed fields use [] instead of None so
        # the frontend sees an empty collection rather than a null it might skip.
        if strict:
            _list_fields = {
                name for name, field in RecallParameter.model_fields.items() if "list" in str(field.annotation).lower()
            }
            provided_params = {
                k: ([] if v is None and k in _list_fields else v) for k, v in parameters.model_dump().items()
            }
        else:
            provided_params = {k: v for k, v in parameters.model_dump().items() if v is not None}

        if not provided_params:
            return {"status": "no_parameters_provided", "updated_count": 0}

        # Store each parameter in client state scoped to the current user
        updated_count = 0
        for param_key, param_value in provided_params.items():
            # Convert parameter values to JSON strings for storage
            value_str = json.dumps(param_value)
            try:
                ApiDependencies.invoker.services.client_state_persistence.set_by_key(
                    current_user.user_id, f"recall_{param_key}", value_str
                )
                updated_count += 1
            except Exception as e:
                logger.error(f"Error setting recall parameter {param_key}: {e}")
                raise HTTPException(
                    status_code=500,
                    detail=f"Error setting recall parameter {param_key}",
                )

        logger.info(f"Updated {updated_count} recall parameters for queue {queue_id}")

        # Resolve model name to key if a model was provided
        if "model" in provided_params and isinstance(provided_params["model"], str):
            model_name = provided_params["model"]
            model_key = resolve_model_name_to_key(model_name, ModelType.Main)

            if model_key:
                logger.info(f"Resolved model name '{model_name}' to key '{model_key}'")
                provided_params["model"] = model_key
            else:
                logger.warning(f"Could not resolve model name '{model_name}' to a model key")
                # Remove model from parameters if we couldn't resolve it
                del provided_params["model"]

        # Process LoRAs if provided
        if "loras" in provided_params:
            loras_param = parameters.loras
            if loras_param is not None:
                resolved_loras = resolve_lora_models(loras_param)
                provided_params["loras"] = resolved_loras
                logger.info(f"Resolved {len(resolved_loras)} LoRA(s)")

        # Process control layers if provided
        if "control_layers" in provided_params:
            control_layers_param = parameters.control_layers
            if control_layers_param is not None:
                resolved_controls = resolve_control_models(control_layers_param)
                provided_params["control_layers"] = resolved_controls
                logger.info(f"Resolved {len(resolved_controls)} control layer(s)")

        # Process IP adapters if provided
        if "ip_adapters" in provided_params:
            ip_adapters_param = parameters.ip_adapters
            if ip_adapters_param is not None:
                resolved_adapters = resolve_ip_adapter_models(ip_adapters_param)
                provided_params["ip_adapters"] = resolved_adapters
                logger.info(f"Resolved {len(resolved_adapters)} IP adapter(s)")

        # Process model-free reference images if provided
        if "reference_images" in provided_params:
            reference_images_param = parameters.reference_images
            if reference_images_param is not None:
                resolved_refs = resolve_reference_images(reference_images_param)
                provided_params["reference_images"] = resolved_refs
                logger.info(f"Resolved {len(resolved_refs)} reference image(s)")

        # Append mode rides along inside the event's parameters dict rather
        # than as a new event field so the generated client schema (which
        # types parameters as a free-form object) doesn't need regenerating.
        # Added after the persistence loop above, so the flag itself is never
        # stored as a recall parameter.
        if append:
            provided_params["append"] = True

        # Emit event to notify frontend of parameter updates
        try:
            logger.info(
                f"Emitting recall_parameters_updated event for queue {queue_id} with {len(provided_params)} parameters"
            )
            ApiDependencies.invoker.services.events.emit_recall_parameters_updated(
                queue_id, current_user.user_id, provided_params
            )
            logger.info("Successfully emitted recall_parameters_updated event")
        except Exception as e:
            logger.error(f"Error emitting recall parameters event: {e}", exc_info=True)
            # Don't fail the request if event emission fails, just log it

        return {
            "status": "success",
            "queue_id": queue_id,
            "updated_count": updated_count,
            "parameters": provided_params,
        }

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error updating recall parameters: {e}")
        raise HTTPException(
            status_code=500,
            detail="Error updating recall parameters",
        )


@recall_parameters_router.get(
    "/{queue_id}",
    operation_id="get_recall_parameters",
    response_model=dict[str, Any],
)
async def get_recall_parameters(
    current_user: CurrentUserOrDefault,
    queue_id: str = Path(..., description="The queue id to retrieve parameters for"),
) -> dict[str, Any]:
    """
    Retrieve all stored recall parameters for a given queue.

    Returns a dictionary of all recall parameters that have been set for the queue.

    Args:
        queue_id: The queue ID to retrieve parameters for

    Returns:
        A dictionary containing all stored recall parameters
    """
    logger = ApiDependencies.invoker.services.logger

    try:
        # Retrieve all recall parameters by iterating through expected keys
        # Since client_state_persistence doesn't have a "get_all" method, we'll
        # return an informative response
        return {
            "status": "success",
            "queue_id": queue_id,
            "note": "Use the frontend to access stored recall parameters, or set specific parameters using POST",
        }

    except Exception as e:
        logger.error(f"Error retrieving recall parameters: {e}")
        raise HTTPException(
            status_code=500,
            detail="Error retrieving recall parameters",
        )