mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
250 lines
11 KiB
Python
250 lines
11 KiB
Python
from typing import Any, Literal, Optional, Union
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
|
|
from invokeai.app.invocations.fields import (
|
|
FieldDescriptions,
|
|
ImageField,
|
|
InputField,
|
|
MetadataField,
|
|
OutputField,
|
|
UIType,
|
|
)
|
|
from invokeai.app.invocations.model import ModelIdentifierField
|
|
from invokeai.app.services.shared.invocation_context import InvocationContext
|
|
from invokeai.app.util.controlnet_utils import CONTROLNET_MODE_VALUES, CONTROLNET_RESIZE_VALUES
|
|
from invokeai.version.invokeai_version import __version__
|
|
|
|
|
|
class MetadataItemField(BaseModel):
|
|
label: str = Field(description=FieldDescriptions.metadata_item_label)
|
|
value: Any = Field(description=FieldDescriptions.metadata_item_value)
|
|
|
|
|
|
class LoRAMetadataField(BaseModel):
|
|
"""LoRA Metadata Field"""
|
|
|
|
model: ModelIdentifierField = Field(description=FieldDescriptions.lora_model)
|
|
weight: float = Field(description=FieldDescriptions.lora_weight)
|
|
|
|
|
|
class IPAdapterMetadataField(BaseModel):
|
|
"""IP Adapter Field, minus the CLIP Vision Encoder model"""
|
|
|
|
image: ImageField = Field(description="The IP-Adapter image prompt.")
|
|
ip_adapter_model: ModelIdentifierField = Field(description="The IP-Adapter model.")
|
|
clip_vision_model: Literal["ViT-H", "ViT-G"] = Field(description="The CLIP Vision model")
|
|
method: Literal["full", "style", "composition"] = Field(description="Method to apply IP Weights with")
|
|
weight: Union[float, list[float]] = Field(description="The weight given to the IP-Adapter")
|
|
begin_step_percent: float = Field(description="When the IP-Adapter is first applied (% of total steps)")
|
|
end_step_percent: float = Field(description="When the IP-Adapter is last applied (% of total steps)")
|
|
|
|
|
|
class T2IAdapterMetadataField(BaseModel):
|
|
image: ImageField = Field(description="The control image.")
|
|
processed_image: Optional[ImageField] = Field(default=None, description="The control image, after processing.")
|
|
t2i_adapter_model: ModelIdentifierField = Field(description="The T2I-Adapter model to use.")
|
|
weight: Union[float, list[float]] = Field(default=1, description="The weight given to the T2I-Adapter")
|
|
begin_step_percent: float = Field(
|
|
default=0, ge=0, le=1, description="When the T2I-Adapter is first applied (% of total steps)"
|
|
)
|
|
end_step_percent: float = Field(
|
|
default=1, ge=0, le=1, description="When the T2I-Adapter is last applied (% of total steps)"
|
|
)
|
|
resize_mode: CONTROLNET_RESIZE_VALUES = Field(default="just_resize", description="The resize mode to use")
|
|
|
|
|
|
class ControlNetMetadataField(BaseModel):
|
|
image: ImageField = Field(description="The control image")
|
|
processed_image: Optional[ImageField] = Field(default=None, description="The control image, after processing.")
|
|
control_model: ModelIdentifierField = Field(description="The ControlNet model to use")
|
|
control_weight: Union[float, list[float]] = Field(default=1, description="The weight given to the ControlNet")
|
|
begin_step_percent: float = Field(
|
|
default=0, ge=0, le=1, description="When the ControlNet is first applied (% of total steps)"
|
|
)
|
|
end_step_percent: float = Field(
|
|
default=1, ge=0, le=1, description="When the ControlNet is last applied (% of total steps)"
|
|
)
|
|
control_mode: CONTROLNET_MODE_VALUES = Field(default="balanced", description="The control mode to use")
|
|
resize_mode: CONTROLNET_RESIZE_VALUES = Field(default="just_resize", description="The resize mode to use")
|
|
|
|
|
|
@invocation_output("metadata_item_output")
|
|
class MetadataItemOutput(BaseInvocationOutput):
|
|
"""Metadata Item Output"""
|
|
|
|
item: MetadataItemField = OutputField(description="Metadata Item")
|
|
|
|
|
|
@invocation("metadata_item", title="Metadata Item", tags=["metadata"], category="metadata", version="1.0.1")
|
|
class MetadataItemInvocation(BaseInvocation):
|
|
"""Used to create an arbitrary metadata item. Provide "label" and make a connection to "value" to store that data as the value."""
|
|
|
|
label: str = InputField(description=FieldDescriptions.metadata_item_label)
|
|
value: Any = InputField(description=FieldDescriptions.metadata_item_value, ui_type=UIType.Any)
|
|
|
|
def invoke(self, context: InvocationContext) -> MetadataItemOutput:
|
|
return MetadataItemOutput(item=MetadataItemField(label=self.label, value=self.value))
|
|
|
|
|
|
@invocation_output("metadata_output")
|
|
class MetadataOutput(BaseInvocationOutput):
|
|
metadata: MetadataField = OutputField(description="Metadata Dict")
|
|
|
|
|
|
@invocation("metadata", title="Metadata", tags=["metadata"], category="metadata", version="1.0.1")
|
|
class MetadataInvocation(BaseInvocation):
|
|
"""Takes a MetadataItem or collection of MetadataItems and outputs a MetadataDict."""
|
|
|
|
items: Union[list[MetadataItemField], MetadataItemField] = InputField(
|
|
description=FieldDescriptions.metadata_item_polymorphic
|
|
)
|
|
|
|
def invoke(self, context: InvocationContext) -> MetadataOutput:
|
|
if isinstance(self.items, MetadataItemField):
|
|
# single metadata item
|
|
data = {self.items.label: self.items.value}
|
|
else:
|
|
# collection of metadata items
|
|
data = {item.label: item.value for item in self.items}
|
|
|
|
# add app version
|
|
data.update({"app_version": __version__})
|
|
return MetadataOutput(metadata=MetadataField.model_validate(data))
|
|
|
|
|
|
@invocation("merge_metadata", title="Metadata Merge", tags=["metadata"], category="metadata", version="1.0.1")
|
|
class MergeMetadataInvocation(BaseInvocation):
|
|
"""Merged a collection of MetadataDict into a single MetadataDict."""
|
|
|
|
collection: list[MetadataField] = InputField(description=FieldDescriptions.metadata_collection)
|
|
|
|
def invoke(self, context: InvocationContext) -> MetadataOutput:
|
|
data = {}
|
|
for item in self.collection:
|
|
data.update(item.model_dump())
|
|
|
|
return MetadataOutput(metadata=MetadataField.model_validate(data))
|
|
|
|
|
|
GENERATION_MODES = Literal[
|
|
"txt2img", "img2img", "inpaint", "outpaint", "sdxl_txt2img", "sdxl_img2img", "sdxl_inpaint", "sdxl_outpaint"
|
|
]
|
|
|
|
|
|
@invocation("core_metadata", title="Core Metadata", tags=["metadata"], category="metadata", version="2.0.0")
|
|
class CoreMetadataInvocation(BaseInvocation):
|
|
"""Collects core generation metadata into a MetadataField"""
|
|
|
|
generation_mode: Optional[GENERATION_MODES] = InputField(
|
|
default=None,
|
|
description="The generation mode that output this image",
|
|
)
|
|
positive_prompt: Optional[str] = InputField(default=None, description="The positive prompt parameter")
|
|
negative_prompt: Optional[str] = InputField(default=None, description="The negative prompt parameter")
|
|
width: Optional[int] = InputField(default=None, description="The width parameter")
|
|
height: Optional[int] = InputField(default=None, description="The height parameter")
|
|
seed: Optional[int] = InputField(default=None, description="The seed used for noise generation")
|
|
rand_device: Optional[str] = InputField(default=None, description="The device used for random number generation")
|
|
cfg_scale: Optional[float] = InputField(default=None, description="The classifier-free guidance scale parameter")
|
|
cfg_rescale_multiplier: Optional[float] = InputField(
|
|
default=None, description=FieldDescriptions.cfg_rescale_multiplier
|
|
)
|
|
steps: Optional[int] = InputField(default=None, description="The number of steps used for inference")
|
|
scheduler: Optional[str] = InputField(default=None, description="The scheduler used for inference")
|
|
seamless_x: Optional[bool] = InputField(default=None, description="Whether seamless tiling was used on the X axis")
|
|
seamless_y: Optional[bool] = InputField(default=None, description="Whether seamless tiling was used on the Y axis")
|
|
clip_skip: Optional[int] = InputField(
|
|
default=None,
|
|
description="The number of skipped CLIP layers",
|
|
)
|
|
model: Optional[ModelIdentifierField] = InputField(default=None, description="The main model used for inference")
|
|
controlnets: Optional[list[ControlNetMetadataField]] = InputField(
|
|
default=None, description="The ControlNets used for inference"
|
|
)
|
|
ipAdapters: Optional[list[IPAdapterMetadataField]] = InputField(
|
|
default=None, description="The IP Adapters used for inference"
|
|
)
|
|
t2iAdapters: Optional[list[T2IAdapterMetadataField]] = InputField(
|
|
default=None, description="The IP Adapters used for inference"
|
|
)
|
|
loras: Optional[list[LoRAMetadataField]] = InputField(default=None, description="The LoRAs used for inference")
|
|
strength: Optional[float] = InputField(
|
|
default=None,
|
|
description="The strength used for latents-to-latents",
|
|
)
|
|
init_image: Optional[str] = InputField(
|
|
default=None,
|
|
description="The name of the initial image",
|
|
)
|
|
vae: Optional[ModelIdentifierField] = InputField(
|
|
default=None,
|
|
description="The VAE used for decoding, if the main model's default was not used",
|
|
)
|
|
|
|
# High resolution fix metadata.
|
|
hrf_enabled: Optional[bool] = InputField(
|
|
default=None,
|
|
description="Whether or not high resolution fix was enabled.",
|
|
)
|
|
# TODO: should this be stricter or do we just let the UI handle it?
|
|
hrf_method: Optional[str] = InputField(
|
|
default=None,
|
|
description="The high resolution fix upscale method.",
|
|
)
|
|
hrf_strength: Optional[float] = InputField(
|
|
default=None,
|
|
description="The high resolution fix img2img strength used in the upscale pass.",
|
|
)
|
|
|
|
# SDXL
|
|
positive_style_prompt: Optional[str] = InputField(
|
|
default=None,
|
|
description="The positive style prompt parameter",
|
|
)
|
|
negative_style_prompt: Optional[str] = InputField(
|
|
default=None,
|
|
description="The negative style prompt parameter",
|
|
)
|
|
|
|
# SDXL Refiner
|
|
refiner_model: Optional[ModelIdentifierField] = InputField(
|
|
default=None,
|
|
description="The SDXL Refiner model used",
|
|
)
|
|
refiner_cfg_scale: Optional[float] = InputField(
|
|
default=None,
|
|
description="The classifier-free guidance scale parameter used for the refiner",
|
|
)
|
|
refiner_steps: Optional[int] = InputField(
|
|
default=None,
|
|
description="The number of steps used for the refiner",
|
|
)
|
|
refiner_scheduler: Optional[str] = InputField(
|
|
default=None,
|
|
description="The scheduler used for the refiner",
|
|
)
|
|
refiner_positive_aesthetic_score: Optional[float] = InputField(
|
|
default=None,
|
|
description="The aesthetic score used for the refiner",
|
|
)
|
|
refiner_negative_aesthetic_score: Optional[float] = InputField(
|
|
default=None,
|
|
description="The aesthetic score used for the refiner",
|
|
)
|
|
refiner_start: Optional[float] = InputField(
|
|
default=None,
|
|
description="The start value used for refiner denoising",
|
|
)
|
|
|
|
def invoke(self, context: InvocationContext) -> MetadataOutput:
|
|
"""Collects and outputs a CoreMetadata object"""
|
|
|
|
as_dict = self.model_dump(exclude_none=True, exclude={"id", "type", "is_intermediate", "use_cache"})
|
|
as_dict["app_version"] = __version__
|
|
|
|
return MetadataOutput(metadata=MetadataField.model_validate(as_dict))
|
|
|
|
model_config = ConfigDict(extra="allow")
|