mlflow--mlflow/tests/diffusers/test_diffusers_model_export.py

import importlib.util
import json
from pathlib import Path

import numpy as np
import pytest
import yaml

pytest.importorskip("diffusers")
pytest.importorskip("safetensors")

from unittest.mock import MagicMock, Mock, patch

from safetensors.numpy import save_file

import mlflow
import mlflow.diffusers
from mlflow.diffusers import FLAVOR_NAME, DiffusersAdapterModel
from mlflow.exceptions import MlflowException
from mlflow.models.model import MLMODEL_FILE_NAME

BASE_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"


def _create_fake_adapter(tmp_path, filename="adapter.safetensors"):
    adapter_dir = tmp_path / "fake_adapter"
    adapter_dir.mkdir(exist_ok=True)
    tensors = {"lora_weight": np.random.randn(4, 4).astype(np.float32)}
    save_file(tensors, str(adapter_dir / filename))
    return adapter_dir


@pytest.fixture
def adapter_dir(tmp_path):
    return _create_fake_adapter(tmp_path)


@pytest.fixture
def adapter_file(tmp_path):
    adapter_dir = _create_fake_adapter(tmp_path)
    return adapter_dir / "adapter.safetensors"


@pytest.fixture
def model_path(tmp_path):
    return tmp_path / "model_output"


def test_save_model_from_directory(adapter_dir, model_path):
    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
    )

    mlmodel_path = model_path / MLMODEL_FILE_NAME
    assert mlmodel_path.exists()

    with open(mlmodel_path) as f:
        mlmodel = yaml.safe_load(f)

    assert FLAVOR_NAME in mlmodel["flavors"]
    assert "python_function" in mlmodel["flavors"]

    flavor_conf = mlmodel["flavors"][FLAVOR_NAME]
    assert flavor_conf["base_model"] == BASE_MODEL_ID
    assert flavor_conf["adapter_type"] == "lora"
    assert flavor_conf["adapter_weights"] == "adapter_weights"


def test_save_model_normalizes_single_file_input(adapter_file, model_path):
    mlflow.diffusers.save_model(
        adapter_path=str(adapter_file),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
    )

    weights_dir = model_path / "adapter_weights"
    assert weights_dir.exists()
    assert (weights_dir / "pytorch_lora_weights.safetensors").exists()


def test_save_model_normalizes_single_file_directory(adapter_dir, model_path):
    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
    )

    weights_dir = model_path / "adapter_weights"
    assert weights_dir.exists()
    assert (weights_dir / "pytorch_lora_weights.safetensors").exists()


def test_save_model_writes_environment_files(adapter_dir, model_path):
    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
    )

    assert (model_path / "conda.yaml").exists()
    assert (model_path / "requirements.txt").exists()
    assert (model_path / "python_env.yaml").exists()


def test_save_model_default_signature(adapter_dir, model_path):
    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
    )

    with open(model_path / MLMODEL_FILE_NAME) as f:
        mlmodel = yaml.safe_load(f)

    assert "signature" in mlmodel
    sig = mlmodel["signature"]
    inputs = json.loads(sig["inputs"])
    assert inputs[0]["name"] == "prompt"
    assert inputs[0]["type"] == "string"


def test_save_model_custom_signature(adapter_dir, model_path):
    from mlflow.types import DataType, Schema
    from mlflow.types.schema import ColSpec

    custom_sig = mlflow.models.ModelSignature(
        inputs=Schema([ColSpec(type=DataType.string, name="text")]),
        outputs=Schema([ColSpec(type=DataType.binary, name="img")]),
    )

    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
        signature=custom_sig,
    )

    with open(model_path / MLMODEL_FILE_NAME) as f:
        mlmodel = yaml.safe_load(f)

    inputs = json.loads(mlmodel["signature"]["inputs"])
    assert inputs[0]["name"] == "text"


def test_save_model_invalid_adapter_type(adapter_dir, model_path):
    with pytest.raises(mlflow.exceptions.MlflowException, match="Unsupported adapter type"):
        mlflow.diffusers.save_model(
            adapter_path=str(adapter_dir),
            path=str(model_path),
            base_model=BASE_MODEL_ID,
            adapter_type="invalid",
        )


def test_save_model_nonexistent_path(model_path):
    with pytest.raises(mlflow.exceptions.MlflowException, match="does not exist"):
        mlflow.diffusers.save_model(
            adapter_path="/nonexistent/path",
            path=str(model_path),
            base_model=BASE_MODEL_ID,
        )


