| """ |
| 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},) |
