128 lines
3.8 KiB
Python
128 lines
3.8 KiB
Python
from .core_sampling import (
|
|
SampleChunkConfig,
|
|
decode_latent_batch_chunked,
|
|
encode_image_batch_chunked,
|
|
sample_latent_batch_chunked,
|
|
)
|
|
|
|
|
|
class VAEEncodeBatchChunked:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pixels": ("IMAGE",),
|
|
"vae": ("VAE",),
|
|
"chunk_size": ("INT", {"default": 8, "min": 1, "max": 4096}),
|
|
"clear_cache_between_chunks": ("BOOLEAN", {"default": False, "advanced": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("latent",)
|
|
FUNCTION = "encode"
|
|
CATEGORY = "latent"
|
|
|
|
def encode(self, pixels, vae, chunk_size, clear_cache_between_chunks=False):
|
|
return (
|
|
encode_image_batch_chunked(
|
|
pixels,
|
|
vae,
|
|
chunk_size,
|
|
clear_cache_between_chunks=clear_cache_between_chunks,
|
|
),
|
|
)
|
|
|
|
|
|
class VAEDecodeBatchChunked:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"samples": ("LATENT",),
|
|
"vae": ("VAE",),
|
|
"chunk_size": ("INT", {"default": 8, "min": 1, "max": 4096}),
|
|
"clear_cache_between_chunks": ("BOOLEAN", {"default": False, "advanced": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "decode"
|
|
CATEGORY = "latent"
|
|
|
|
def decode(self, samples, vae, chunk_size, clear_cache_between_chunks=False):
|
|
return (
|
|
decode_latent_batch_chunked(
|
|
samples,
|
|
vae,
|
|
chunk_size,
|
|
clear_cache_between_chunks=clear_cache_between_chunks,
|
|
),
|
|
)
|
|
|
|
|
|
class SamplerCustomAdvancedChunked:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"noise": ("NOISE",),
|
|
"guider": ("GUIDER",),
|
|
"sampler": ("SAMPLER",),
|
|
"sigmas": ("SIGMAS",),
|
|
"latent_image": ("LATENT",),
|
|
"chunk_size": ("INT", {"default": 2, "min": 1, "max": 4096}),
|
|
"auto_reduce_on_oom": ("BOOLEAN", {"default": True}),
|
|
"min_chunk_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
|
"clear_cache_on_retry": ("BOOLEAN", {"default": True, "advanced": True}),
|
|
"clear_cache_between_chunks": ("BOOLEAN", {"default": False, "advanced": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "LATENT")
|
|
RETURN_NAMES = ("output", "denoised_output")
|
|
FUNCTION = "sample"
|
|
CATEGORY = "sampling/custom_sampling"
|
|
|
|
def sample(
|
|
self,
|
|
noise,
|
|
guider,
|
|
sampler,
|
|
sigmas,
|
|
latent_image,
|
|
chunk_size,
|
|
auto_reduce_on_oom=True,
|
|
min_chunk_size=1,
|
|
clear_cache_on_retry=True,
|
|
clear_cache_between_chunks=False,
|
|
):
|
|
return sample_latent_batch_chunked(
|
|
noise,
|
|
guider,
|
|
sampler,
|
|
sigmas,
|
|
latent_image,
|
|
config=SampleChunkConfig(
|
|
chunk_size=chunk_size,
|
|
auto_reduce_on_oom=auto_reduce_on_oom,
|
|
min_chunk_size=min_chunk_size,
|
|
clear_cache_on_retry=clear_cache_on_retry,
|
|
clear_cache_between_chunks=clear_cache_between_chunks,
|
|
),
|
|
)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"VAEEncodeBatchChunked": VAEEncodeBatchChunked,
|
|
"VAEDecodeBatchChunked": VAEDecodeBatchChunked,
|
|
"SamplerCustomAdvancedChunked": SamplerCustomAdvancedChunked,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"VAEEncodeBatchChunked": "VAE Encode Batch Chunked",
|
|
"VAEDecodeBatchChunked": "VAE Decode Batch Chunked",
|
|
"SamplerCustomAdvancedChunked": "SamplerCustomAdvanced Chunked",
|
|
}
|