Create __init__.py

2024-11-10 19:07:25 -05:00
commit 7a2fe139c4
1 changed files with 45 additions and 0 deletions
@@ -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}