116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
import torch
|
|
import random
|
|
from server import PromptServer
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# 0. GLOBAL PRINT TOGGLE
|
|
# ----------------------------------------------------------------------------
|
|
ENABLE_PRINT = False # <-- Set to True to enable all prints
|
|
|
|
def debug_print(*args, **kwargs):
|
|
if ENABLE_PRINT:
|
|
print(*args, **kwargs)
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# 1. THE MATH
|
|
# ----------------------------------------------------------------------------
|
|
def get_custom_noise_function(latent_image, seed):
|
|
debug_print(f"[PYTHON-MATH] Generating noise with seed: {seed}")
|
|
generator = torch.manual_seed(seed)
|
|
noise = torch.randn(
|
|
latent_image.size(),
|
|
dtype=latent_image.dtype,
|
|
layout=latent_image.layout,
|
|
generator=generator,
|
|
device="cpu"
|
|
)
|
|
return noise.to(latent_image.device)
|
|
|
|
class CustomNoiseObject:
|
|
def __init__(self, seed):
|
|
self.seed = seed
|
|
|
|
def generate_noise(self, input_latent):
|
|
return get_custom_noise_function(input_latent["samples"], self.seed)
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# 2. THE NODE
|
|
# ----------------------------------------------------------------------------
|
|
class RandomNoiseCustom:
|
|
last_seeds = {} # Store per-node
|
|
MAX_SAFE_INTEGER = 9007199254740991 # JavaScript's MAX_SAFE_INTEGER (2^53 - 1)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"seed_value": ("INT", {
|
|
"default": 42,
|
|
"min": 0,
|
|
"max": s.MAX_SAFE_INTEGER
|
|
}),
|
|
"mode": (["Fixed", "Increment", "Decrement", "Randomize"], {
|
|
"default": "Randomize"
|
|
}),
|
|
},
|
|
"hidden": {"unique_id": "UNIQUE_ID"}
|
|
}
|
|
|
|
RETURN_TYPES = ("NOISE",)
|
|
FUNCTION = "get_noise"
|
|
CATEGORY = "sampling/custom_sampling/noise"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, seed_value, mode, unique_id):
|
|
if mode == "Fixed":
|
|
return seed_value
|
|
return float("NaN")
|
|
|
|
def get_noise(self, seed_value, mode, unique_id):
|
|
# Use the current seed_value from the widget as the starting point
|
|
current_seed = seed_value
|
|
|
|
# Get the last seed for increment/decrement operations
|
|
last_seed = self.last_seeds.get(unique_id, current_seed)
|
|
|
|
if mode == "Fixed":
|
|
seed_to_use = current_seed
|
|
elif mode == "Randomize":
|
|
seed_to_use = random.randint(0, self.MAX_SAFE_INTEGER)
|
|
elif mode == "Increment":
|
|
# Increment but cap at MAX_SAFE_INTEGER
|
|
seed_to_use = min(current_seed + 1, self.MAX_SAFE_INTEGER)
|
|
elif mode == "Decrement":
|
|
# Decrement but don't go below 0
|
|
seed_to_use = max(current_seed - 1, 0)
|
|
else:
|
|
seed_to_use = current_seed
|
|
|
|
# Store this seed for next time
|
|
self.last_seeds[unique_id] = seed_to_use
|
|
|
|
debug_print(f"[PYTHON-NODE] Current Widget: {current_seed}, Calculated Seed: {seed_to_use}")
|
|
|
|
if mode != "Fixed":
|
|
PromptServer.instance.send_sync(
|
|
"comfy_server_noise_update",
|
|
{
|
|
"node_id": unique_id,
|
|
"seed": str(seed_to_use)
|
|
}
|
|
)
|
|
|
|
return (CustomNoiseObject(seed_to_use),)
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# NODE REGISTRATION
|
|
# ----------------------------------------------------------------------------
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RandomNoiseCustom": RandomNoiseCustom
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RandomNoiseCustom": "RandomNoiseCustom"
|
|
}
|
|
|
|
WEB_DIRECTORY = "./js" |