def METHOD_NAME(self, model, ratio):
    self.u = None
    def tomesd_m(q, k, v, extra_options):
        #NOTE: In the reference code get_functions takes x (input of the transformer block) as the argument instead of q
        #however from my basic testing it seems that using q instead gives better results
        m, self.u = get_functions(q, ratio, extra_options["original_shape"])
        return m(q), k, v
    def tomesd_u(n, extra_options):
        return self.u(n)
    m = model.clone()
    m.set_model_attn1_patch(tomesd_m)
    m.set_model_attn1_output_patch(tomesd_u)
    return (m, )