import os
import yaml

from huggingface_hub import hf_hub_download

from src.models.sdxl.adapter.ip_adapter import (
    SDXLModelWithIPAdapter,
    SDXLModelWithIPAdapterConfig,
)


def _load_config(config_path: str, model_path: str, adapter_path: str):
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    config = SDXLModelWithIPAdapterConfig(**config)
    config.checkpoint_path = model_path
    config.adapter.checkpoint_weight = adapter_path

    return config


def load_ip_adapter_model(
    model_path: str,
    config_path: str,
    adapter_path: str,
):
    config = _load_config(config_path, model_path, adapter_path)
    model = SDXLModelWithIPAdapter.from_checkpoint(config)

    return model


def get_file_path(repo_id: str, filename: str, revision: str = "main") -> str:
    if os.path.exists(filename):
        return filename

    file_path = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        revision=revision,
    )

    return file_path