def test_save_model_with_metadata(adapter_dir, model_path):
    test_metadata = {"training_dataset": "my-dataset", "lora_rank": 16}

    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
        metadata=test_metadata,
    )

    with open(model_path / MLMODEL_FILE_NAME) as f:
        mlmodel = yaml.safe_load(f)

    assert mlmodel["metadata"] == test_metadata


def test_log_model(adapter_dir):
    with mlflow.start_run() as run:
        model_info = mlflow.diffusers.log_model(
            adapter_path=str(adapter_dir),
            base_model=BASE_MODEL_ID,
            name="test_adapter",
        )

    assert model_info is not None

    client = mlflow.MlflowClient()
    artifacts = [a.path for a in client.list_artifacts(run.info.run_id, "test_adapter")]
    assert any("adapter_weights" in a for a in artifacts)
    assert any(MLMODEL_FILE_NAME in a for a in artifacts)


def test_save_load_model_direct_roundtrip(adapter_dir, model_path):
    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
    )

    loaded = mlflow.diffusers.load_model(str(model_path))

    assert isinstance(loaded, DiffusersAdapterModel)
    assert loaded.base_model == BASE_MODEL_ID
    assert loaded.adapter_type == "lora"
    assert Path(loaded.adapter_path).exists()


def test_base_model_revision_roundtrip(adapter_dir, model_path):

    fake_revision = "abc123def456"
    with patch("mlflow.diffusers._resolve_base_model_revision", return_value=fake_revision):
        mlflow.diffusers.save_model(
            adapter_path=str(adapter_dir),
            path=str(model_path),
            base_model=BASE_MODEL_ID,
        )

    # Verify revision stored in flavor config
    mlmodel_path = model_path / MLMODEL_FILE_NAME
    with open(mlmodel_path) as f:
        mlmodel = yaml.safe_load(f)
    assert mlmodel["flavors"][FLAVOR_NAME]["base_model_revision"] == fake_revision

    # Verify revision survives load_model roundtrip
    loaded = mlflow.diffusers.load_model(str(model_path))
    assert loaded.base_model_revision == fake_revision


def test_load_model_via_tracking_roundtrip(adapter_dir):
    with mlflow.start_run() as run:
        mlflow.diffusers.log_model(
            adapter_path=str(adapter_dir),
            base_model=BASE_MODEL_ID,
            name="test_adapter",
        )

    model_uri = f"runs:/{run.info.run_id}/test_adapter"
    loaded = mlflow.diffusers.load_model(model_uri)

    assert isinstance(loaded, DiffusersAdapterModel)
    assert loaded.base_model == BASE_MODEL_ID
    assert loaded.adapter_type == "lora"
    assert Path(loaded.adapter_path).exists()
    assert (Path(loaded.adapter_path) / "pytorch_lora_weights.safetensors").exists()


def test_load_pyfunc_returns_wrapper(adapter_dir):
    from mlflow.diffusers.wrapper import _DiffusersAdapterWrapper

    with mlflow.start_run() as run:
        mlflow.diffusers.log_model(
            adapter_path=str(adapter_dir),
            base_model=BASE_MODEL_ID,
            name="test_adapter",
        )

    model_uri = f"runs:/{run.info.run_id}/test_adapter"
    loaded_pyfunc = mlflow.pyfunc.load_model(model_uri)

    assert hasattr(loaded_pyfunc, "predict")
    wrapper = loaded_pyfunc._model_impl
    assert isinstance(wrapper, _DiffusersAdapterWrapper)
    assert wrapper._flavor_conf["base_model"] == BASE_MODEL_ID


def test_load_pipeline_base_model_override(adapter_dir, model_path):

    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
    )

    loaded = mlflow.diffusers.load_model(str(model_path))
    override = "other-org/other-model"

    with patch("diffusers.DiffusionPipeline.from_pretrained") as mock_fp:
        mock_pipe = MagicMock()
        mock_fp.return_value = mock_pipe
        loaded.load_pipeline(base_model=override)
        mock_fp.assert_called_once()
        assert mock_fp.call_args[0][0] == override


def test_load_pipeline_wraps_oserror(adapter_dir, model_path):

    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
    )

    loaded = mlflow.diffusers.load_model(str(model_path))

    with patch(
        "diffusers.DiffusionPipeline.from_pretrained",
        side_effect=OSError("is not a local folder"),
    ):
        with pytest.raises(MlflowException, match="Failed to load base model"):
            loaded.load_pipeline()


