Qwen-Image-Layered-MRP-MLX/tests/model_config/test_model_config.py

import pytest

from mflux.config.model_config import ModelConfig
from mflux.utils.exceptions import InvalidBaseModel, ModelConfigError


def test_bfl_dev():
    model = ModelConfig.from_name("dev")
    assert model.model_name == "black-forest-labs/FLUX.1-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_bfl_dev_full_name():
    model = ModelConfig.from_name("black-forest-labs/FLUX.1-dev")
    assert model.model_name == "black-forest-labs/FLUX.1-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_bfl_krea_dev():
    model = ModelConfig.from_name("krea-dev")
    assert model.model_name == "black-forest-labs/FLUX.1-Krea-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_bfl_dev_krea_alias():
    model = ModelConfig.from_name("dev-krea")
    assert model.model_name == "black-forest-labs/FLUX.1-Krea-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_bfl_krea_dev_full_name():
    model = ModelConfig.from_name("black-forest-labs/FLUX.1-Krea-dev")
    assert model.model_name == "black-forest-labs/FLUX.1-Krea-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_bfl_schnell():
    model = ModelConfig.from_name("schnell")
    assert model.model_name == "black-forest-labs/FLUX.1-schnell"
    assert model.max_sequence_length == 256
    assert model.num_train_steps == 1000
    assert model.supports_guidance is False
    assert model.requires_sigma_shift is False


def test_bfl_schnell_full_name():
    model = ModelConfig.from_name("black-forest-labs/FLUX.1-schnell")
    assert model.model_name == "black-forest-labs/FLUX.1-schnell"
    assert model.max_sequence_length == 256
    assert model.num_train_steps == 1000
    assert model.supports_guidance is False
    assert model.requires_sigma_shift is False


def test_bfl_dev_fill():
    model = ModelConfig.from_name("dev-fill")
    assert model.model_name == "black-forest-labs/FLUX.1-Fill-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_bfl_dev_fill_full_name():
    model = ModelConfig.from_name("black-forest-labs/FLUX.1-Fill-dev")
    assert model.model_name == "black-forest-labs/FLUX.1-Fill-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_bfl_dev_depth():
    model = ModelConfig.from_name("dev-depth")
    assert model.model_name == "black-forest-labs/FLUX.1-Depth-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_bfl_dev_depth_full_name():
    model = ModelConfig.from_name("black-forest-labs/FLUX.1-Depth-dev")
    assert model.model_name == "black-forest-labs/FLUX.1-Depth-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_bfl_dev_redux():
    model = ModelConfig.from_name("dev-redux")
    assert model.model_name == "black-forest-labs/FLUX.1-Redux-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_bfl_dev_redux_full_name():
    model = ModelConfig.from_name("black-forest-labs/FLUX.1-Redux-dev")
    assert model.model_name == "black-forest-labs/FLUX.1-Redux-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_bfl_dev_controlnet_canny():
    model = ModelConfig.from_name("dev-controlnet-canny")
    assert model.model_name == "black-forest-labs/FLUX.1-dev"
    assert model.controlnet_model == "InstantX/FLUX.1-dev-Controlnet-Canny"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True
    assert model.is_canny() is True


def test_bfl_schnell_controlnet_canny():
    model = ModelConfig.from_name("schnell-controlnet-canny")
    assert model.model_name == "black-forest-labs/FLUX.1-schnell"
    assert model.controlnet_model == "InstantX/FLUX.1-dev-Controlnet-Canny"
    assert model.max_sequence_length == 256
    assert model.num_train_steps == 1000
    assert model.supports_guidance is False
    assert model.requires_sigma_shift is False
    assert model.is_canny() is True


def test_bfl_dev_controlnet_upscaler():
    model = ModelConfig.from_name("dev-controlnet-upscaler")
    assert model.model_name == "black-forest-labs/FLUX.1-dev"
    assert model.controlnet_model == "jasperai/Flux.1-dev-Controlnet-Upscaler"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is False
    assert model.requires_sigma_shift is False
    assert model.is_canny() is False


def test_community_dev_fill_implicit_base_model():
    model = ModelConfig.from_name("acme-lab/some-dev-fill-model")
    assert model.base_model == "black-forest-labs/FLUX.1-Fill-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_community_dev_fill_explicit_base_model():
    model = ModelConfig.from_name("acme-lab/some-model", base_model="dev-fill")
    assert model.base_model == "black-forest-labs/FLUX.1-Fill-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_implicit_base_model_prefers_dev_fill_over_dev():
    model = ModelConfig.from_name("acme-lab/dev-fill-based-model")
    assert model.base_model == "black-forest-labs/FLUX.1-Fill-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.requires_sigma_shift is True


def test_community_dev_implicit_base_model():
    model = ModelConfig.from_name("acme-lab/some-awesome-dev-model")
    assert model.base_model == "black-forest-labs/FLUX.1-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_community_schnell_implicit_base_model():
    model = ModelConfig.from_name("acme-lab/some-quick-schnell-model")
    assert model.base_model == "black-forest-labs/FLUX.1-schnell"
    assert model.max_sequence_length == 256
    assert model.num_train_steps == 1000
    assert model.supports_guidance is False
    assert model.requires_sigma_shift is False


def test_community_dev_explicit_base_model():
    model = ModelConfig.from_name("acme-lab/some-awesome-model", base_model="dev")
    assert model.base_model == "black-forest-labs/FLUX.1-dev"
    assert model.max_sequence_length == 512
    assert model.num_train_steps == 1000
    assert model.supports_guidance is True
    assert model.requires_sigma_shift is True


def test_community_schnell_explicit_base_model():
    model = ModelConfig.from_name("acme-lab/some-awesome-model", base_model="schnell")
    assert model.base_model == "black-forest-labs/FLUX.1-schnell"
    assert model.max_sequence_length == 256
    assert model.num_train_steps == 1000
    assert model.supports_guidance is False
    assert model.requires_sigma_shift is False


def test_model_config_error():
    with pytest.raises(ModelConfigError):
        ModelConfig.from_name("acme-lab/some-model-who-knows-what-its-based-on")


def test_invalid_base_model_error():
    with pytest.raises(InvalidBaseModel):
        ModelConfig.from_name("acme-lab/some-model-who-knows-what-its-based-on", base_model="something_unknown")


def test_qwen_image():
    model = ModelConfig.from_name("qwen")
    assert model.model_name == "Qwen/Qwen-Image"
    assert model.max_sequence_length is None
    assert model.num_train_steps is None
    assert model.supports_guidance is None
    assert model.requires_sigma_shift is None


def test_qwen_image_alias():
    model = ModelConfig.from_name("qwen-image")
    assert model.model_name == "Qwen/Qwen-Image"
    assert model.max_sequence_length is None
    assert model.num_train_steps is None
    assert model.supports_guidance is None
    assert model.requires_sigma_shift is None


def test_qwen_image_full_name():
    model = ModelConfig.from_name("Qwen/Qwen-Image")
    assert model.model_name == "Qwen/Qwen-Image"
    assert model.max_sequence_length is None
    assert model.num_train_steps is None
    assert model.supports_guidance is None
    assert model.requires_sigma_shift is None