try:
    import smea_sampling_kohak
    from smea_sampling_kohak import (
        sample_euler_dy_old, 
        sample_euler_smea_dy_old, 
        sample_euler_negative_old, 
        sample_euler_dy_negative_old, 
        sample_Kohaku_LoNyu_Yog_old,
        kohaku_lonyu_yog_stochastic_old,
        kohaku_lonyu_yog_compatible_old,
        sample_Kohaku_LoNyu_Yog_v2_old,
        kohaku_lonyu_yog_geo_compatible_old,
        kohaku_lonyu_yog_dy_old
    )

    if smea_sampling_kohak.BACKEND == "WebUI":
        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):
                if not smea_sampling_kohak.INITIALIZED:
                    samplers_smea = [
                        ("Euler Dy Old", sample_euler_dy_old, ["k_euler_dy_old"], {}),
                        ("Euler SMEA Dy Old", sample_euler_smea_dy_old, ["k_euler_smea_dy_old"], {}),
                        ("Euler Negative Old", sample_euler_negative_old, ["k_euler_negative_old"], {}),
                        ("Euler Negative Dy Old", sample_euler_dy_negative_old, ["k_euler_negative_dy_old"], {}),
                        ("Kohaku LoNyu Yog Old", sample_Kohaku_LoNyu_Yog_old, ["k_euler_Kohaku_LoNyu_Yog_old"], {}),
                        ("Kohaku LoNyu Yog Stochastic Old", kohaku_lonyu_yog_stochastic_old, ["k_euler_Kohaku_LoNyu_Yog_stochastic_old"], {}),
                        ("Kohaku LoNyu Yog Compatible Old", kohaku_lonyu_yog_compatible_old, ["k_euler_Kohaku_LoNyu_Yog_compatible_old"], {}),
                        ("Kohaku LoNyu Yog v2 Old", sample_Kohaku_LoNyu_Yog_v2_old, ["k_euler_Kohaku_LoNyu_Yog_v2_old"], {}),
                        ("Kohaku LoNyu Yog Geo Compatible Old", kohaku_lonyu_yog_geo_compatible_old, ["k_euler_Kohaku_LoNyu_Yog_geo_compatible_old"], {}),
                        ("Kohaku LoNyu Yog Dy Old", kohaku_lonyu_yog_dy_old, ["k_euler_Kohaku_LoNyu_Yog_dy_old"], {}),
                    ]
                    samplers_data_smea = [
                        sd_samplers_common.SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases, options)
                        for label, funcname, aliases, options in samplers_smea
                        if callable(funcname)
                    ]
                    sampler_extra_params["sample_euler_dy_old"] = ["s_churn", "s_tmin", "s_tmax", "s_noise"]
                    sampler_extra_params["sample_euler_smea_dy_old"] = ["s_churn", "s_tmin", "s_tmax", "s_noise"]
                    sampler_extra_params["sample_euler_negative_old"] = ["s_churn", "s_tmin", "s_tmax", "s_noise"]
                    sampler_extra_params["sample_euler_dy_negative_old"] = ["s_churn", "s_tmin", "s_tmax", "s_noise"]
                    sampler_extra_params["sample_Kohaku_LoNyu_Yog_old"] = ["s_churn", "s_tmin", "s_tmax", "s_noise", "eta"]
                    sampler_extra_params["kohaku_lonyu_yog_stochastic_old"] = ["langevin_strength"]
                    sampler_extra_params["kohaku_lonyu_yog_compatible_old"] = []
                    sampler_extra_params["sample_Kohaku_LoNyu_Yog_v2_old"] = ["s_churn", "s_tmin", "s_tmax", "s_noise", "eta_start", "eta_end", "use_normals"]
                    sampler_extra_params["kohaku_lonyu_yog_geo_compatible_old"] = []
                    sampler_extra_params["kohaku_lonyu_yog_dy_old"] = ["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()
                    smea_sampling_kohak.INITIALIZED = True

except ImportError as e:
    print(f"Failed to initialize SMEA samplers: {e}")