commit 7a2fe139c4c008a07bbf6b2974f905e91e56c254 Author: Troy McCall Date: Sun Nov 10 19:07:25 2024 -0500 Create __init__.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..f7cbc56 --- /dev/null +++ b/__init__.py @@ -0,0 +1,45 @@ +import torch + +class ScaleToTargetMegapixels: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "megapixels": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100.0, "step": 0.1}) + } + } + + CATEGORY = "image" + RETURN_TYPES = ("FLOAT", "FLOAT", "FLOAT", "FLOAT") + RETURN_NAMES = ("upscale", "downscale", "upscale_if_need", "downscale_if_need", ) + OUTPUT_NODE = True + + FUNCTION = "scale" + + def scale(self, image, megapixels): + height, width = image.shape[1], image.shape[2] + current_megapixels = (height * width) / 1_000_000 + + if megapixels <= 0: + raise ValueError("Megapixels must be greater than 0") + + upscale = (megapixels / current_megapixels) ** 0.5 + downscale = (current_megapixels / megapixels) ** 0.5 + + upscale_if_need = upscale + downscale_if_need = downscale + + if current_megapixels < megapixels: + upscale_if_need = 1.0 + downscale_if_need = 1.0 + + + text = [f"upscale: {upscale} downscale: {downscale} upscale_if_need: {upscale_if_need} downscale_if_need: {downscale_if_need}"] + + print(text) + + return {"ui": {"text": text}, "result": (upscale, downscale, upscale_if_need, downscale_if_need)} + + +NODE_CLASS_MAPPINGS = {"ScaleToTargetMegapixels": ScaleToTargetMegapixels}