def test_wrapper_model_config_base_model_override():
    from mlflow.diffusers.wrapper import _DiffusersAdapterWrapper

    w = _DiffusersAdapterWrapper(
        adapter_path="/fake",
        flavor_conf={"base_model": "original/model", "adapter_type": "lora"},
        model_config={"base_model": "override/model"},
    )

    with patch("diffusers.DiffusionPipeline.from_pretrained") as mock_fp:
        mock_pipe = MagicMock()
        mock_fp.return_value = mock_pipe
        w._load_pipeline()
        mock_fp.assert_called_once()
        assert mock_fp.call_args[0][0] == "override/model"


@pytest.mark.parametrize("package", ["diffusers", "transformers", "torch", "peft", "safetensors"])
def test_default_pip_requirements_contains_core(package):
    reqs = mlflow.diffusers.get_default_pip_requirements()
    req_names = [r.split("==")[0] for r in reqs]
    assert package in req_names


@pytest.mark.parametrize("package", ["accelerate"])
def test_default_pip_requirements_contains_optional_if_installed(package):
    reqs = mlflow.diffusers.get_default_pip_requirements()
    req_names = [r.split("==")[0] for r in reqs]
    if importlib.util.find_spec(package):
        assert package in req_names
    else:
        assert package not in req_names


def test_default_conda_env():
    env = mlflow.diffusers.get_default_conda_env()
    assert "dependencies" in env


# -- predict() tests (mock-based, no GPU) --


class _FakeImage:
    def save(self, buf, format=None):
        buf.write(b"\x89PNG\r\n\x1a\n" + b"\x00" * 100)


class _FakePipelineOutput:
    def __init__(self, n=1):
        self.images = [_FakeImage() for _ in range(n)]


@pytest.fixture
def wrapper():
    from mlflow.diffusers.wrapper import _DiffusersAdapterWrapper

    w = _DiffusersAdapterWrapper(
        adapter_path="/fake",
        flavor_conf={"base_model": "test/model", "adapter_type": "lora"},
    )
    # Inject a mock pipeline so _load_pipeline is never called
    w._pipeline = lambda prompt, **kwargs: _FakePipelineOutput(len(prompt))
    return w


def test_get_raw_model(wrapper):
    model = wrapper.get_raw_model()
    assert model is not None
    assert callable(model)


def test_predict_dataframe(wrapper):
    import pandas as pd

    result = wrapper.predict(pd.DataFrame({"prompt": ["a cat", "a dog"]}))
    assert len(result) == 2
    assert all(isinstance(r, bytes) for r in result)


def test_predict_string(wrapper):
    result = wrapper.predict("a cat")
    assert len(result) == 1
    assert result[0][:4] == b"\x89PNG"


def test_predict_dict(wrapper):
    result = wrapper.predict({"prompt": "a cat"})
    assert len(result) == 1


def test_predict_dict_list(wrapper):
    result = wrapper.predict({"prompt": ["a cat", "a dog"]})
    assert len(result) == 2


def test_predict_list(wrapper):
    result = wrapper.predict(["a cat", "a dog"])
    assert len(result) == 2


def test_predict_with_params(wrapper):

    mock_pipeline = Mock(return_value=_FakePipelineOutput(1))
    wrapper._pipeline = mock_pipeline

    result = wrapper.predict(
        "a cat",
        params={"num_inference_steps": 10, "height": 256, "negative_prompt": "blurry"},
    )
    assert len(result) == 1

    mock_pipeline.assert_called_once()
    call_kwargs = mock_pipeline.call_args.kwargs
    assert call_kwargs["num_inference_steps"] == 10
    assert call_kwargs["height"] == 256
    assert call_kwargs["negative_prompt"] == "blurry"
    assert call_kwargs["prompt"] == ["a cat"]


def test_predict_dataframe_missing_prompt(wrapper):
    import pandas as pd

    with pytest.raises(MlflowException, match="prompt"):
        wrapper.predict(pd.DataFrame({"text": ["a cat"], "other": ["extra"]}))


def test_predict_dict_missing_prompt(wrapper):
    with pytest.raises(MlflowException, match="prompt"):
        wrapper.predict({"text": "a cat"})


def test_predict_empty_prompts(wrapper):
    with pytest.raises(MlflowException, match="No prompts"):
        wrapper.predict([])


def test_predict_dict_invalid_prompt_type(wrapper):
    with pytest.raises(MlflowException, match="must be a string or list of strings"):
        wrapper.predict({"prompt": 42})


def test_predict_unsupported_type(wrapper):
    with pytest.raises(MlflowException, match="Unsupported input type"):
        wrapper.predict(12345)


