diff --git a/comfy_extras/v3/nodes_ace.py b/comfy_extras/v3/nodes_ace.py index 988e0ed5a..fdad8f800 100644 --- a/comfy_extras/v3/nodes_ace.py +++ b/comfy_extras/v3/nodes_ace.py @@ -52,6 +52,6 @@ class EmptyAceStepLatentAudio(io.ComfyNode): NODES_LIST: list[type[io.ComfyNode]] = [ - TextEncodeAceStepAudio, EmptyAceStepLatentAudio, + TextEncodeAceStepAudio, ] diff --git a/comfy_extras/v3/nodes_advanced_samplers.py b/comfy_extras/v3/nodes_advanced_samplers.py index 91512effb..ecbe7094f 100644 --- a/comfy_extras/v3/nodes_advanced_samplers.py +++ b/comfy_extras/v3/nodes_advanced_samplers.py @@ -122,7 +122,7 @@ class SamplerEulerCFGpp(io.ComfyNode): return io.NodeOutput(sampler) -NODES_LIST = [ - SamplerLCMUpscale, +NODES_LIST: list[type[io.ComfyNode]] = [ SamplerEulerCFGpp, + SamplerLCMUpscale, ] diff --git a/comfy_extras/v3/nodes_align_your_steps.py b/comfy_extras/v3/nodes_align_your_steps.py index c2a211c99..acb71c631 100644 --- a/comfy_extras/v3/nodes_align_your_steps.py +++ b/comfy_extras/v3/nodes_align_your_steps.py @@ -5,6 +5,18 @@ import torch from comfy_api.latest import io + +def loglinear_interp(t_steps, num_steps): + """Performs log-linear interpolation of a given array of decreasing numbers.""" + xs = np.linspace(0, 1, len(t_steps)) + ys = np.log(t_steps[::-1]) + + new_xs = np.linspace(0, 1, num_steps) + new_ys = np.interp(new_xs, xs, ys) + + return np.exp(new_ys)[::-1].copy() + + NOISE_LEVELS = { "SD1": [ 14.6146412293, @@ -36,17 +48,6 @@ NOISE_LEVELS = { } -def loglinear_interp(t_steps, num_steps): - """Performs log-linear interpolation of a given array of decreasing numbers.""" - xs = np.linspace(0, 1, len(t_steps)) - ys = np.log(t_steps[::-1]) - - new_xs = np.linspace(0, 1, num_steps) - new_ys = np.interp(new_xs, xs, ys) - - return np.exp(new_ys)[::-1].copy() - - class AlignYourStepsScheduler(io.ComfyNode): @classmethod def define_schema(cls) -> io.Schema: @@ -78,6 +79,6 @@ class AlignYourStepsScheduler(io.ComfyNode): return io.NodeOutput(torch.FloatTensor(sigmas)) -NODES_LIST = [ +NODES_LIST: list[type[io.ComfyNode]] = [ AlignYourStepsScheduler, ] diff --git a/comfy_extras/v3/nodes_apg.py b/comfy_extras/v3/nodes_apg.py index 961bdddb3..f9fc208d0 100644 --- a/comfy_extras/v3/nodes_apg.py +++ b/comfy_extras/v3/nodes_apg.py @@ -93,6 +93,6 @@ class APG(io.ComfyNode): return io.NodeOutput(m) -NODES_LIST = [ +NODES_LIST: list[type[io.ComfyNode]] = [ APG, ] diff --git a/comfy_extras/v3/nodes_attention_multiply.py b/comfy_extras/v3/nodes_attention_multiply.py index 9fb86714f..dfcb85568 100644 --- a/comfy_extras/v3/nodes_attention_multiply.py +++ b/comfy_extras/v3/nodes_attention_multiply.py @@ -131,9 +131,9 @@ class UNetTemporalAttentionMultiply(io.ComfyNode): return io.NodeOutput(m) -NODES_LIST = [ - UNetSelfAttentionMultiply, - UNetCrossAttentionMultiply, +NODES_LIST: list[type[io.ComfyNode]] = [ CLIPAttentionMultiply, + UNetCrossAttentionMultiply, + UNetSelfAttentionMultiply, UNetTemporalAttentionMultiply, ] diff --git a/comfy_extras/v3/nodes_audio.py b/comfy_extras/v3/nodes_audio.py index 3fb2fcc7d..089c2cb73 100644 --- a/comfy_extras/v3/nodes_audio.py +++ b/comfy_extras/v3/nodes_audio.py @@ -13,6 +13,28 @@ import node_helpers from comfy_api.latest import io, ui +class EmptyLatentAudio(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="EmptyLatentAudio_V3", + category="latent/audio", + inputs=[ + io.Float.Input("seconds", default=47.6, min=1.0, max=1000.0, step=0.1), + io.Int.Input( + "batch_size", default=1, min=1, max=4096, tooltip="The number of latent images in the batch." + ), + ], + outputs=[io.Latent.Output()], + ) + + @classmethod + def execute(cls, seconds, batch_size) -> io.NodeOutput: + length = round((seconds * 44100 / 2048) / 2) * 2 + latent = torch.zeros([batch_size, 64, length], device=comfy.model_management.intermediate_device()) + return io.NodeOutput({"samples": latent, "type": "audio"}) + + class ConditioningStableAudio(io.ComfyNode): @classmethod def define_schema(cls): @@ -43,26 +65,140 @@ class ConditioningStableAudio(io.ComfyNode): ) -class EmptyLatentAudio(io.ComfyNode): +class VAEEncodeAudio(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( - node_id="EmptyLatentAudio_V3", + node_id="VAEEncodeAudio_V3", category="latent/audio", inputs=[ - io.Float.Input("seconds", default=47.6, min=1.0, max=1000.0, step=0.1), - io.Int.Input( - "batch_size", default=1, min=1, max=4096, tooltip="The number of latent images in the batch." - ), + io.Audio.Input("audio"), + io.Vae.Input("vae"), ], outputs=[io.Latent.Output()], ) @classmethod - def execute(cls, seconds, batch_size) -> io.NodeOutput: - length = round((seconds * 44100 / 2048) / 2) * 2 - latent = torch.zeros([batch_size, 64, length], device=comfy.model_management.intermediate_device()) - return io.NodeOutput({"samples": latent, "type": "audio"}) + def execute(cls, vae, audio) -> io.NodeOutput: + sample_rate = audio["sample_rate"] + if 44100 != sample_rate: + waveform = torchaudio.functional.resample(audio["waveform"], sample_rate, 44100) + else: + waveform = audio["waveform"] + return io.NodeOutput({"samples": vae.encode(waveform.movedim(1, -1))}) + + +class VAEDecodeAudio(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="VAEDecodeAudio_V3", + category="latent/audio", + inputs=[ + io.Latent.Input("samples"), + io.Vae.Input("vae"), + ], + outputs=[io.Audio.Output()], + ) + + @classmethod + def execute(cls, vae, samples) -> io.NodeOutput: + audio = vae.decode(samples["samples"]).movedim(-1, 1) + std = torch.std(audio, dim=[1, 2], keepdim=True) * 5.0 + std[std < 1.0] = 1.0 + audio /= std + return io.NodeOutput({"waveform": audio, "sample_rate": 44100}) + + +class SaveAudio(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SaveAudio_V3", # frontend expects "SaveAudio" to work + display_name="Save Audio _V3", # frontend ignores "display_name" for this node + category="audio", + inputs=[ + io.Audio.Input("audio"), + io.String.Input("filename_prefix", default="audio/ComfyUI"), + ], + hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo], + is_output_node=True, + ) + + @classmethod + def execute(cls, audio, filename_prefix="ComfyUI", format="flac") -> io.NodeOutput: + return io.NodeOutput( + ui=ui.AudioSaveHelper.get_save_audio_ui(audio, filename_prefix=filename_prefix, cls=cls, format=format) + ) + + +class SaveAudioMP3(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SaveAudioMP3_V3", # frontend expects "SaveAudioMP3" to work + display_name="Save Audio(MP3) _V3", # frontend ignores "display_name" for this node + category="audio", + inputs=[ + io.Audio.Input("audio"), + io.String.Input("filename_prefix", default="audio/ComfyUI"), + io.Combo.Input("quality", options=["V0", "128k", "320k"], default="V0"), + ], + hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo], + is_output_node=True, + ) + + @classmethod + def execute(cls, audio, filename_prefix="ComfyUI", format="mp3", quality="V0") -> io.NodeOutput: + return io.NodeOutput( + ui=ui.AudioSaveHelper.get_save_audio_ui( + audio, filename_prefix=filename_prefix, cls=cls, format=format, quality=quality + ) + ) + + +class SaveAudioOpus(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SaveAudioOpus_V3", # frontend expects "SaveAudioOpus" to work + display_name="Save Audio(Opus) _V3", # frontend ignores "display_name" for this node + category="audio", + inputs=[ + io.Audio.Input("audio"), + io.String.Input("filename_prefix", default="audio/ComfyUI"), + io.Combo.Input("quality", options=["64k", "96k", "128k", "192k", "320k"], default="128k"), + ], + hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo], + is_output_node=True, + ) + + @classmethod + def execute(cls, audio, filename_prefix="ComfyUI", format="opus", quality="128k") -> io.NodeOutput: + return io.NodeOutput( + ui=ui.AudioSaveHelper.get_save_audio_ui( + audio, filename_prefix=filename_prefix, cls=cls, format=format, quality=quality + ) + ) + + +class PreviewAudio(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="PreviewAudio_V3", # frontend expects "PreviewAudio" to work + display_name="Preview Audio _V3", # frontend ignores "display_name" for this node + category="audio", + inputs=[ + io.Audio.Input("audio"), + ], + hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo], + is_output_node=True, + ) + + @classmethod + def execute(cls, audio) -> io.NodeOutput: + return io.NodeOutput(ui=ui.PreviewAudio(audio, cls=cls)) class LoadAudio(io.ComfyNode): @@ -141,150 +277,14 @@ class LoadAudio(io.ComfyNode): return True -class PreviewAudio(io.ComfyNode): - @classmethod - def define_schema(cls): - return io.Schema( - node_id="PreviewAudio_V3", # frontend expects "PreviewAudio" to work - display_name="Preview Audio _V3", # frontend ignores "display_name" for this node - category="audio", - inputs=[ - io.Audio.Input("audio"), - ], - hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo], - is_output_node=True, - ) - - @classmethod - def execute(cls, audio) -> io.NodeOutput: - return io.NodeOutput(ui=ui.PreviewAudio(audio, cls=cls)) - - -class SaveAudioMP3(io.ComfyNode): - @classmethod - def define_schema(cls): - return io.Schema( - node_id="SaveAudioMP3_V3", # frontend expects "SaveAudioMP3" to work - display_name="Save Audio(MP3) _V3", # frontend ignores "display_name" for this node - category="audio", - inputs=[ - io.Audio.Input("audio"), - io.String.Input("filename_prefix", default="audio/ComfyUI"), - io.Combo.Input("quality", options=["V0", "128k", "320k"], default="V0"), - ], - hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo], - is_output_node=True, - ) - - @classmethod - def execute(cls, audio, filename_prefix="ComfyUI", format="mp3", quality="V0") -> io.NodeOutput: - return io.NodeOutput( - ui=ui.AudioSaveHelper.get_save_audio_ui( - audio, filename_prefix=filename_prefix, cls=cls, format=format, quality=quality - ) - ) - - -class SaveAudioOpus(io.ComfyNode): - @classmethod - def define_schema(cls): - return io.Schema( - node_id="SaveAudioOpus_V3", # frontend expects "SaveAudioOpus" to work - display_name="Save Audio(Opus) _V3", # frontend ignores "display_name" for this node - category="audio", - inputs=[ - io.Audio.Input("audio"), - io.String.Input("filename_prefix", default="audio/ComfyUI"), - io.Combo.Input("quality", options=["64k", "96k", "128k", "192k", "320k"], default="128k"), - ], - hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo], - is_output_node=True, - ) - - @classmethod - def execute(cls, audio, filename_prefix="ComfyUI", format="opus", quality="128k") -> io.NodeOutput: - return io.NodeOutput( - ui=ui.AudioSaveHelper.get_save_audio_ui( - audio, filename_prefix=filename_prefix, cls=cls, format=format, quality=quality - ) - ) - - -class SaveAudio(io.ComfyNode): - @classmethod - def define_schema(cls): - return io.Schema( - node_id="SaveAudio_V3", # frontend expects "SaveAudio" to work - display_name="Save Audio _V3", # frontend ignores "display_name" for this node - category="audio", - inputs=[ - io.Audio.Input("audio"), - io.String.Input("filename_prefix", default="audio/ComfyUI"), - ], - hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo], - is_output_node=True, - ) - - @classmethod - def execute(cls, audio, filename_prefix="ComfyUI", format="flac") -> io.NodeOutput: - return io.NodeOutput( - ui=ui.AudioSaveHelper.get_save_audio_ui(audio, filename_prefix=filename_prefix, cls=cls, format=format) - ) - - -class VAEDecodeAudio(io.ComfyNode): - @classmethod - def define_schema(cls): - return io.Schema( - node_id="VAEDecodeAudio_V3", - category="latent/audio", - inputs=[ - io.Latent.Input("samples"), - io.Vae.Input("vae"), - ], - outputs=[io.Audio.Output()], - ) - - @classmethod - def execute(cls, vae, samples) -> io.NodeOutput: - audio = vae.decode(samples["samples"]).movedim(-1, 1) - std = torch.std(audio, dim=[1, 2], keepdim=True) * 5.0 - std[std < 1.0] = 1.0 - audio /= std - return io.NodeOutput({"waveform": audio, "sample_rate": 44100}) - - -class VAEEncodeAudio(io.ComfyNode): - @classmethod - def define_schema(cls): - return io.Schema( - node_id="VAEEncodeAudio_V3", - category="latent/audio", - inputs=[ - io.Audio.Input("audio"), - io.Vae.Input("vae"), - ], - outputs=[io.Latent.Output()], - ) - - @classmethod - def execute(cls, vae, audio) -> io.NodeOutput: - sample_rate = audio["sample_rate"] - if 44100 != sample_rate: - waveform = torchaudio.functional.resample(audio["waveform"], sample_rate, 44100) - else: - waveform = audio["waveform"] - return io.NodeOutput({"samples": vae.encode(waveform.movedim(1, -1))}) - - NODES_LIST: list[type[io.ComfyNode]] = [ ConditioningStableAudio, EmptyLatentAudio, LoadAudio, PreviewAudio, + SaveAudio, SaveAudioMP3, SaveAudioOpus, - SaveAudio, VAEDecodeAudio, VAEEncodeAudio, ] diff --git a/comfy_extras/v3/nodes_camera_trajectory.py b/comfy_extras/v3/nodes_camera_trajectory.py index edc159591..40fb1dcf9 100644 --- a/comfy_extras/v3/nodes_camera_trajectory.py +++ b/comfy_extras/v3/nodes_camera_trajectory.py @@ -212,6 +212,6 @@ class WanCameraEmbedding(io.ComfyNode): return io.NodeOutput(control_camera_video, width, height, length) -NODES_LIST = [ +NODES_LIST: list[type[io.ComfyNode]] = [ WanCameraEmbedding, ] diff --git a/comfy_extras/v3/nodes_canny.py b/comfy_extras/v3/nodes_canny.py index e24b0df38..0b68db381 100644 --- a/comfy_extras/v3/nodes_canny.py +++ b/comfy_extras/v3/nodes_canny.py @@ -27,6 +27,6 @@ class Canny(io.ComfyNode): return io.NodeOutput(img_out) -NODES_LIST = [ +NODES_LIST: list[type[io.ComfyNode]] = [ Canny, ] diff --git a/comfy_extras/v3/nodes_cfg.py b/comfy_extras/v3/nodes_cfg.py index 66ec27f9a..e8e84a2bd 100644 --- a/comfy_extras/v3/nodes_cfg.py +++ b/comfy_extras/v3/nodes_cfg.py @@ -5,6 +5,7 @@ import torch from comfy_api.latest import io +# https://github.com/WeichenFan/CFG-Zero-star def optimized_scale(positive, negative): positive_flat = positive.reshape(positive.shape[0], -1) negative_flat = negative.reshape(negative.shape[0], -1) @@ -21,6 +22,36 @@ def optimized_scale(positive, negative): return st_star.reshape([positive.shape[0]] + [1] * (positive.ndim - 1)) +class CFGZeroStar(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="CFGZeroStar_V3", + category="advanced/guidance", + inputs=[ + io.Model.Input("model"), + ], + outputs=[io.Model.Output(display_name="patched_model")], + ) + + @classmethod + def execute(cls, model) -> io.NodeOutput: + m = model.clone() + + def cfg_zero_star(args): + guidance_scale = args['cond_scale'] + x = args['input'] + cond_p = args['cond_denoised'] + uncond_p = args['uncond_denoised'] + out = args["denoised"] + alpha = optimized_scale(x - cond_p, x - uncond_p) + + return out + uncond_p * (alpha - 1.0) + guidance_scale * uncond_p * (1.0 - alpha) + + m.set_model_sampler_post_cfg_function(cfg_zero_star) + return io.NodeOutput(m) + + class CFGNorm(io.ComfyNode): @classmethod def define_schema(cls) -> io.Schema: @@ -52,37 +83,7 @@ class CFGNorm(io.ComfyNode): return io.NodeOutput(m) -class CFGZeroStar(io.ComfyNode): - @classmethod - def define_schema(cls) -> io.Schema: - return io.Schema( - node_id="CFGZeroStar_V3", - category="advanced/guidance", - inputs=[ - io.Model.Input("model"), - ], - outputs=[io.Model.Output(display_name="patched_model")], - ) - - @classmethod - def execute(cls, model) -> io.NodeOutput: - m = model.clone() - - def cfg_zero_star(args): - guidance_scale = args['cond_scale'] - x = args['input'] - cond_p = args['cond_denoised'] - uncond_p = args['uncond_denoised'] - out = args["denoised"] - alpha = optimized_scale(x - cond_p, x - uncond_p) - - return out + uncond_p * (alpha - 1.0) + guidance_scale * uncond_p * (1.0 - alpha) - - m.set_model_sampler_post_cfg_function(cfg_zero_star) - return io.NodeOutput(m) - - -NODES_LIST = [ +NODES_LIST: list[type[io.ComfyNode]] = [ CFGNorm, CFGZeroStar, ] diff --git a/comfy_extras/v3/nodes_clip_sdxl.py b/comfy_extras/v3/nodes_clip_sdxl.py index 54b83dc16..3d05b7595 100644 --- a/comfy_extras/v3/nodes_clip_sdxl.py +++ b/comfy_extras/v3/nodes_clip_sdxl.py @@ -4,6 +4,31 @@ import nodes from comfy_api.latest import io +class CLIPTextEncodeSDXLRefiner(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="CLIPTextEncodeSDXLRefiner_V3", + category="advanced/conditioning", + inputs=[ + io.Float.Input("ascore", default=6.0, min=0.0, max=1000.0, step=0.01), + io.Int.Input("width", default=1024, min=0, max=nodes.MAX_RESOLUTION), + io.Int.Input("height", default=1024, min=0, max=nodes.MAX_RESOLUTION), + io.String.Input("text", multiline=True, dynamic_prompts=True), + io.Clip.Input("clip"), + ], + outputs=[io.Conditioning.Output()], + ) + + @classmethod + def execute(cls, ascore, width, height, text, clip) -> io.NodeOutput: + tokens = clip.tokenize(text) + conditioning = clip.encode_from_tokens_scheduled( + tokens, add_dict={"aesthetic_score": ascore, "width": width, "height": height} + ) + return io.NodeOutput(conditioning) + + class CLIPTextEncodeSDXL(io.ComfyNode): @classmethod def define_schema(cls): @@ -48,32 +73,7 @@ class CLIPTextEncodeSDXL(io.ComfyNode): return io.NodeOutput(conditioning) -class CLIPTextEncodeSDXLRefiner(io.ComfyNode): - @classmethod - def define_schema(cls): - return io.Schema( - node_id="CLIPTextEncodeSDXLRefiner_V3", - category="advanced/conditioning", - inputs=[ - io.Float.Input("ascore", default=6.0, min=0.0, max=1000.0, step=0.01), - io.Int.Input("width", default=1024, min=0, max=nodes.MAX_RESOLUTION), - io.Int.Input("height", default=1024, min=0, max=nodes.MAX_RESOLUTION), - io.String.Input("text", multiline=True, dynamic_prompts=True), - io.Clip.Input("clip"), - ], - outputs=[io.Conditioning.Output()], - ) - - @classmethod - def execute(cls, ascore, width, height, text, clip) -> io.NodeOutput: - tokens = clip.tokenize(text) - conditioning = clip.encode_from_tokens_scheduled( - tokens, add_dict={"aesthetic_score": ascore, "width": width, "height": height} - ) - return io.NodeOutput(conditioning) - - -NODES_LIST = [ +NODES_LIST: list[type[io.ComfyNode]] = [ CLIPTextEncodeSDXL, CLIPTextEncodeSDXLRefiner, ] diff --git a/comfy_extras/v3/nodes_compositing.py b/comfy_extras/v3/nodes_compositing.py index cfe195148..b1e59ec78 100644 --- a/comfy_extras/v3/nodes_compositing.py +++ b/comfy_extras/v3/nodes_compositing.py @@ -112,32 +112,6 @@ def porter_duff_composite( return out_image, out_alpha -class JoinImageWithAlpha(io.ComfyNode): - @classmethod - def define_schema(cls): - return io.Schema( - node_id="JoinImageWithAlpha_V3", - display_name="Join Image with Alpha _V3", - category="mask/compositing", - inputs=[ - io.Image.Input("image"), - io.Mask.Input("alpha"), - ], - outputs=[io.Image.Output()], - ) - - @classmethod - def execute(cls, image: torch.Tensor, alpha: torch.Tensor) -> io.NodeOutput: - batch_size = min(len(image), len(alpha)) - out_images = [] - - alpha = 1.0 - resize_mask(alpha, image.shape[1:]) - for i in range(batch_size): - out_images.append(torch.cat((image[i][:, :, :3], alpha[i].unsqueeze(2)), dim=2)) - - return io.NodeOutput(torch.stack(out_images)) - - class PorterDuffImageComposite(io.ComfyNode): @classmethod def define_schema(cls): @@ -219,7 +193,33 @@ class SplitImageWithAlpha(io.ComfyNode): return io.NodeOutput(torch.stack(out_images), 1.0 - torch.stack(out_alphas)) -NODES_LIST = [ +class JoinImageWithAlpha(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="JoinImageWithAlpha_V3", + display_name="Join Image with Alpha _V3", + category="mask/compositing", + inputs=[ + io.Image.Input("image"), + io.Mask.Input("alpha"), + ], + outputs=[io.Image.Output()], + ) + + @classmethod + def execute(cls, image: torch.Tensor, alpha: torch.Tensor) -> io.NodeOutput: + batch_size = min(len(image), len(alpha)) + out_images = [] + + alpha = 1.0 - resize_mask(alpha, image.shape[1:]) + for i in range(batch_size): + out_images.append(torch.cat((image[i][:, :, :3], alpha[i].unsqueeze(2)), dim=2)) + + return io.NodeOutput(torch.stack(out_images)) + + +NODES_LIST: list[type[io.ComfyNode]] = [ JoinImageWithAlpha, PorterDuffImageComposite, SplitImageWithAlpha, diff --git a/comfy_extras/v3/nodes_controlnet.py b/comfy_extras/v3/nodes_controlnet.py index 4788113a4..a4656fad2 100644 --- a/comfy_extras/v3/nodes_controlnet.py +++ b/comfy_extras/v3/nodes_controlnet.py @@ -3,6 +3,33 @@ from comfy.cldm.control_types import UNION_CONTROLNET_TYPES from comfy_api.latest import io +class SetUnionControlNetType(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SetUnionControlNetType_V3", + category="conditioning/controlnet", + inputs=[ + io.ControlNet.Input("control_net"), + io.Combo.Input("type", options=["auto"] + list(UNION_CONTROLNET_TYPES.keys())), + ], + outputs=[ + io.ControlNet.Output(), + ], + ) + + @classmethod + def execute(cls, control_net, type) -> io.NodeOutput: + control_net = control_net.copy() + type_number = UNION_CONTROLNET_TYPES.get(type, -1) + if type_number >= 0: + control_net.set_extra_arg("control_type", [type_number]) + else: + control_net.set_extra_arg("control_type", []) + + return io.NodeOutput(control_net) + + class ControlNetApplyAdvanced(io.ComfyNode): @classmethod def define_schema(cls): @@ -60,33 +87,6 @@ class ControlNetApplyAdvanced(io.ComfyNode): return io.NodeOutput(out[0], out[1]) -class SetUnionControlNetType(io.ComfyNode): - @classmethod - def define_schema(cls): - return io.Schema( - node_id="SetUnionControlNetType_V3", - category="conditioning/controlnet", - inputs=[ - io.ControlNet.Input("control_net"), - io.Combo.Input("type", options=["auto"] + list(UNION_CONTROLNET_TYPES.keys())), - ], - outputs=[ - io.ControlNet.Output(), - ], - ) - - @classmethod - def execute(cls, control_net, type) -> io.NodeOutput: - control_net = control_net.copy() - type_number = UNION_CONTROLNET_TYPES.get(type, -1) - if type_number >= 0: - control_net.set_extra_arg("control_type", [type_number]) - else: - control_net.set_extra_arg("control_type", []) - - return io.NodeOutput(control_net) - - class ControlNetInpaintingAliMamaApply(ControlNetApplyAdvanced): @classmethod def define_schema(cls): diff --git a/comfy_extras/v3/nodes_cosmos.py b/comfy_extras/v3/nodes_cosmos.py index 9779e0ffe..a32c192e8 100644 --- a/comfy_extras/v3/nodes_cosmos.py +++ b/comfy_extras/v3/nodes_cosmos.py @@ -9,6 +9,29 @@ import nodes from comfy_api.latest import io +class EmptyCosmosLatentVideo(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="EmptyCosmosLatentVideo_V3", + category="latent/video", + inputs=[ + io.Int.Input("width", default=1280, min=16, max=nodes.MAX_RESOLUTION, step=16), + io.Int.Input("height", default=704, min=16, max=nodes.MAX_RESOLUTION, step=16), + io.Int.Input("length", default=121, min=1, max=nodes.MAX_RESOLUTION, step=8), + io.Int.Input("batch_size", default=1, min=1, max=4096), + ], + outputs=[io.Latent.Output()], + ) + + @classmethod + def execute(cls, width, height, length, batch_size) -> io.NodeOutput: + latent = torch.zeros( + [batch_size, 16, ((length - 1) // 8) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device() + ) + return io.NodeOutput({"samples": latent}) + + def vae_encode_with_padding(vae, image, width, height, length, padding=0): pixels = comfy.utils.common_upscale(image[..., :3].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1) pixel_len = min(pixels.shape[0], length) @@ -116,30 +139,7 @@ class CosmosPredict2ImageToVideoLatent(io.ComfyNode): return io.NodeOutput(out_latent) -class EmptyCosmosLatentVideo(io.ComfyNode): - @classmethod - def define_schema(cls) -> io.Schema: - return io.Schema( - node_id="EmptyCosmosLatentVideo_V3", - category="latent/video", - inputs=[ - io.Int.Input("width", default=1280, min=16, max=nodes.MAX_RESOLUTION, step=16), - io.Int.Input("height", default=704, min=16, max=nodes.MAX_RESOLUTION, step=16), - io.Int.Input("length", default=121, min=1, max=nodes.MAX_RESOLUTION, step=8), - io.Int.Input("batch_size", default=1, min=1, max=4096), - ], - outputs=[io.Latent.Output()], - ) - - @classmethod - def execute(cls, width, height, length, batch_size) -> io.NodeOutput: - latent = torch.zeros( - [batch_size, 16, ((length - 1) // 8) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device() - ) - return io.NodeOutput({"samples": latent}) - - -NODES_LIST = [ +NODES_LIST: list[type[io.ComfyNode]] = [ CosmosImageToVideoLatent, CosmosPredict2ImageToVideoLatent, EmptyCosmosLatentVideo,