1 line
3.1 KiB
Python
1 line
3.1 KiB
Python
import torch
|
|
import nodes # Èìïîðòèðóåì ñòàíäàðòíûå íîäû äëÿ äîñòóïà ê common_ksampler
|
|
import comfy.sample
|
|
|
|
class BatchSeedSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent_image": ("LATENT",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"control_after_generate": (["fixed", "increment", "decrement", "randomize"],),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "sample_batch"
|
|
CATEGORY = "sampling"
|
|
|
|
def sample_batch(self, model, positive, negative, latent_image, seed, control_after_generate, steps, cfg, sampler_name, scheduler, denoise):
|
|
samples = latent_image["samples"]
|
|
batch_size = samples.shape[0]
|
|
|
|
print(f"\n[BatchSampler] Íà÷èíàþ îáðàáîòêó áàò÷à. Ðàçìåð: {batch_size}")
|
|
|
|
output_samples = []
|
|
|
|
for i in range(batch_size):
|
|
# Ðàññ÷èòûâàåì ñèä äëÿ òåêóùåãî ýëåìåíòà
|
|
current_seed = seed
|
|
if control_after_generate == "increment":
|
|
current_seed = seed + i
|
|
elif control_after_generate == "decrement":
|
|
current_seed = seed - i
|
|
elif control_after_generate == "randomize":
|
|
# Ãåíåðèðóåì ïñåâäîðàíäîì íà îñíîâå áàçîâîãî ñèäà è èíäåêñà
|
|
current_seed = hash(f"{seed}_{i}") % 0xffffffffffffffff
|
|
|
|
print(f"[BatchSampler] Îáðàáîòêà ýëåìåíòà {i+1}/{batch_size} ñ Seed: {current_seed}")
|
|
|
|
# Âûäåëÿåì i-é ëàòåíò èç áàò÷à (ñîõðàíÿÿ ðàçìåðíîñòü [1, C, H, W])
|
|
single_latent = {"samples": samples[i:i+1]}
|
|
|
|
# Èñïîëüçóåì ñèñòåìíóþ ôóíêöèþ ñýìïëèðîâàíèÿ ComfyUI
|
|
# Îíà âîçâðàùàåò êîðòåæ, ãäå ïåðâûé ýëåìåíò - ýòî ñëîâàðü ñ ëàòåíòîì
|
|
try:
|
|
result = nodes.common_ksampler(
|
|
model,
|
|
current_seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
single_latent,
|
|
denoise=denoise
|
|
)
|
|
output_samples.append(result[0]["samples"])
|
|
except Exception as e:
|
|
print(f"[BatchSampler] ÎØÈÁÊÀ íà ýëåìåíòå {i}: {str(e)}")
|
|
raise e
|
|
|
|
# Îáúåäèíÿåì âñå ðåçóëüòàòû îáðàòíî â îäèí òåíçîð áàò÷à
|
|
final_samples = torch.cat(output_samples, dim=0)
|
|
|
|
print(f"[BatchSampler] Çàâåðøåíî. Èòîãîâûé òåíçîð: {final_samples.shape}\n")
|
|
|
|
return ({"samples": final_samples},) |