|
1 | 1 | # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) |
2 | 2 |
|
| 3 | +import math |
3 | 4 | from contextlib import ExitStack |
4 | 5 | from functools import singledispatchmethod |
5 | 6 | from typing import List, Literal, Optional, Union |
@@ -1228,3 +1229,54 @@ def invoke(self, context: InvocationContext) -> LatentsOutput: |
1228 | 1229 | context.services.latents.save(name, cropped_latents) |
1229 | 1230 |
|
1230 | 1231 | return build_latents_output(latents_name=name, latents=cropped_latents) |
| 1232 | + |
| 1233 | + |
| 1234 | +@invocation_output("ideal_size_output") |
| 1235 | +class IdealSizeOutput(BaseInvocationOutput): |
| 1236 | + """Base class for invocations that output an image""" |
| 1237 | + |
| 1238 | + width: int = OutputField(description="The ideal width of the image in pixels") |
| 1239 | + height: int = OutputField(description="The ideal height of the image in pixels") |
| 1240 | + |
| 1241 | + |
| 1242 | +@invocation( |
| 1243 | + "ideal_size", |
| 1244 | + title="Ideal Size", |
| 1245 | + tags=["latents", "math", "ideal_size"], |
| 1246 | + version="1.0.2", |
| 1247 | +) |
| 1248 | +class IdealSizeInvocation(BaseInvocation): |
| 1249 | + """Calculates the ideal size for generation to avoid duplication""" |
| 1250 | + |
| 1251 | + width: int = InputField(default=1024, description="Target width") |
| 1252 | + height: int = InputField(default=576, description="Target height") |
| 1253 | + unet: UNetField = InputField(default=None, description="UNet submodel") |
| 1254 | + multiplier: float = InputField(default=1.0, description="Dimensional multiplier") |
| 1255 | + |
| 1256 | + def trim_to_multiple_of(self, *args, multiple_of=LATENT_SCALE_FACTOR): |
| 1257 | + return tuple((x - x % multiple_of) for x in args) |
| 1258 | + |
| 1259 | + def invoke(self, context: InvocationContext) -> IdealSizeOutput: |
| 1260 | + aspect = self.width / self.height |
| 1261 | + dimension = 512 |
| 1262 | + if self.unet.unet.base_model == BaseModelType.StableDiffusion2: |
| 1263 | + dimension = 768 |
| 1264 | + elif self.unet.unet.base_model == BaseModelType.StableDiffusionXL: |
| 1265 | + dimension = 1024 |
| 1266 | + dimension = dimension * self.multiplier |
| 1267 | + min_dimension = math.floor(dimension * 0.5) |
| 1268 | + model_area = dimension * dimension # hardcoded for now since all models are trained on square images |
| 1269 | + |
| 1270 | + if aspect > 1.0: |
| 1271 | + init_height = max(min_dimension, math.sqrt(model_area / aspect)) |
| 1272 | + init_width = init_height * aspect |
| 1273 | + else: |
| 1274 | + init_width = max(min_dimension, math.sqrt(model_area * aspect)) |
| 1275 | + init_height = init_width / aspect |
| 1276 | + |
| 1277 | + scaled_width, scaled_height = self.trim_to_multiple_of( |
| 1278 | + math.floor(init_width), |
| 1279 | + math.floor(init_height), |
| 1280 | + ) |
| 1281 | + |
| 1282 | + return IdealSizeOutput(width=scaled_width, height=scaled_height) |
0 commit comments