cddb07a176
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
import torch
|
|
from pydantic import field_validator
|
|
|
|
from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
|
|
from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
|
|
from invokeai.app.invocations.fields import FieldDescriptions, InputField, LatentsField, OutputField
|
|
from invokeai.app.invocations.latent_noise import (
|
|
LatentNoiseType,
|
|
generate_noise_tensor,
|
|
)
|
|
from invokeai.app.services.shared.invocation_context import InvocationContext
|
|
from invokeai.app.util.misc import SEED_MAX
|
|
from invokeai.backend.util.devices import TorchDevice
|
|
|
|
|
|
@invocation_output("noise_output")
|
|
class NoiseOutput(BaseInvocationOutput):
|
|
"""Invocation noise output"""
|
|
|
|
noise: LatentsField = OutputField(description=FieldDescriptions.noise)
|
|
width: int = OutputField(description=FieldDescriptions.width)
|
|
height: int = OutputField(description=FieldDescriptions.height)
|
|
|
|
@classmethod
|
|
def build(cls, latents_name: str, latents: torch.Tensor, seed: int) -> "NoiseOutput":
|
|
return cls(
|
|
noise=LatentsField(latents_name=latents_name, seed=seed),
|
|
width=latents.shape[-1] * LATENT_SCALE_FACTOR,
|
|
height=latents.shape[-2] * LATENT_SCALE_FACTOR,
|
|
)
|
|
|
|
|
|
@invocation(
|
|
"noise",
|
|
title="Create Latent Noise",
|
|
tags=["latents", "noise"],
|
|
category="latents",
|
|
version="1.1.0",
|
|
)
|
|
class NoiseInvocation(BaseInvocation):
|
|
"""Generates latent noise for supported denoiser architectures."""
|
|
|
|
noise_type: LatentNoiseType = InputField(default="SD", description="Architecture-specific noise type.")
|
|
|
|
seed: int = InputField(
|
|
default=0,
|
|
ge=0,
|
|
le=SEED_MAX,
|
|
description=FieldDescriptions.seed,
|
|
)
|
|
width: int = InputField(
|
|
default=512,
|
|
multiple_of=LATENT_SCALE_FACTOR,
|
|
gt=0,
|
|
description=FieldDescriptions.width,
|
|
)
|
|
height: int = InputField(
|
|
default=512,
|
|
multiple_of=LATENT_SCALE_FACTOR,
|
|
gt=0,
|
|
description=FieldDescriptions.height,
|
|
)
|
|
use_cpu: bool = InputField(
|
|
default=True,
|
|
description="Use CPU for noise generation (for reproducible results across platforms)",
|
|
)
|
|
|
|
@field_validator("seed", mode="before")
|
|
def modulo_seed(cls, v):
|
|
"""Return the seed modulo (SEED_MAX + 1) to ensure it is within the valid range."""
|
|
return v % (SEED_MAX + 1)
|
|
|
|
def invoke(self, context: InvocationContext) -> NoiseOutput:
|
|
noise = generate_noise_tensor(
|
|
noise_type=self.noise_type,
|
|
width=self.width,
|
|
height=self.height,
|
|
device=TorchDevice.choose_torch_device(),
|
|
seed=self.seed,
|
|
dtype=TorchDevice.choose_torch_dtype(),
|
|
use_cpu=self.use_cpu,
|
|
)
|
|
name = context.tensors.save(tensor=noise)
|
|
return NoiseOutput.build(latents_name=name, latents=noise, seed=self.seed)
|