def test_predict_no_images(wrapper):
    class _NoImagesOutput:
        images = None

    wrapper._pipeline = lambda prompt, **kwargs: _NoImagesOutput()

    with pytest.raises(MlflowException, match="Pipeline returned no images"):
        wrapper.predict("a cat")


def test_predict_empty_images_list(wrapper):
    class _EmptyImagesOutput:
        images = []

    wrapper._pipeline = lambda prompt, **kwargs: _EmptyImagesOutput()

    with pytest.raises(MlflowException, match="Pipeline returned no images"):
        wrapper.predict("a cat")


def test_predict_none_prompt_rejected(wrapper):
    with pytest.raises(MlflowException, match="must be strings, not None"):
        wrapper.predict([None])


def test_predict_none_in_dataframe_rejected(wrapper):
    import pandas as pd

    with pytest.raises(MlflowException, match="must be strings, not None"):
        wrapper.predict(pd.DataFrame({"prompt": [None]}))


# -- _resolve_base_model_revision tests --


def test_resolve_revision_absolute_path_returns_none():
    from mlflow.diffusers import _resolve_base_model_revision

    assert _resolve_base_model_revision("/absolute/path/to/model") is None


def test_resolve_revision_dot_relative_path_returns_none():
    from mlflow.diffusers import _resolve_base_model_revision

    assert _resolve_base_model_revision("./local/model") is None
    assert _resolve_base_model_revision("../parent/model") is None


def test_resolve_revision_hf_hub_success():

    from mlflow.diffusers import _resolve_base_model_revision

    fake_sha = "abc123"
    with patch(
        "mlflow.utils.huggingface_utils.get_latest_commit_for_repo",
        return_value=fake_sha,
    ):
        result = _resolve_base_model_revision("org/model-name")
    assert result == fake_sha


def test_resolve_revision_hf_hub_failure_returns_none():

    from mlflow.diffusers import _resolve_base_model_revision

    with patch(
        "mlflow.utils.huggingface_utils.get_latest_commit_for_repo",
        side_effect=Exception("network error"),
    ):
        result = _resolve_base_model_revision("org/model-name")
    assert result is None


def test_resolve_revision_cwd_collision_still_resolves(tmp_path, monkeypatch):

    from mlflow.diffusers import _resolve_base_model_revision

    (tmp_path / "org" / "model").mkdir(parents=True)
    monkeypatch.chdir(tmp_path)

    fake_sha = "deadbeef"
    with patch(
        "mlflow.utils.huggingface_utils.get_latest_commit_for_repo",
        return_value=fake_sha,
    ):
        result = _resolve_base_model_revision("org/model")
    assert result == fake_sha


def test_save_model_multi_file_directory_preserved(tmp_path, model_path):
    adapter_dir = tmp_path / "multi_adapter"
    adapter_dir.mkdir()
    tensors = {"w": np.random.randn(4, 4).astype(np.float32)}
    save_file(tensors, str(adapter_dir / "pytorch_lora_weights.safetensors"))
    # Companion file (e.g., text encoder LoRA or adapter_config from PEFT)
    (adapter_dir / "adapter_config.json").write_text('{"type": "lora"}')

    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
    )

    weights_dir = model_path / "adapter_weights"
    assert (weights_dir / "pytorch_lora_weights.safetensors").exists()
    assert (weights_dir / "adapter_config.json").exists()


def test_save_model_records_weight_name_for_nonstandard_files(tmp_path, model_path):
    adapter_dir = tmp_path / "multi_nonstandard"
    adapter_dir.mkdir()
    tensors = {"w": np.random.randn(4, 4).astype(np.float32)}
    save_file(tensors, str(adapter_dir / "alpha_weights.safetensors"))
    save_file(tensors, str(adapter_dir / "beta_weights.safetensors"))

    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
    )

    mlmodel = yaml.safe_load((model_path / MLMODEL_FILE_NAME).read_text())
    flavor_conf = mlmodel["flavors"]["diffusers"]
    assert flavor_conf["weight_name"] == "alpha_weights.safetensors"

    loaded = mlflow.diffusers.load_model(str(model_path))
    assert loaded.weight_name == "alpha_weights.safetensors"


def test_save_model_records_weight_name_for_peft_adapter(tmp_path, model_path):
    adapter_dir = tmp_path / "peft_adapter"
    adapter_dir.mkdir()
    tensors = {"w": np.random.randn(4, 4).astype(np.float32)}
    save_file(tensors, str(adapter_dir / "adapter_model.safetensors"))
    (adapter_dir / "adapter_config.json").write_text('{"type": "lora"}')

    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
    )

    mlmodel = yaml.safe_load((model_path / MLMODEL_FILE_NAME).read_text())
    flavor_conf = mlmodel["flavors"]["diffusers"]
    assert flavor_conf["weight_name"] == "adapter_model.safetensors"

    loaded = mlflow.diffusers.load_model(str(model_path))
    assert loaded.weight_name == "adapter_model.safetensors"


