How to create customized sampler in extension script in version 1.6.1?

[Carzit]

I am trying to register a sampler in my extension script https://github.com/Carzit/sd-webui-samplers-scheduler, but there's something wrong when i upgrade SD WebUI from 1.5.2 to 1.6.1. (You can try this extension. It works in 1.5.2. Change the version to 1.6.1 and you will see the bug i described below)

It works well in 1.5.2, but in 1.6.1 the same script code will make all sampler section button text S. Regardless of which one is selected, the program will not run correctly. I am checking a1111's changelog. While the sampler section seems to just add a couple of new samplers according to him, he seems to have tweaked some of the architecture so that the original function to register sampler just produces confusing results.

My code implementation was based on https://github.com/hnmr293/sd-webui-sd-webui-DPMPP_2M_Karras_v2, which fails after the 1.6.1 update as described above.

Does anyone has a solution to add customized sampler by extension in 1.6.1? QAQ

related code :

def seniorious():
    # the sampler function

class KDiffusionSamplerLocal(K.KDiffusionSampler):
    
    def __init__(self,funcname: str,original_funcname: str,func,sd_model: nn.Module):
        denoiser = k_diffusion.external.CompVisVDenoiser if sd_model.parameterization == "v" else k_diffusion.external.CompVisDenoiser

        self.model_wrap = denoiser(sd_model, quantize=shared.opts.enable_quantization)
        self.funcname = funcname
        self.func = func
        self.extra_params = K.sampler_extra_params.get(original_funcname, [])
        self.model_wrap_cfg = K.CFGDenoiser(self.model_wrap)
        self.sampler_noises = None
        self.stop_at = None
        self.eta = None
        self.config = None
        self.last_latent = None

        self.conditioning_key = sd_model.model.conditioning_key # type: ignore


def add_seniorious():
    original = [ x for x in K.samplers_k_diffusion if x[0] == 'Heun' ][0]
    o_label, o_constructor, o_aliases, o_options = original
    
    label = 'Seniorious'
    funcname = seniorious.__name__
    
    def constructor(model: nn.Module):
        return KDiffusionSamplerLocal(funcname, o_constructor, seniorious, model)
    
    aliases = ['seniorious']
    options = { **o_options }

    data = sd_samplers_common.SamplerData(label, constructor, aliases, options)
    if len([ x for x in sd_samplers.all_samplers if x.name == label ]) == 0:
        sd_samplers.all_samplers.append(data)


def hook(fn):
    
    @functools.wraps(fn)
    def f(*args, **kwargs):
        old_samplers, mode, *rest = args
        
        if mode not in ['txt2img', 'img2img']:
            print(f'unknown mode: {mode}', file=sys.stderr)
            return fn(*args, **kwargs)
        
        update_samplers()
        
        new_samplers = (
            sd_samplers.samplers if mode == 'txt2img' else
            sd_samplers.samplers_for_img2img
        )
        
        return fn(new_samplers, mode, *rest, **kwargs)
    
    return f


# register new sampler
add_seniorious()
update_samplers()


# hook Sampler textbox creation
from modules import ui
ui.create_sampler_and_steps_selection = hook(ui.create_sampler_and_steps_selection)

[catboxanon]

I haven't taken a look at the code yet, but my guess is it would be related to the refactor done to pass samplers by name rather than index. See 70c63c1

[Carzit]

Thank you so much. I will check the relevant code. (●'◡'●)