From 6ade5df25cffa26285586b0295b343a294449c53 Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Wed, 5 Jun 2024 10:30:29 -0400 Subject: [PATCH] Move IdealSizeInvocation to its own file. No functional changes. --- invokeai/app/invocations/ideal_size.py | 65 ++++++++++++++++++++++++++ invokeai/app/invocations/latent.py | 56 ---------------------- 2 files changed, 65 insertions(+), 56 deletions(-) create mode 100644 invokeai/app/invocations/ideal_size.py diff --git a/invokeai/app/invocations/ideal_size.py b/invokeai/app/invocations/ideal_size.py new file mode 100644 index 0000000000..120f8c1ba0 --- /dev/null +++ b/invokeai/app/invocations/ideal_size.py @@ -0,0 +1,65 @@ +import math +from typing import Tuple + +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, OutputField +from invokeai.app.invocations.model import UNetField +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.config import BaseModelType + + +@invocation_output("ideal_size_output") +class IdealSizeOutput(BaseInvocationOutput): + """Base class for invocations that output an image""" + + width: int = OutputField(description="The ideal width of the image (in pixels)") + height: int = OutputField(description="The ideal height of the image (in pixels)") + + +@invocation( + "ideal_size", + title="Ideal Size", + tags=["latents", "math", "ideal_size"], + version="1.0.3", +) +class IdealSizeInvocation(BaseInvocation): + """Calculates the ideal size for generation to avoid duplication""" + + width: int = InputField(default=1024, description="Final image width") + height: int = InputField(default=576, description="Final image height") + unet: UNetField = InputField(default=None, description=FieldDescriptions.unet) + multiplier: float = InputField( + default=1.0, + description="Amount to multiply the model's dimensions by when calculating the ideal size (may result in " + "initial generation artifacts if too large)", + ) + + def trim_to_multiple_of(self, *args: int, multiple_of: int = LATENT_SCALE_FACTOR) -> Tuple[int, ...]: + return tuple((x - x % multiple_of) for x in args) + + def invoke(self, context: InvocationContext) -> IdealSizeOutput: + unet_config = context.models.get_config(self.unet.unet.key) + aspect = self.width / self.height + dimension: float = 512 + if unet_config.base == BaseModelType.StableDiffusion2: + dimension = 768 + elif unet_config.base == BaseModelType.StableDiffusionXL: + dimension = 1024 + dimension = dimension * self.multiplier + min_dimension = math.floor(dimension * 0.5) + model_area = dimension * dimension # hardcoded for now since all models are trained on square images + + if aspect > 1.0: + init_height = max(min_dimension, math.sqrt(model_area / aspect)) + init_width = init_height * aspect + else: + init_width = max(min_dimension, math.sqrt(model_area * aspect)) + init_height = init_width / aspect + + scaled_width, scaled_height = self.trim_to_multiple_of( + math.floor(init_width), + math.floor(init_height), + ) + + return IdealSizeOutput(width=scaled_width, height=scaled_height) diff --git a/invokeai/app/invocations/latent.py b/invokeai/app/invocations/latent.py index ede2443307..0fceec9c1b 100644 --- a/invokeai/app/invocations/latent.py +++ b/invokeai/app/invocations/latent.py @@ -1,6 +1,5 @@ # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) import inspect -import math from contextlib import ExitStack from functools import singledispatchmethod from typing import Any, Dict, Iterator, List, Literal, Optional, Tuple, Union @@ -1444,58 +1443,3 @@ class CropLatentsCoreInvocation(BaseInvocation): name = context.tensors.save(tensor=cropped_latents) return LatentsOutput.build(latents_name=name, latents=cropped_latents) - - -@invocation_output("ideal_size_output") -class IdealSizeOutput(BaseInvocationOutput): - """Base class for invocations that output an image""" - - width: int = OutputField(description="The ideal width of the image (in pixels)") - height: int = OutputField(description="The ideal height of the image (in pixels)") - - -@invocation( - "ideal_size", - title="Ideal Size", - tags=["latents", "math", "ideal_size"], - version="1.0.3", -) -class IdealSizeInvocation(BaseInvocation): - """Calculates the ideal size for generation to avoid duplication""" - - width: int = InputField(default=1024, description="Final image width") - height: int = InputField(default=576, description="Final image height") - unet: UNetField = InputField(default=None, description=FieldDescriptions.unet) - multiplier: float = InputField( - default=1.0, - description="Amount to multiply the model's dimensions by when calculating the ideal size (may result in initial generation artifacts if too large)", - ) - - def trim_to_multiple_of(self, *args: int, multiple_of: int = LATENT_SCALE_FACTOR) -> Tuple[int, ...]: - return tuple((x - x % multiple_of) for x in args) - - def invoke(self, context: InvocationContext) -> IdealSizeOutput: - unet_config = context.models.get_config(self.unet.unet.key) - aspect = self.width / self.height - dimension: float = 512 - if unet_config.base == BaseModelType.StableDiffusion2: - dimension = 768 - elif unet_config.base == BaseModelType.StableDiffusionXL: - dimension = 1024 - dimension = dimension * self.multiplier - min_dimension = math.floor(dimension * 0.5) - model_area = dimension * dimension # hardcoded for now since all models are trained on square images - - if aspect > 1.0: - init_height = max(min_dimension, math.sqrt(model_area / aspect)) - init_width = init_height * aspect - else: - init_width = max(min_dimension, math.sqrt(model_area * aspect)) - init_height = init_width / aspect - - scaled_width, scaled_height = self.trim_to_multiple_of( - math.floor(init_width), - math.floor(init_height), - ) - - return IdealSizeOutput(width=scaled_width, height=scaled_height)