mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
138 lines
4.4 KiB
Python
138 lines
4.4 KiB
Python
# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) and the InvokeAI Team
|
|
|
|
from typing import Literal
|
|
|
|
import numpy as np
|
|
from pydantic import Field, validator
|
|
|
|
from invokeai.app.models.image import ImageField
|
|
from invokeai.app.util.misc import SEED_MAX, get_random_seed
|
|
|
|
from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationConfig, InvocationContext, UIConfig
|
|
|
|
|
|
class IntCollectionOutput(BaseInvocationOutput):
|
|
"""A collection of integers"""
|
|
|
|
type: Literal["int_collection"] = "int_collection"
|
|
|
|
# Outputs
|
|
collection: list[int] = Field(default=[], description="The int collection")
|
|
|
|
|
|
class FloatCollectionOutput(BaseInvocationOutput):
|
|
"""A collection of floats"""
|
|
|
|
type: Literal["float_collection"] = "float_collection"
|
|
|
|
# Outputs
|
|
collection: list[float] = Field(default=[], description="The float collection")
|
|
|
|
|
|
class ImageCollectionOutput(BaseInvocationOutput):
|
|
"""A collection of images"""
|
|
|
|
type: Literal["image_collection"] = "image_collection"
|
|
|
|
# Outputs
|
|
collection: list[ImageField] = Field(default=[], description="The output images")
|
|
|
|
class Config:
|
|
schema_extra = {"required": ["type", "collection"]}
|
|
|
|
|
|
class RangeInvocation(BaseInvocation):
|
|
"""Creates a range of numbers from start to stop with step"""
|
|
|
|
type: Literal["range"] = "range"
|
|
|
|
# Inputs
|
|
start: int = Field(default=0, description="The start of the range")
|
|
stop: int = Field(default=10, description="The stop of the range")
|
|
step: int = Field(default=1, description="The step of the range")
|
|
|
|
class Config(InvocationConfig):
|
|
schema_extra = {
|
|
"ui": {"title": "Range", "tags": ["range", "integer", "collection"]},
|
|
}
|
|
|
|
@validator("stop")
|
|
def stop_gt_start(cls, v, values):
|
|
if "start" in values and v <= values["start"]:
|
|
raise ValueError("stop must be greater than start")
|
|
return v
|
|
|
|
def invoke(self, context: InvocationContext) -> IntCollectionOutput:
|
|
return IntCollectionOutput(collection=list(range(self.start, self.stop, self.step)))
|
|
|
|
|
|
class RangeOfSizeInvocation(BaseInvocation):
|
|
"""Creates a range from start to start + size with step"""
|
|
|
|
type: Literal["range_of_size"] = "range_of_size"
|
|
|
|
# Inputs
|
|
start: int = Field(default=0, description="The start of the range")
|
|
size: int = Field(default=1, description="The number of values")
|
|
step: int = Field(default=1, description="The step of the range")
|
|
|
|
class Config(InvocationConfig):
|
|
schema_extra = {
|
|
"ui": {"title": "Sized Range", "tags": ["range", "integer", "size", "collection"]},
|
|
}
|
|
|
|
def invoke(self, context: InvocationContext) -> IntCollectionOutput:
|
|
return IntCollectionOutput(collection=list(range(self.start, self.start + self.size, self.step)))
|
|
|
|
|
|
class RandomRangeInvocation(BaseInvocation):
|
|
"""Creates a collection of random numbers"""
|
|
|
|
type: Literal["random_range"] = "random_range"
|
|
|
|
# Inputs
|
|
low: int = Field(default=0, description="The inclusive low value")
|
|
high: int = Field(default=np.iinfo(np.int32).max, description="The exclusive high value")
|
|
size: int = Field(default=1, description="The number of values to generate")
|
|
seed: int = Field(
|
|
ge=0,
|
|
le=SEED_MAX,
|
|
description="The seed for the RNG (omit for random)",
|
|
default_factory=get_random_seed,
|
|
)
|
|
|
|
class Config(InvocationConfig):
|
|
schema_extra = {
|
|
"ui": {"title": "Random Range", "tags": ["range", "integer", "random", "collection"]},
|
|
}
|
|
|
|
def invoke(self, context: InvocationContext) -> IntCollectionOutput:
|
|
rng = np.random.default_rng(self.seed)
|
|
return IntCollectionOutput(collection=list(rng.integers(low=self.low, high=self.high, size=self.size)))
|
|
|
|
|
|
class ImageCollectionInvocation(BaseInvocation):
|
|
"""Load a collection of images and provide it as output."""
|
|
|
|
# fmt: off
|
|
type: Literal["image_collection"] = "image_collection"
|
|
|
|
# Inputs
|
|
images: list[ImageField] = Field(
|
|
default=[], description="The image collection to load"
|
|
)
|
|
# fmt: on
|
|
|
|
def invoke(self, context: InvocationContext) -> ImageCollectionOutput:
|
|
return ImageCollectionOutput(collection=self.images)
|
|
|
|
class Config(InvocationConfig):
|
|
schema_extra = {
|
|
"ui": {
|
|
"type_hints": {
|
|
"title": "Image Collection",
|
|
"images": "image_collection",
|
|
}
|
|
},
|
|
}
|