| from copy import copy
|
|
|
| from .nodes_registry import comfy_node
|
|
|
|
|
| @comfy_node(name="Set VAE Decoder Noise")
|
| class DecoderNoise:
|
| @classmethod
|
| def INPUT_TYPES(cls):
|
| return {
|
| "required": {
|
| "vae": ("VAE",),
|
| "timestep": (
|
| "FLOAT",
|
| {
|
| "default": 0.05,
|
| "min": 0.0,
|
| "max": 1.0,
|
| "step": 0.001,
|
| "tooltip": "The timestep used for decoding the noise.",
|
| },
|
| ),
|
| "scale": (
|
| "FLOAT",
|
| {
|
| "default": 0.025,
|
| "min": 0.0,
|
| "max": 1.0,
|
| "step": 0.001,
|
| "tooltip": "The scale of the noise added to the decoder.",
|
| },
|
| ),
|
| "seed": (
|
| "INT",
|
| {
|
| "default": 42,
|
| "min": 0,
|
| "max": 0xFFFFFFFFFFFFFFFF,
|
| "tooltip": "The random seed used for creating the noise.",
|
| },
|
| ),
|
| }
|
| }
|
|
|
| FUNCTION = "add_noise"
|
| RETURN_TYPES = ("VAE",)
|
| CATEGORY = "lightricks/LTXV"
|
|
|
| def add_noise(self, vae, timestep, scale, seed):
|
| result = copy(vae)
|
| if hasattr(result, "first_stage_model"):
|
| result.first_stage_model.decode_timestep = timestep
|
| result.first_stage_model.decode_noise_scale = scale
|
| result._decode_timestep = timestep
|
| result.decode_noise_scale = scale
|
| result.seed = seed
|
| return (result,)
|
|
|
