├── .github └── workflows │ └── publish.yml ├── LICENSE ├── README.md ├── __init__.py ├── nodes.py ├── pyproject.toml └── utilities.py /.github/workflows/publish.yml: -------------------------------------------------------------------------------- 1 | name: Publish to Comfy registry 2 | on: 3 | workflow_dispatch: 4 | push: 5 | branches: 6 | - main 7 | paths: 8 | - "pyproject.toml" 9 | 10 | jobs: 11 | publish-node: 12 | name: Publish Custom Node to registry 13 | runs-on: ubuntu-latest 14 | steps: 15 | - name: Check out code 16 | uses: actions/checkout@v4 17 | - name: Publish Custom Node 18 | uses: Comfy-Org/publish-node-action@main 19 | with: 20 | ## Add your own personal access token to your Github Repository secrets and reference it here. 21 | personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 bmad4ever 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ### comfyui_ab_sampler 2 | 3 | Experimental sampler node. 4 | 5 | Sampling alternates between A and B inputs until only one remains, starting with A. 6 | 7 | B steps run over a 2x2 grid, where 3/4's of the grid are copies of the original input latent. 8 | 9 | When the optional mask is used, the region outside the defined roi is copied from the 10 | original latent at the end of every step. 11 | 12 | **Disclaimer: its applications should be rather niche and there are likely better alternative approaches.** 13 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | from .nodes import NODE_CLASS_MAPPINGS 2 | __all__ = [NODE_CLASS_MAPPINGS] 3 | -------------------------------------------------------------------------------- /nodes.py: -------------------------------------------------------------------------------- 1 | import comfy 2 | import random 3 | import torch 4 | from .utilities import * 5 | import comfy_extras.nodes_mask 6 | from comfy_extras.nodes_custom_sampler import SamplerCustom 7 | 8 | 9 | class AB_SamplerCustom: 10 | @classmethod 11 | def INPUT_TYPES(s): 12 | types = SamplerCustom.INPUT_TYPES() 13 | types["required"].pop("positive") 14 | types["required"].pop("negative") 15 | types["required"]["cfgA"] = types["required"]["cfg"] 16 | types["required"]["cfgB"] = types["required"]["cfg"] 17 | types["required"]["sigmasA"] = types["required"]["sigmas"] 18 | types["required"]["sigmasB"] = types["required"]["sigmas"] 19 | types["required"].pop("cfg") 20 | types["required"].pop("model") 21 | types["required"].pop("sigmas") 22 | types["optional"] = {} 23 | 24 | types["required"]["modelA"] = ("MODEL",) 25 | types["optional"]["modelB"] = ("MODEL",) 26 | types["required"]["positive_A"] = ("CONDITIONING",) 27 | types["required"]["negative_A"] = ("CONDITIONING",) 28 | types["required"]["positive_B"] = ("CONDITIONING",) 29 | types["required"]["negative_B"] = ("CONDITIONING",) 30 | 31 | types["optional"]["roi_mask"] = ("MASK",) 32 | 33 | return types 34 | 35 | RETURN_TYPES = ("LATENT", "LATENT") 36 | RETURN_NAMES = ("output", "denoised_output") 37 | FUNCTION = "sample" 38 | CATEGORY = "Bmad/experimental" 39 | 40 | def sample(self, modelA, add_noise, noise_seed, 41 | cfgA, cfgB, positive_A, negative_A, positive_B, negative_B, 42 | sampler, sigmasA, sigmasB, latent_image, modelB=None, roi_mask=None): 43 | if modelB is None: 44 | modelB = modelA 45 | 46 | latent = latent_image 47 | latent_image = latent["samples"] 48 | latent_image_o = None if roi_mask is None else latent_image.clone() 49 | latent_grid = repeat_into_grid(latent, 2, 2) 50 | latent_grid_o = latent_grid.clone() 51 | _, _, height, width = latent_image.size() 52 | 53 | empty_noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, 54 | device="cpu") 55 | noise = empty_noise 56 | fake_noise_grid = torch.zeros(latent_grid.size(), dtype=latent_image.dtype, layout=latent_image.layout, 57 | device="cpu") 58 | if add_noise: 59 | batch_inds = latent["batch_index"] if "batch_index" in latent else None 60 | noise = comfy.sample.prepare_noise(latent_image, noise_seed, batch_inds) 61 | 62 | noise_mask = None 63 | if "noise_mask" in latent: 64 | noise_mask = latent["noise_mask"] 65 | 66 | x0_output = {} 67 | # callback = latent_preview.prepare_callback(modelA, sigmasA.shape[-1]+sigmasB.shape[-1] - 2, x0_output) 68 | total_steps = sigmasA.shape[-1] + sigmasB.shape[-1] - 2 # last sigma is zero 69 | pbar = comfy.utils.ProgressBar(total_steps) 70 | 71 | disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED 72 | 73 | composite_node = None if roi_mask is None else comfy_extras.nodes_mask.LatentCompositeMasked() 74 | random.seed(noise_seed) 75 | # random_seeds = [random.randint(0, sys.maxsize) for _ in range(sigmasA.shape[-1]+sigmasB.shape[-1])] 76 | min_steps_ab = min(sigmasA.shape[-1], sigmasB.shape[-1]) - 1 77 | for i in range(0, min_steps_ab): 78 | # A step 79 | target_latent = comfy.sample.sample_custom( 80 | modelA, noise, cfgA, sampler, sigmasA[i:i + 2], positive_A, negative_A, latent_image, 81 | noise_mask=noise_mask, callback=None, disable_pbar=disable_pbar, seed=noise_seed) 82 | # noise_seed = random_seeds.pop() 83 | 84 | if i == 0: 85 | noise = empty_noise 86 | 87 | latent_image, latent_grid = setup_latents_ab(composite_node, width, height, 88 | latent_image, latent_image_o, 89 | latent_grid, latent_grid_o, 90 | roi_mask, target_latent, 91 | this_step=STEP_A, next_step=STEP_B) 92 | 93 | # B step 94 | target_latent = comfy.sample.sample_custom( 95 | modelB, fake_noise_grid, cfgB, sampler, sigmasB[i:i + 2], positive_B, negative_B, latent_grid, 96 | noise_mask=noise_mask, callback=None, disable_pbar=disable_pbar, seed=noise_seed) 97 | 98 | # check if last step, and, if so, change to STEP_B if B has more steps 99 | next_step = STEP_A 100 | if i + 1 == min_steps_ab and sigmasB.shape[-1] > sigmasA.shape[-1]: 101 | next_step = STEP_B 102 | latent_image, latent_grid = setup_latents_ab(composite_node, width, height, 103 | latent_image, latent_image_o, 104 | latent_grid, latent_grid_o, 105 | roi_mask, target_latent, 106 | this_step=STEP_B, next_step=next_step) 107 | # noise_seed = random_seeds.pop() 108 | pbar.update_absolute(i * 2, total_steps) 109 | 110 | # TAIL (only missing A or B steps) 111 | if sigmasA.shape[-1] != sigmasB.shape[-1]: # tail, only A or B steps 112 | tail_step_type = STEP_A if sigmasA.shape[-1] > sigmasB.shape[-1] else STEP_B 113 | tail_sigmas, tail_model, tail_cfg, tail_pos, tail_neg, tail_noise = \ 114 | (sigmasA, modelA, cfgA, positive_A, negative_A, noise) \ 115 | if tail_step_type is STEP_A \ 116 | else (sigmasB, modelB, cfgB, positive_B, negative_B, fake_noise_grid) 117 | 118 | for i in range(min_steps_ab, max(sigmasA.shape[-1], sigmasB.shape[-1]) - 1): 119 | target_latent = latent_image if tail_step_type is STEP_A else latent_grid 120 | target_latent = comfy.sample.sample_custom( 121 | tail_model, tail_noise, tail_cfg, sampler, tail_sigmas[i:i + 2], tail_pos, tail_neg, target_latent, 122 | noise_mask=noise_mask, callback=None, disable_pbar=disable_pbar, seed=noise_seed) 123 | 124 | latent_image, latent_grid = setup_latents_ab(composite_node, width, height, 125 | latent_image, latent_image_o, 126 | latent_grid, latent_grid_o, 127 | roi_mask, target_latent, 128 | this_step=tail_step_type, next_step=tail_step_type) 129 | 130 | pbar.update_absolute(min_steps_ab + i, total_steps) 131 | # noise_seed = random_seeds.pop() 132 | 133 | out = latent.copy() 134 | out["samples"] = latent_image 135 | if "x0" in x0_output: 136 | out_denoised = latent.copy() 137 | out_denoised["samples"] = modelA.model.process_latent_out(x0_output["x0"].cpu()) 138 | else: 139 | out_denoised = out 140 | return (out, out_denoised) 141 | 142 | 143 | NODE_CLASS_MAPPINGS = { 144 | "AB SamplerCustom (experimental)": AB_SamplerCustom, 145 | } 146 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "comfyui_ab_samplercustom" 3 | description = "Experimental sampler node. Sampling alternates between A and B inputs until only one remains, starting with A. B steps run over a 2x2 grid, where 3/4's of the grid are copies of the original input latent. When the optional mask is used, the region outside the defined roi is copied from the original latent at the end of every step. Disclaimer: its applications should be rather niche and there are likely better alternative approaches." 4 | version = "1.0.0" 5 | license = { file = "LICENSE" } 6 | 7 | [project.urls] 8 | Repository = "https://github.com/bmad4ever/comfyui_ab_samplercustom" 9 | # Used by Comfy Registry https://comfyregistry.org 10 | 11 | [tool.comfy] 12 | PublisherId = "bmad4ever" 13 | DisplayName = "comfyui_ab_samplercustom" 14 | Icon = "" 15 | -------------------------------------------------------------------------------- /utilities.py: -------------------------------------------------------------------------------- 1 | STEP_A = True 2 | STEP_B = False 3 | 4 | 5 | def repeat_into_grid(samples, columns, rows): 6 | return samples['samples'].repeat(1, 1, rows, columns) 7 | 8 | 9 | def setup_latents_ab(composite_node, width, height, latent_image, latent_image_o, latent_grid, latent_grid_o, 10 | roi_mask, target_latent, this_step, next_step): 11 | # get current step result in the same size auxiliary latent 12 | if this_step is STEP_A: 13 | latent_image = target_latent 14 | else: 15 | latent_grid = target_latent 16 | latent_image[0, :, :, :] = latent_grid[0, :, 0:height, 0:width] 17 | 18 | # composite with the original using the provided mask 19 | if roi_mask is not None: 20 | latent_image = composite_node.composite( 21 | {"samples": latent_image_o}, {"samples": latent_image}, 0, 0, False, roi_mask)[0]["samples"] 22 | 23 | # NOTE: latent_image is already ready for A step 24 | 25 | # setup grid latent for B step 26 | if next_step is STEP_B: 27 | latent_grid[:, :, :, :] = latent_grid_o[:, :, :, :] 28 | latent_grid[0, :, 0:height, 0:width] = latent_image[0, :, :, :] 29 | 30 | return latent_image, latent_grid 31 | --------------------------------------------------------------------------------