File size: 1,766 Bytes
695fbf0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
"""
Simple latent generator for ComfyUI.
The ``MagicSeedLatent`` class creates a random latent tensor of the specified size.
If ``mix_image`` is enabled, the input image is encoded with a VAE and mixed with noise.
"""
from __future__ import annotations
import torch
class MagicSeedLatent:
"""Generate a latent tensor with optional image mixing."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"width": ("INT", {"default": 512, "min": 8, "max": 4096, "step": 8}),
"height": ("INT", {"default": 512, "min": 8, "max": 4096, "step": 8}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
"sigma": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}),
"bias": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.1}),
"mix_image": ("BOOLEAN", {"default": False}),
},
"optional": {
"vae": ("VAE", {}),
"image": ("IMAGE", {}),
},
}
RETURN_TYPES = ("LATENT",)
RETURN_NAMES = ("LATENT",)
FUNCTION = "generate"
CATEGORY = "MagicNodes"
def generate(
self,
width: int,
height: int,
batch_size: int,
sigma: float,
bias: float,
mix_image: bool = False,
vae=None,
image=None,
):
"""Generate a random latent tensor and optionally mix it with an image."""
lat = torch.randn(batch_size, 4, height // 8, width // 8) * sigma + bias
if mix_image and vae is not None and image is not None:
encoded = vae.encode(image[:, :, :, :3])
lat = encoded + lat
return ({"samples": lat},) |