try:
    from smea_sampling_test import (
        sample_Kohaku_LoNyu_Yog_v1_test,
        kohaku_lonyu_yog_stochastic_v1_test,
        kohaku_lonyu_yog_compatible_v1_test,
        sample_Kohaku_LoNyu_Yog_v2_v1_test,
        kohaku_lonyu_yog_geo_compatible_v1_test,
        kohaku_lonyu_yog_dy_v1_test
    )

    from modules import scripts, sd_samplers_common, sd_samplers
    from modules.sd_samplers_kdiffusion import sampler_extra_params, KDiffusionSampler

    class SMEA(scripts.Script):
        def title(self):
            return "SMEA Samplers"

        def show(self, is_img2img):
            return False

        def __init__(self):
            smea_sampling_test = [
                ("Kohaku LoNyu Yog v1", sample_Kohaku_LoNyu_Yog_v1_test, ["k_kohaku_lonyu_yog_v1"], {}),
                ("Kohaku LoNyu Yog Stochastic v1", kohaku_lonyu_yog_stochastic_v1_test, ["k_kohaku_lonyu_yog_stochastic_v1"], {}),
                ("Kohaku LoNyu Yog Compatible v1", kohaku_lonyu_yog_compatible_v1_test, ["k_kohaku_lonyu_yog_compatible_v1"], {}),
                ("Kohaku LoNyu Yog v2 v1", sample_Kohaku_LoNyu_Yog_v2_v1_test, ["k_kohaku_lonyu_yog_v2_v1"], {}),
                ("Kohaku LoNyu Yog Geo Compatible v1", kohaku_lonyu_yog_geo_compatible_v1_test, ["k_kohaku_lonyu_yog_geo_compatible_v1"], {}),
                ("Kohaku LoNyu Yog Dy v1", kohaku_lonyu_yog_dy_v1_test, ["k_kohaku_lonyu_yog_dy_v1"], {}),
            ]
            samplers_data_smea = [
                sd_samplers_common.SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases, options)
                for label, funcname, aliases, options in smea_sampling_test
                if callable(funcname)
            ]
            sampler_extra_params["sample_Kohaku_LoNyu_Yog_v1_test"] = ["s_churn", "s_tmin", "s_tmax", "s_noise", "eta"]
            sampler_extra_params["kohaku_lonyu_yog_stochastic_v1_test"] = ["langevin_strength"]
            sampler_extra_params["kohaku_lonyu_yog_compatible_v1_test"] = []
            sampler_extra_params["sample_Kohaku_LoNyu_Yog_v2_v1_test"] = ["s_churn", "s_tmin", "s_tmax", "s_noise", "eta_start", "eta_end", "use_normals"]
            sampler_extra_params["kohaku_lonyu_yog_geo_compatible_v1_test"] = []
            sampler_extra_params["kohaku_lonyu_yog_dy_v1_test"] = ["s_churn", "s_tmin", "s_tmax", "s_noise"]
            sd_samplers.all_samplers.extend(samplers_data_smea)
            sd_samplers.all_samplers_map = {x.name: x for x in sd_samplers.all_samplers}
            sd_samplers.set_samplers()

except ImportError as _:
    pass