156 lines
4.3 KiB
Python
156 lines
4.3 KiB
Python
"""
|
|
KSampler Control Nodes for ComfyUI
|
|
Nodes for managing KSampler parameters and tracking generation time.
|
|
"""
|
|
|
|
import time
|
|
from typing import Tuple, Dict, Any
|
|
|
|
import torch
|
|
import comfy.samplers
|
|
|
|
|
|
class KSamplerControl:
|
|
"""
|
|
Master Controller for KSampler.
|
|
Captures parameters and start time, formats them for text overlay.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"latent": ("LATENT",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"steps": ("INT", {"default": 9, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"STRING",
|
|
"LATENT",
|
|
"INT",
|
|
"INT",
|
|
"FLOAT",
|
|
comfy.samplers.KSampler.SAMPLERS,
|
|
comfy.samplers.KSampler.SCHEDULERS,
|
|
"FLOAT"
|
|
)
|
|
RETURN_NAMES = (
|
|
"text_overlay",
|
|
"latent",
|
|
"seed",
|
|
"steps",
|
|
"cfg",
|
|
"sampler_name",
|
|
"scheduler",
|
|
"start_time"
|
|
)
|
|
|
|
FUNCTION = "control"
|
|
CATEGORY = "sampling/control"
|
|
|
|
def control(
|
|
self,
|
|
latent: Dict[str, Any],
|
|
seed: int,
|
|
steps: int,
|
|
cfg: float,
|
|
sampler_name: str,
|
|
scheduler: str
|
|
) -> Tuple[str, Dict[str, Any], int, int, float, str, str, float]:
|
|
"""
|
|
Capture KSampler parameters and format them for display.
|
|
|
|
Args:
|
|
latent: The latent image dictionary.
|
|
seed: Random seed used for generation.
|
|
steps: Number of sampling steps.
|
|
cfg: Classifier Free Guidance scale.
|
|
sampler_name: Name of the sampler.
|
|
scheduler: Name of the scheduler.
|
|
|
|
Returns:
|
|
Tuple containing text overlay string, passed-through parameters, and start time.
|
|
"""
|
|
# Capture start time for generation tracking
|
|
start_time = time.time()
|
|
|
|
# Extract resolution from latent tensor
|
|
# Latent shape: [batch, channels, height, width]
|
|
latent_samples = latent["samples"]
|
|
latent_height = latent_samples.shape[2]
|
|
latent_width = latent_samples.shape[3]
|
|
|
|
# Calculate actual image resolution (VAE uses 8x downsampling)
|
|
width = latent_width * 8
|
|
height = latent_height * 8
|
|
|
|
# Format text overlay
|
|
text = (
|
|
f"Seed: {seed}\n"
|
|
f"Steps: {steps} | CFG: {cfg:.1f}\n"
|
|
f"{sampler_name} | {scheduler}\n"
|
|
f"Resolution: {width}x{height}"
|
|
)
|
|
|
|
return (text, latent, seed, steps, cfg, sampler_name, scheduler, start_time)
|
|
|
|
|
|
class GenerationTime:
|
|
"""
|
|
Appends generation time to text overlay.
|
|
Should be placed AFTER image generation is complete.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"text": ("STRING", {"forceInput": True}),
|
|
"start_time": ("FLOAT", {"forceInput": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "text")
|
|
FUNCTION = "append_time"
|
|
CATEGORY = "sampling/control"
|
|
|
|
def append_time(
|
|
self,
|
|
image: torch.Tensor,
|
|
text: str,
|
|
start_time: float
|
|
) -> Tuple[torch.Tensor, str]:
|
|
"""
|
|
Calculate and append generation time to text.
|
|
|
|
Args:
|
|
image: The generated image tensor.
|
|
text: The existing text overlay string.
|
|
start_time: The start time timestamp.
|
|
|
|
Returns:
|
|
Tuple containing the image and the updated text string.
|
|
"""
|
|
end_time = time.time()
|
|
duration = end_time - start_time
|
|
new_text = f"{text}\nTime: {duration:.2f}s"
|
|
return (image, new_text)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"KSamplerControl": KSamplerControl,
|
|
"GenerationTime": GenerationTime,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"KSamplerControl": "KSampler Bridge to Text (SEB)",
|
|
"GenerationTime": "Generation Time (SEB)",
|
|
}
|