diff --git a/__init__.py b/__init__.py index e7a3980..f432dbd 100644 --- a/__init__.py +++ b/__init__.py @@ -1,15 +1,17 @@ from .flux_resolution_cal_node import FluxResolutionNode from .flux_sampler_node import FluxSampler +from .flux_controlnet_node import FluxControlNetApply -print("Initializing Flux Resolution Node and Flux Sampler Node...") +print("Initializing Flux Resolution Node, Flux Sampler Node, and Flux ControlNet Apply Node") -# Registering the nodes NODE_CLASS_MAPPINGS = { "FluxResolutionNode": FluxResolutionNode, "FluxSampler": FluxSampler, + "FluxControlNetApply": FluxControlNetApply, } NODE_DISPLAY_NAME_MAPPINGS = { "FluxResolutionNode": "Flux Resolution Calculator", "FluxSampler": "Flux Sampler", + "FluxControlNetApply": "Flux ControlNet", } \ No newline at end of file diff --git a/flux_controlnet_node.py b/flux_controlnet_node.py new file mode 100644 index 0000000..4c62d47 --- /dev/null +++ b/flux_controlnet_node.py @@ -0,0 +1,40 @@ +class FluxControlNetApply: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "conditioning": ("CONDITIONING", ), + "control_net": ("CONTROL_NET", ), + "image": ("IMAGE", ), + "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}) + } + } + + RETURN_TYPES = ("CONDITIONING",) + FUNCTION = "flux_controlnet" + CATEGORY = "ControlAltAI Nodes/Flux" + + def flux_controlnet(self, conditioning, control_net, image, strength): + if strength == 0: + return (conditioning,) + + c = [] + control_hint = image.movedim(-1, 1) + for t in conditioning: + n = [t[0], t[1].copy()] + c_net = control_net.copy().set_cond_hint(control_hint, strength) + if 'control' in t[1]: + c_net.set_previous_controlnet(t[1]['control']) + n[1]['control'] = c_net + n[1]['control_apply_to_uncond'] = False # This ensures it's only applied to positive + c.append(n) + + return (c,) + +NODE_CLASS_MAPPINGS = { + "FluxControlNetApply": FluxControlNetApply, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FluxControlNetApply": "Flux ControlNet", +} \ No newline at end of file diff --git a/flux_sampler_node.py b/flux_sampler_node.py index 5e9c40b..2cd3449 100644 --- a/flux_sampler_node.py +++ b/flux_sampler_node.py @@ -4,7 +4,7 @@ import comfy.sample import latent_preview FLUX_SAMPLER_NAMES = [ - "euler", "heun", "heunpp2", "dpm_2", "dpm_2_ancestral", "lms", "dpmpp_2m", + "euler", "heun", "heunpp2", "dpm_2", "lms", "dpmpp_2m", "ipndm", "ipndm_v", "deis", "ddim", "uni_pc", "uni_pc_bh2" ] @@ -18,11 +18,11 @@ class FluxSampler: "model": ("MODEL",), "conditioning": ("CONDITIONING",), "latent_image": ("LATENT",), - "sampler_name": (FLUX_SAMPLER_NAMES,), - "scheduler": (FLUX_SCHEDULER_NAMES,), - "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + "sampler_name": (FLUX_SAMPLER_NAMES, {"default": "heunpp2"}), + "scheduler": (FLUX_SCHEDULER_NAMES, {"default": "beta"}), + "steps": ("INT", {"default": 42, "min": 1, "max": 10000}), "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "noise_seed": ("INT", {"default": 143220275975594, "min": 0, "max": 0xffffffffffffffff}), } } @@ -33,32 +33,28 @@ class FluxSampler: def sample(self, model, conditioning, latent_image, sampler_name, scheduler, steps, denoise, noise_seed): device = comfy.model_management.get_torch_device() - - # Create sampler sampler = comfy.samplers.KSampler(model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=denoise) - - # Prepare latent + latent = latent_image.copy() latent_image = latent["samples"] - - # Prepare noise + + # Handle noise_mask if present + noise_mask = latent.get("noise_mask", None) + noise = comfy.sample.prepare_noise(latent_image, noise_seed) - # Prepare conditioning positive = conditioning negative = [] # Empty list for negative conditioning - - # Sampling + callback = latent_preview.prepare_callback(model, steps) disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED - samples = sampler.sample(noise, positive, negative, cfg=1.0, latent_image=latent_image, - force_full_denoise=True, denoise_mask=None, callback=callback, + samples = sampler.sample(noise, positive, negative, cfg=1.0, latent_image=latent_image, + force_full_denoise=True, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed) out = latent.copy() out["samples"] = samples - return (out,) NODE_CLASS_MAPPINGS = {