def test_save_model_rejects_non_safetensors_file(tmp_path, model_path):
    bad_file = tmp_path / "model.bin"
    bad_file.write_bytes(b"\x00" * 100)

    with pytest.raises(MlflowException, match=".safetensors"):
        mlflow.diffusers.save_model(
            adapter_path=str(bad_file),
            path=str(model_path),
            base_model=BASE_MODEL_ID,
        )


def test_save_model_rejects_invalid_safetensors_content(tmp_path, model_path):
    fake_file = tmp_path / "bad.safetensors"
    fake_file.write_bytes(b"this is not safetensors format")

    with pytest.raises(MlflowException, match="not a valid safetensors"):
        mlflow.diffusers.save_model(
            adapter_path=str(fake_file),
            path=str(model_path),
            base_model=BASE_MODEL_ID,
        )


@pytest.mark.parametrize(
    ("base_model", "match"),
    [
        ("", "non-empty"),
        ("   ", "non-empty"),
        (None, "must be a"),
        (123, "must be a"),
    ],
)
def test_save_model_rejects_invalid_base_model(adapter_dir, model_path, base_model, match):
    with pytest.raises(MlflowException, match=match):
        mlflow.diffusers.save_model(
            adapter_path=str(adapter_dir),
            path=str(model_path),
            base_model=base_model,
        )


@pytest.mark.parametrize("adapter_type", [None, 123, []])
def test_save_model_rejects_non_string_adapter_type(adapter_dir, model_path, adapter_type):
    with pytest.raises(MlflowException, match="adapter_type must be a string"):
        mlflow.diffusers.save_model(
            adapter_path=str(adapter_dir),
            path=str(model_path),
            base_model=BASE_MODEL_ID,
            adapter_type=adapter_type,
        )


# -- _detect_device tests --


def test_detect_device_explicit():
    from mlflow.diffusers import _detect_device

    assert _detect_device("cpu") == "cpu"
    assert _detect_device("cuda:1") == "cuda:1"


def test_detect_device_env_var(monkeypatch):
    from mlflow.diffusers import _detect_device

    monkeypatch.setenv("MLFLOW_DEFAULT_PREDICTION_DEVICE", "cpu")
    assert _detect_device() == "cpu"


def test_predict_single_column_dataframe(wrapper):
    import pandas as pd

    result = wrapper.predict(pd.DataFrame(["a cat", "a dog"]))
    assert len(result) == 2


def test_save_model_rejects_empty_directory(tmp_path, model_path):
    empty_dir = tmp_path / "empty_adapter"
    empty_dir.mkdir()
    (empty_dir / "readme.txt").write_text("no weights here")

    with pytest.raises(MlflowException, match="no .safetensors"):
        mlflow.diffusers.save_model(
            adapter_path=str(empty_dir),
            path=str(model_path),
            base_model=BASE_MODEL_ID,
        )


def test_save_model_ignores_hidden_files(tmp_path, model_path):
    adapter_dir = tmp_path / "adapter_with_ds_store"
    adapter_dir.mkdir()
    tensors = {"w": np.random.randn(4, 4).astype(np.float32)}
    save_file(tensors, str(adapter_dir / "my_lora.safetensors"))
    (adapter_dir / ".DS_Store").write_bytes(b"\x00" * 10)

    mlflow.diffusers.save_model(
        adapter_path=str(adapter_dir),
        path=str(model_path),
        base_model=BASE_MODEL_ID,
    )

    weights_dir = model_path / "adapter_weights"
    assert (weights_dir / "pytorch_lora_weights.safetensors").exists()


def test_save_model_validates_all_safetensors_in_multi_file_dir(tmp_path, model_path):
    adapter_dir = tmp_path / "corrupt_multi"
    adapter_dir.mkdir()
    tensors = {"w": np.random.randn(4, 4).astype(np.float32)}
    save_file(tensors, str(adapter_dir / "pytorch_lora_weights.safetensors"))
    (adapter_dir / "corrupt.safetensors").write_bytes(b"not valid safetensors")

    with pytest.raises(MlflowException, match="not a valid safetensors"):
        mlflow.diffusers.save_model(
            adapter_path=str(adapter_dir),
            path=str(model_path),
            base_model=BASE_MODEL_ID,
        )