diff --git a/nodes/resolution_machine.py b/nodes/resolution_machine.py index 9407f81..7e0f241 100644 --- a/nodes/resolution_machine.py +++ b/nodes/resolution_machine.py @@ -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() diff --git a/nodes/upscale_machine.py b/nodes/upscale_machine.py index 61d60a4..44e2f3e 100644 --- a/nodes/upscale_machine.py +++ b/nodes/upscale_machine.py @@ -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) \ No newline at end of file + return (images_out,) \ No newline at end of file diff --git a/prompts/lora.csv b/prompts/lora.csv index b09b2cf..7a3b94a 100644 --- a/prompts/lora.csv +++ b/prompts/lora.csv @@ -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 ",, diff --git a/prompts/models.csv b/prompts/models.csv index 318f8c9..73a4dce 100644 --- a/prompts/models.csv +++ b/prompts/models.csv @@ -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" diff --git a/requirements.txt b/requirements.txt index e8c0ea1..1a2332e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,5 @@ opencv-python>=4.5.0 -spandrel \ No newline at end of file +spandrel +nvidia-tensorrt>=8.0 +onnx>=1.12.0 +onnxruntime-gpu>=1.12.0 \ No newline at end of file