updated upscale machine

2025-10-10 20:33:06 +05:30
parent 827eee49e3
commit dcd00d6c45
5 changed files with 128 additions and 36 deletions
@@ -46,7 +46,7 @@ class Resolution_Machine:
    RETURN_TYPES = ("INT", "INT")
    RETURN_NAMES = ("width", "height")
    FUNCTION = "get_resolution"
-    CATEGORY = "SATA Utility"
+    CATEGORY = "SATA_UtilityNode"

    def get_resolution(self, model, resolution, custom_width, custom_height):
        config = load_config()
@@ -6,6 +6,12 @@ import comfy.utils
 from comfy import model_management
 from spandrel import ModelLoader

+try:
+    import tensorrt as trt
+    TENSORRT_AVAILABLE = True
+except ImportError:
+    TENSORRT_AVAILABLE = False
+

 def resize_tensor_opencv(tensor_chw, target_width, target_height, supersample='true', factor=2.0):
    """
@@ -51,25 +57,75 @@ class Upscale_Machine:
        return {
            "required": {
                "image": ("IMAGE",),
-                "upscale_model": (folder_paths.get_filename_list("upscale_models"),),
+                "upscale_model": (folder_paths.get_filename_list("upscale_models"), {"default": None}),
+                "chained_model": (["None"] + folder_paths.get_filename_list("upscale_models"), {"default": "None"}),
                "rescale_factor": ("FLOAT", {"default": 2.0, "min": 0.01, "max": 16.0, "step": 0.01}),
                "supersample": (["true", "false"],),
                "rounding_modulus": ("INT", {"default": 8, "min": 1, "max": 1024, "step": 1}),
+                "use_tensorrt": ("BOOLEAN", {"default": False}),
            }
        }

-    RETURN_TYPES = ("IMAGE", "STRING",)
-    RETURN_NAMES = ("IMAGE", "show_help",)
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("IMAGE",)
    FUNCTION = "upscale"
    CATEGORY = "SATA_UtilityNode"

-    def load_model(self, model_name):
-        """Load ESRGAN/RealESRGAN/AESRGAN via spandrel"""
+    def load_model(self, model_name, use_tensorrt=False):
+        """Load ESRGAN/RealESRGAN/AESRGAN with optional TensorRT optimization"""
        model_path = folder_paths.get_full_path("upscale_models", model_name)
        model_loader = ModelLoader()
        model = model_loader.load_from_file(model_path)
+        
        if model is None:
            raise RuntimeError(f"Failed to load upscale model: {model_name}")
+            
+        if use_tensorrt and TENSORRT_AVAILABLE:
+            try:
+                # Convert model to TensorRT
+                logger = trt.Logger(trt.Logger.WARNING)
+                builder = trt.Builder(logger)
+                network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+                parser = trt.OnnxParser(network, logger)
+                
+                # Export model to ONNX first
+                torch.onnx.export(model, torch.randn(1, 3, 64, 64), "/tmp/temp.onnx")
+                
+                with open("/tmp/temp.onnx", 'rb') as f:
+                    parser.parse(f.read())
+                
+                config = builder.create_builder_config()
+                config.max_workspace_size = 1 << 30  # 1GB
+                engine = builder.build_engine(network, config)
+                
+                if engine is None:
+                    raise RuntimeError("Failed to build TensorRT engine")
+                
+                context = engine.create_execution_context()
+                
+                # Wrap TensorRT engine in a pytorch module
+                class TRTWrapper(torch.nn.Module):
+                    def __init__(self, engine, context):
+                        super().__init__()
+                        self.engine = engine
+                        self.context = context
+                        
+                    def forward(self, x):
+                        # Convert input to TensorRT
+                        inp = x.cpu().numpy()
+                        # Execute inference
+                        output = np.empty((1, 3, x.shape[2]*4, x.shape[3]*4), dtype=np.float32)
+                        self.context.execute_v2([inp.ctypes.data, output.ctypes.data])
+                        # Convert back to PyTorch
+                        return torch.from_numpy(output).to(x.device)
+                        
+                model = TRTWrapper(engine, context)
+                print(f"Successfully optimized {model_name} with TensorRT")
+                
+            except Exception as e:
+                print(f"TensorRT optimization failed: {e}")
+                print("Falling back to regular model")
+                
        model.eval()
        return model

@@ -117,42 +173,68 @@ class Upscale_Machine:
            return int(max(1, round(value)))
        return max(modulus, int(round(value / modulus)) * modulus)

-    def upscale(self, image, upscale_model, rounding_modulus=8, supersample='true',
-                rescale_factor=2.0):
-
+    def upscale(self, image, upscale_model, chained_model="None", rounding_modulus=8, supersample='true',
+                rescale_factor=2.0, use_tensorrt=False):

        if image.ndim != 4:
            raise ValueError("Expected IMAGE tensor with 4 dims (B,H,W,C).")

-        original_height = int(image.shape[1])  
-        original_width  = int(image.shape[2]) 
-
-
-        target_width  = self._round_to_modulus(original_width  * rescale_factor, rounding_modulus)
+        original_height = int(image.shape[1])
+        original_width = int(image.shape[2])
+        target_width = self._round_to_modulus(original_width * rescale_factor, rounding_modulus)
        target_height = self._round_to_modulus(original_height * rescale_factor, rounding_modulus)

+        current_bhwc = image

-        up_model = self.load_model(upscale_model)
-        up_bchw = self.upscale_with_model(up_model, image)
+        # First upscale
+        if upscale_model:
+            up_model = self.load_model(upscale_model, use_tensorrt)
+            up_bchw = self.upscale_with_model(up_model, current_bhwc)
+            current_bhwc = up_bchw.movedim(1, -1).contiguous()

-        batch_out = []
-        for i in range(up_bchw.shape[0]):
-            chw = up_bchw[i] 
-            resized_chw = resize_tensor_opencv(
-                chw, target_width=target_width, target_height=target_height,
-                supersample=supersample, factor=rescale_factor
-            )
+            # Downscale to target size before chaining
+            batch_downscaled = []
+            for i in range(current_bhwc.shape[0]):
+                chw = current_bhwc[i].movedim(-1, 0)  # (H,W,C) -> (C,H,W)
+                chw = torch.clamp(chw, 0.0, 1.0)
+                resized_chw = resize_tensor_opencv(
+                    chw, target_width=target_width, target_height=target_height,
+                    supersample=supersample, factor=rescale_factor
+                )
+                if resized_chw.shape[0] == 1:
+                    resized_chw = resized_chw.repeat(3, 1, 1)
+                elif resized_chw.shape[0] > 3:
+                    resized_chw = resized_chw[:3, :, :]
+                batch_downscaled.append(resized_chw.unsqueeze(0))
+            current_bhwc = torch.cat(batch_downscaled, dim=0)
+            current_bhwc = current_bhwc.movedim(1, -1).contiguous()

-            if resized_chw.shape[0] == 1:
-                resized_chw = resized_chw.repeat(3, 1, 1)
-            elif resized_chw.shape[0] > 3:
-                resized_chw = resized_chw[:3, :, :]
-            batch_out.append(resized_chw.unsqueeze(0)) 
+        # Second upscale (chained model) if selected
+        if chained_model and chained_model != "None":
+            chain_model = self.load_model(chained_model, use_tensorrt)
+            chain_bchw = self.upscale_with_model(chain_model, current_bhwc)
+            current_bhwc = chain_bchw.movedim(1, -1).contiguous()

-        out_bchw = torch.cat(batch_out, dim=0) 
+        # If no chained model, ensure output is at target size
+        if not (chained_model and chained_model != "None"):
+            # Already downscaled after first upscale, so just output
+            images_out = current_bhwc
+        else:
+            # After chained model, do a final resize to target size
+            batch_out = []
+            for i in range(current_bhwc.shape[0]):
+                chw = current_bhwc[i].movedim(-1, 0)
+                chw = torch.clamp(chw, 0.0, 1.0)
+                resized_chw = resize_tensor_opencv(
+                    chw, target_width=target_width, target_height=target_height,
+                    supersample=supersample, factor=rescale_factor
+                )
+                if resized_chw.shape[0] == 1:
+                    resized_chw = resized_chw.repeat(3, 1, 1)
+                elif resized_chw.shape[0] > 3:
+                    resized_chw = resized_chw[:3, :, :]
+                batch_out.append(resized_chw.unsqueeze(0))
+            out_bchw = torch.cat(batch_out, dim=0)
+            images_out = out_bchw.movedim(1, -1).contiguous()

- 
-        images_out = out_bchw.movedim(1, -1).contiguous()
-
-        show_help = "https://github.com/SatadalAI/help_image"
-        return (images_out, show_help)
+        return (images_out,)
@@ -1,2 +1,8 @@
 name,positive,negative,note
 Chibi,"full body, chibi, highly detailed","photo, black and white, low contrast",chibi_blindbox_v1_mix
+illustrious_masterpieces_v3,"masterpiece, best quality, very aesthetic",,
+RE3-30,"asian, asian_japanese, asian_chinese, asian_korean, european, european_irish, european_french, european_german, european_italian, african, african_east, african_west, latina, thai, arab, indian, iranian, native_american, polynesian",,
+aidmaMJ6.1_v0.5_IL,aidmaMJ6.1,,
+princess_rc_il,"Anna, Ariel, Aurora, Belle, Cinderella, Elsa, Asha, Jasmine, Merida, Moana, Mulan, Namaari, Pocahontas, Rapunzel, Raya, Shank, Snow White, Tiana, Vanellope, Alice, Chel,Esmerada, Jane Porter, Kida, Megera,Mirabel Madrigal, Isabela Madrigal, Dolores Madrigal and some secret character!
+",,
+DoubleExposureIL,"Silhouette, within, blending, blended, merged, filled with / into ",,
@@ -3,4 +3,5 @@ Photon1,,"cartoon, painting, illustration, (worst quality, low quality, normal q
 RealVis,"RAW photo, subject, 8k uhd, dslr, soft lighting, high quality, film grain, Fujifilm XT3","(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, mutated hands and fingers:1.4), (deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, amputation, UnrealisticDream",sd1.5
 RevAnime,"((best quality)), ((masterpiece)), (detailed) ",,sd1.5
 ilustmix,"masterpiece, best quality, amazing quality, very aesthetic, detailed eyes, perfect eyes, realistic eyes",,Sampler: Euler a Steps: 25
-WAI,masterpiece,,Steps: 15-30 CFG scale: 5-7 Sampler: Euler a
+WAI,"masterpiece,best quality,","bad quality,worst quality,worst detail,",Steps: 15-30 CFG scale: 5-7 Sampler: Euler a
+Plantmink_walnut,"masterpiece, best quality",," Euler or Euler a or HeunPP2 or Restart sampler, 3 or ~3-6 cfg, ~28 steps"
@@ -1,2 +1,5 @@
 opencv-python>=4.5.0
-spandrel
+spandrel
+nvidia-tensorrt>=8.0
+onnx>=1.12.0
+onnxruntime-gpu>=1.12.0