Custom ksampler

[Fenish]

When i run this custom node ksampler finishes job in like 10 minutes etc. how can i fix that why it took soo long. im using Ksampler's default sample function. When i use normal ksampler for example text to image promt its fast but when i use ksampler class to generaete custom pipeline i got this slowness. Slowness sometimes took like 30 mins etc.

class AITRepeater(KSampler):
    RETURN_TYPES = (
        "LATENT",
    )

    CATEGORY = "AIT NODES/ Repeater"
    FUNCTION = "execute"

    def __init__(self):    
        super().__init__()
        self.tile_size = 768
        self.overlap_size = 256

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True, "tooltip": "The random seed used for creating the noise."}),
                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The algorithm used when sampling, this can affect the quality, speed, and style of the generated output."}),
                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}),
                "positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to include in the image."}),
                "negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to exclude from the image."}),
                "latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}),
                "width_shift_rate": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
                "height_shift_rate": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
                "brush_size": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10000.0, "step": 0.01}),
            }
        }

    
    def execute(self, model, seed, sampler_name, scheduler, positive, negative, latent_image, width_shift_rate, height_shift_rate, brush_size):
        print("model: ", model)
        print("seed: ", seed)
        print("sampler_name: ", sampler_name)
        print("scheduler: ", scheduler)
        print("latent_image: ", latent_image)
        print("width_shift_rate: ", width_shift_rate)
        print("height_shift_rate: ", height_shift_rate)
        print("brush_size: ", brush_size)
        
        self.model = model
        self.pipe = self.sample
        self.seed = seed
        self.sampler_name = sampler_name
        self.scheduler = scheduler
        self.positive = positive
        self.negative = negative
        self.latent_image = latent_image
        self.width_shift_rate = width_shift_rate
        self.height_shift_rate = height_shift_rate
        self.brush_size = brush_size

        pil_image = tensor2pil(latent_image["samples"])
        np_image = np.array(pil_image)


        # Convert pil image to np.darray
        generated_image = self.repeater_pipeline(np_image, width_shift_rate, height_shift_rate, brush_size, seed)
        # generatedToTensor = pil2tensor(generated_image)
        return (generated_image, )
    
    def _inference(self
                   input_image: PIL.Image,
                   input_mask: PIL.Image,
                   ):
        latent = {}
        latent['samples'] = pil2tensor(input_image)
        latent['noise_mask'] = input_mask
        

        output_image = self.pipe(
            model=self.model,
            seed=self.seed,
            steps=35,
            cfg=15,
            sampler_name=self.sampler_name,
            scheduler=self.scheduler,
            positive=self.positive,
            negative=self.negative,
            latent_image=latent,
            denoise=0.99,
        )
        
        return output_image['samples']

[Fenish]

Also my utility functions

def tensor2pil(image):
    # Check if image has a specific shape pattern (like noise_mask)
    if len(image.shape) == 3 and image.shape[0] == 1 and image.shape[1] == 1:
        # Handle noise mask case (1, 1, width) shape
        return PIL.Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
    # Check dimensions and handle accordingly
    if len(image.shape) == 4:
        # Handle batched images (B, C, H, W)
        image = image.squeeze(0)  # Remove batch dimension
    
    # Convert tensor to numpy array, scale and convert to uint8
    img_np = np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
    
    # If the result is a 3D array with first dimension of size 3 or 4, it's likely in CHW format
    if len(img_np.shape) == 3 and img_np.shape[0] in [3, 4]:
        # Convert from CHW to HWC format
        img_np = np.transpose(img_np, (1, 2, 0))
    
    return PIL.Image.fromarray(img_np)

def pil2tensor(image):
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)

[ltdrdata]

Although you inherited from KSampler, you are not using KSampler's function, but instead using a completely different function called self.repeater_pipeline. Without any information about how repeater_pipeline is implemented, it is impossible to make a judgment.

[Fenish]

repeater_pipeline just calls _inference thats why i didn't put it in the question but i guess i found the issue and fixed it