400 lines
14 KiB
Python
400 lines
14 KiB
Python
# (c) TenserTensor <tenser.tensor@proton.me> || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
|
|
|
from typing import override
|
|
|
|
import torch
|
|
|
|
import latent_preview
|
|
from comfy import sample, model_management, samplers
|
|
from comfy_api.latest import ComfyExtension, io
|
|
from .nodes_context import Context
|
|
from .nodes_latent import SCALE_FACTORS, SCALE_METHODS, scale_latent
|
|
from .nodes_text_encoder import SingleCondCFGGuider
|
|
|
|
CATEGORY = "TenserTensor/Sampling"
|
|
|
|
|
|
def do_sample(**kwargs):
|
|
model = kwargs.get("model", None)
|
|
if model is None:
|
|
raise ValueError("Model is required for sampling")
|
|
|
|
latent = kwargs.get("latent", None)
|
|
if latent is None:
|
|
raise ValueError("Latent image is required for sampling")
|
|
|
|
latent_dict = latent.copy()
|
|
samples, add_noise, seed, steps, device = (
|
|
latent_dict["samples"],
|
|
kwargs.get("add_noise", True),
|
|
kwargs.get("seed", 0),
|
|
kwargs.get("steps", 0),
|
|
kwargs.get("device", "cpu")
|
|
)
|
|
latent_tensor = sample.fix_empty_latent_channels(model, samples, latent.get("downscale_ratio_spacial", None))
|
|
|
|
if add_noise:
|
|
batch_idx = latent.get("batch_index", None)
|
|
noise = sample.prepare_noise(samples, seed, batch_idx)
|
|
else:
|
|
noise = torch.zeros(samples.size(), dtype=samples.dtype, layout=samples.layout, device="cpu")
|
|
|
|
interim_buffer = {}
|
|
if kwargs.get("preview_latent"):
|
|
latent_preview.set_preview_method("latent2rgb")
|
|
callback = latent_preview.prepare_callback(model, steps - 1, interim_buffer)
|
|
else:
|
|
callback = None
|
|
|
|
args = {
|
|
"model": model,
|
|
"noise": noise,
|
|
"steps": steps,
|
|
"cfg": kwargs.get("cfg"),
|
|
"sampler_name": kwargs.get("sampler_name"),
|
|
"scheduler": kwargs.get("scheduler"),
|
|
"positive": kwargs.get("positive"),
|
|
"negative": kwargs.get("negative"),
|
|
"latent_image": latent_tensor,
|
|
"denoise": kwargs.get("denoise", 1.0),
|
|
"disable_noise": not add_noise,
|
|
"start_step": kwargs.get("start_step"),
|
|
"last_step": kwargs.get("last_step"),
|
|
"force_full_denoise": kwargs.get("full_denoise", True),
|
|
"noise_mask": latent.get("noise_mask"),
|
|
"callback": callback,
|
|
"seed": seed,
|
|
}
|
|
|
|
samples = sample.sample(**args)
|
|
sampled_latent = latent_dict.copy()
|
|
sampled_latent.pop("downscale_ratio_spacial", None)
|
|
sampled_latent["samples"] = samples
|
|
|
|
return sampled_latent
|
|
|
|
|
|
class TT_KSamplerNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_KSamplerNode",
|
|
display_name="TT KSampler",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.Conditioning.Input("positive"),
|
|
io.Conditioning.Input("negative"),
|
|
io.Latent.Input("latent"),
|
|
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
|
|
io.Int.Input("steps", default=30, min=1, max=10_000),
|
|
io.Float.Input("cfg", default=3.0, min=0.0, max=100.0, step=0.1),
|
|
io.Combo.Input("sampler_name", options=samplers.KSampler.SAMPLERS),
|
|
io.Combo.Input("scheduler", options=samplers.KSampler.SCHEDULERS),
|
|
io.Boolean.Input("preview_latent", default=False, label_on="Show Preview", label_off="Only Sample"),
|
|
],
|
|
outputs=[
|
|
io.Latent.Output("LATENT"),
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
latent = do_sample(**kwargs)
|
|
|
|
return io.NodeOutput(latent)
|
|
|
|
|
|
class TT_KSamplerAdvancedNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_KSamplerAdvancedNode",
|
|
display_name="TT KSampler (Advanced)",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.Conditioning.Input("positive"),
|
|
io.Conditioning.Input("negative"),
|
|
io.Latent.Input("latent"),
|
|
io.Boolean.Input("add_noise", default=False, label_on="Random Noise", label_off="Zero Noise"),
|
|
io.Boolean.Input("full_denoise", default=False, label_on="Complete", label_off="Partial"),
|
|
io.Float.Input("denoise", default=1.0, min=0.0, max=1.0, step=0.01),
|
|
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
|
|
io.Int.Input("steps", default=30, min=1, max=10_000),
|
|
io.Float.Input("cfg", default=3.0, min=0.0, max=100.0, step=0.1),
|
|
io.Int.Input("start_step", default=0, min=0, max=10_000),
|
|
io.Int.Input("last_step", default=10_000, min=0, max=10_000),
|
|
io.Combo.Input("sampler_name", options=samplers.KSampler.SAMPLERS),
|
|
io.Combo.Input("scheduler", options=samplers.KSampler.SCHEDULERS),
|
|
io.Boolean.Input("preview_latent", default=False, label_on="Show Preview", label_off="Only Sample"),
|
|
],
|
|
outputs=[
|
|
io.Latent.Output("LATENT"),
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
latent = do_sample(**kwargs)
|
|
|
|
return io.NodeOutput(latent)
|
|
|
|
|
|
class TT_KSamplerContextNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_KSamplerContextNode",
|
|
display_name="TT KSampler (Context)",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
Context.Input("context"),
|
|
io.Boolean.Input("preview_latent", default=False, label_on="Show Preview", label_off="Only Sample"),
|
|
],
|
|
outputs=[
|
|
Context.Output("CONTEXT"),
|
|
io.Latent.Output("LATENT"),
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, context) -> io.NodeOutput:
|
|
model = context.get("model")
|
|
if model is None:
|
|
raise ValueError("Model is required for sampling")
|
|
|
|
latent = context.get("latent")
|
|
if latent is None:
|
|
raise ValueError("Latent image is required for sampling")
|
|
|
|
config = context.get("workflow_config")
|
|
if config is None:
|
|
raise ValueError("Workflow Config is required for sampling")
|
|
|
|
args = {
|
|
"model": model,
|
|
"positive": context.get_attr("positive"),
|
|
"negative": context.get_attr("negative"),
|
|
"latent": latent,
|
|
"seed": config.get_attr("seed"),
|
|
"steps": config.get_attr("steps"),
|
|
"cfg": config.get_attr("cfg"),
|
|
"sampler_name": config.get_attr("sampler_name"),
|
|
"scheduler": config.get_attr("scheduler"),
|
|
}
|
|
latent = do_sample(**args)
|
|
context.set_attr("latent", latent)
|
|
|
|
return io.NodeOutput(context, latent)
|
|
|
|
|
|
class TT_KSamplerTwoStageNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_KSamplerTwoStageNode",
|
|
display_name="TT KSampler (Two Stage)",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.Conditioning.Input("positive"),
|
|
io.Conditioning.Input("negative"),
|
|
io.Latent.Input("latent"),
|
|
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
|
|
io.Float.Input("cfg", default=3.0, min=0.0, max=100.0, step=0.1),
|
|
io.Int.Input("draft_steps", default=25, min=1, max=10_000),
|
|
io.Int.Input("refiner_steps", default=25, min=1, max=10_000),
|
|
io.Combo.Input("draft_sampler_name", options=samplers.KSampler.SAMPLERS),
|
|
io.Combo.Input("draft_scheduler", options=samplers.KSampler.SCHEDULERS),
|
|
io.Combo.Input("refiner_sampler_name", options=samplers.KSampler.SAMPLERS),
|
|
io.Combo.Input("refiner_scheduler", options=samplers.KSampler.SCHEDULERS),
|
|
io.Float.Input("draft_denoise", default=1.0, min=0.0, max=1.0, step=0.01),
|
|
io.Float.Input("refiner_denoise", default=1.0, min=0.0, max=1.0, step=0.01),
|
|
io.Boolean.Input("preview_latent", default=False, label_on="Show Preview", label_off="Only Sample"),
|
|
],
|
|
outputs=[
|
|
io.Latent.Output("LATENT"),
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
draft_steps = kwargs["draft_steps"]
|
|
refiner_steps = kwargs["refiner_steps"]
|
|
|
|
args = kwargs.copy()
|
|
args.update({
|
|
"steps": draft_steps,
|
|
"start_step": 0,
|
|
"add_noise": True,
|
|
"sampler_name": kwargs["draft_sampler_name"],
|
|
"scheduler": kwargs["draft_scheduler"],
|
|
"full_denoise": False,
|
|
"denoise": kwargs["draft_denoise"],
|
|
})
|
|
latent = do_sample(**args)
|
|
|
|
args.update({
|
|
"latent": latent,
|
|
"steps": draft_steps + refiner_steps,
|
|
"start_step": draft_steps,
|
|
"add_noise": True,
|
|
"sampler_name": kwargs["refiner_sampler_name"],
|
|
"scheduler": kwargs["refiner_scheduler"],
|
|
"full_denoise": True,
|
|
"denoise": kwargs["refiner_denoise"],
|
|
})
|
|
latent = do_sample(**args)
|
|
|
|
return io.NodeOutput(latent)
|
|
|
|
|
|
def execute_guided_sampling(**kwargs):
|
|
guider, latent, sigmas, random_noise, sampler, preview_latent = (
|
|
kwargs.get("guider"),
|
|
kwargs.get("latent"),
|
|
kwargs.get("sigmas"),
|
|
kwargs.get("random_noise"),
|
|
kwargs.get("sampler"),
|
|
kwargs.get("preview_latent"),
|
|
)
|
|
|
|
if not isinstance(guider, SingleCondCFGGuider):
|
|
guider = SingleCondCFGGuider.from_cfg_guider(guider)
|
|
|
|
model_patcher = guider.model_patcher
|
|
latent_dict = latent.copy()
|
|
latent_tensor = sample.fix_empty_latent_channels(
|
|
model_patcher,
|
|
latent_dict["samples"],
|
|
latent.get("downscale_ratio_spacial", None)
|
|
)
|
|
latent_dict["samples"] = latent_tensor
|
|
noise_mask = latent.get("noise_mask", None)
|
|
interim_buffer = {}
|
|
if preview_latent:
|
|
latent_preview.set_preview_method("latent2rgb")
|
|
callback = latent_preview.prepare_callback(model_patcher, sigmas.shape[-1] - 1, interim_buffer)
|
|
else:
|
|
callback = None
|
|
# disable_progress_bar = not utils.PROGRESS_BAR_ENABLED
|
|
|
|
samples = guider.sample(
|
|
random_noise.generate_noise(latent),
|
|
latent_tensor,
|
|
sampler,
|
|
sigmas,
|
|
denoise_mask=noise_mask,
|
|
callback=callback,
|
|
disable_pbar=False,
|
|
seed=random_noise.seed
|
|
).to(model_management.intermediate_device())
|
|
|
|
sampled_latent = latent_dict.copy()
|
|
sampled_latent.pop("downscale_ratio_spacial", None)
|
|
sampled_latent["samples"] = samples
|
|
|
|
return sampled_latent
|
|
|
|
|
|
class TT_GuidedKSamplerNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_GuidedKSamplerNode",
|
|
display_name="TT Guided KSampler",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
io.Latent.Input("latent"),
|
|
io.Guider.Input("guider"),
|
|
io.Sigmas.Input("sigmas"),
|
|
io.Sampler.Input("sampler"),
|
|
io.Noise.Input("random_noise"),
|
|
io.Boolean.Input("preview_latent", default=False, label_on="Show Preview", label_off="Only Sample"),
|
|
],
|
|
outputs=[
|
|
io.Latent.Output("SAMPLES")
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
samples = execute_guided_sampling(**kwargs)
|
|
|
|
return io.NodeOutput(samples)
|
|
|
|
|
|
class TT_GuidedUpscaleKSamplerNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_GuidedUpscaleKSamplerNode",
|
|
display_name="TT Guided Upscale KSampler",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
io.Latent.Input("latent"),
|
|
io.Guider.Input("guider"),
|
|
io.Sigmas.Input("sigmas"),
|
|
io.Sampler.Input("sampler"),
|
|
io.Noise.Input("random_noise"),
|
|
io.Boolean.Input("scale_latent", default=False, label_on="Scale", label_off="Skip"),
|
|
io.Combo.Input("scale_factor", options=SCALE_FACTORS, default="1x", advanced=True),
|
|
io.Combo.Input("scale_method", options=SCALE_METHODS, default="bicubic", advanced=True),
|
|
io.Boolean.Input("preview_latent", default=False, label_on="Show Preview", label_off="Only Sample"),
|
|
],
|
|
outputs=[
|
|
io.Latent.Output("SAMPLES")
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
samples = execute_guided_sampling(**kwargs)
|
|
|
|
if kwargs.get("scale_latent"):
|
|
samples = scale_latent(**kwargs)
|
|
|
|
return io.NodeOutput(samples)
|
|
|
|
|
|
# ==============================================================================
|
|
# V3 entrypoint — registers context nodes with ComfyUI
|
|
# ==============================================================================
|
|
|
|
class SamplingNodesExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
TT_KSamplerNode,
|
|
TT_KSamplerAdvancedNode,
|
|
TT_KSamplerContextNode,
|
|
TT_KSamplerTwoStageNode,
|
|
TT_GuidedKSamplerNode,
|
|
TT_GuidedUpscaleKSamplerNode,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> SamplingNodesExtension:
|
|
return SamplingNodesExtension()
|
|
|
|
|
|
# ==============================================================================
|
|
# Re-exports for backward compatibility
|
|
# ==============================================================================
|
|
|
|
__all__ = [
|
|
"TT_KSamplerNode",
|
|
"TT_KSamplerAdvancedNode",
|
|
"TT_KSamplerContextNode",
|
|
"TT_KSamplerTwoStageNode",
|
|
"TT_GuidedKSamplerNode",
|
|
"TT_GuidedUpscaleKSamplerNode",
|
|
]
|