"""Model construction helpers for Clover Image Tiny inpainting."""

from __future__ import annotations

from pathlib import Path
from typing import Any

import torch
from diffusers import UNet2DConditionModel


def _load_kwargs(revision: str | None) -> dict[str, Any]:
    return {"revision": revision} if revision else {}


def make_inpainting_unet(
    base_model: str | Path,
    *,
    revision: str | None = None,
    zero_initialize_conditioning: bool = True,
) -> UNet2DConditionModel:
    """Create a 9-channel U-Net from the 4-channel Clover checkpoint.

    The first four input channels retain Clover's original weights. The extra
    channels receive the binary inpainting mask and the masked-image latent.
    Zero-initializing them preserves a stable text-to-image starting point while
    the inpainting fine-tune learns how to use the new conditioning channels.
    """

    load_kwargs = _load_kwargs(revision)
    base = UNet2DConditionModel.from_pretrained(
        str(base_model),
        subfolder="unet",
        low_cpu_mem_usage=False,
        **load_kwargs,
    )
    config = dict(base.config)
    config["in_channels"] = 9
    inpaint = UNet2DConditionModel.from_config(config)

    state = base.state_dict()
    conv_in_weight = state.pop("conv_in.weight")
    inpaint.load_state_dict(state, strict=False)

    with torch.no_grad():
        if zero_initialize_conditioning:
            inpaint.conv_in.weight.zero_()
            inpaint.conv_in.weight[:, :4].copy_(conv_in_weight)
        else:
            # Keep the pretrained channels and leave the five new channels at
            # the framework's default initialization.
            inpaint.conv_in.weight[:, :4].copy_(conv_in_weight)
        if "conv_in.bias" in base.state_dict():
            inpaint.conv_in.bias.copy_(base.conv_in.bias)

    del base
    return inpaint