Merge branch 'kijai:main' into main

Author: liubo0902

nodes/image_nodes.py

@@ -2468,14 +2468,15 @@ def INPUT_TYPES(s):
                 "width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1, }),
                 "height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1, }),
                 "upscale_method": (s.upscale_methods,),
-                "keep_proportion": (["stretch", "resize", "pad", "pad_edge", "pad_edge_pixel", "crop", "pillarbox_blur"], { "default": False }),
+                "keep_proportion": (["stretch", "resize", "pad", "pad_edge", "pad_edge_pixel", "crop", "pillarbox_blur", "total_pixels"], { "default": False }),
                 "pad_color": ("STRING", { "default": "0, 0, 0", "tooltip": "Color to use for padding."}),
                 "crop_position": (["center", "top", "bottom", "left", "right"], { "default": "center" }),
                 "divisible_by": ("INT", { "default": 2, "min": 0, "max": 512, "step": 1, }),
             },
             "optional" : {
                 "mask": ("MASK",),
                 "device": (["cpu", "gpu"],),
+                #"per_batch": ("INT", { "default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1, "tooltip": "Process images in sub-batches to reduce memory usage. 0 disables sub-batching."}),
             },
              "hidden": {
                 "unique_id": "UNIQUE_ID",
@@ -2494,7 +2495,7 @@ def INPUT_TYPES(s):
 highest dimension.  
 """
 
-    def resize(self, image, width, height, keep_proportion, upscale_method, divisible_by, pad_color, crop_position, unique_id, device="cpu", mask=None):
+    def resize(self, image, width, height, keep_proportion, upscale_method, divisible_by, pad_color, crop_position, unique_id, device="cpu", mask=None, per_batch=64):
         B, H, W, C = image.shape
 
         if device == "gpu":
@@ -2504,15 +2505,23 @@ def resize(self, image, width, height, keep_proportion, upscale_method, divisibl
         else:
             device = torch.device("cpu")
 
-        if width == 0:
-            width = W
-        if height == 0:
-            height = H
-
         pillarbox_blur = keep_proportion == "pillarbox_blur"
-        if keep_proportion == "resize" or keep_proportion.startswith("pad") or pillarbox_blur:
+        
+        # Initialize padding variables
+        pad_left = pad_right = pad_top = pad_bottom = 0
+
+        if keep_proportion in ["resize", "total_pixels"] or keep_proportion.startswith("pad") or pillarbox_blur:
+            if keep_proportion == "total_pixels":
+                total_pixels = width * height
+                aspect_ratio = W / H
+                new_height = int(math.sqrt(total_pixels / aspect_ratio))
+                new_width = int(math.sqrt(total_pixels * aspect_ratio))
+                
             # If one of the dimensions is zero, calculate it to maintain the aspect ratio
-            if width == 0 and height != 0:
+            elif width == 0 and height == 0:
+                new_width = W
+                new_height = H
+            elif width == 0 and height != 0:
                 ratio = height / H
                 new_width = round(W * ratio)
                 new_height = height
@@ -2528,7 +2537,6 @@ def resize(self, image, width, height, keep_proportion, upscale_method, divisibl
                 new_width = width
                 new_height = height
 
-            pad_left = pad_right = pad_top = pad_bottom = 0
             if keep_proportion.startswith("pad") or pillarbox_blur:
                 # Calculate padding based on position
                 if crop_position == "center":
@@ -2559,76 +2567,136 @@ def resize(self, image, width, height, keep_proportion, upscale_method, divisibl
 
             width = new_width
             height = new_height
+        else:
+            if width == 0:
+                width = W
+            if height == 0:
+                height = H
 
         if divisible_by > 1:
             width = width - (width % divisible_by)
             height = height - (height % divisible_by)
 
-        out_image = image.clone().to(device)
-        if mask is not None:
-            out_mask = mask.clone().to(device)
-        else:
-            out_mask = None
-
-        # Crop logic
-        if keep_proportion == "crop":
-            old_width = W
-            old_height = H
-            old_aspect = old_width / old_height
-            new_aspect = width / height
-            if old_aspect > new_aspect:
-                crop_w = round(old_height * new_aspect)
-                crop_h = old_height
-            else:
-                crop_w = old_width
-                crop_h = round(old_width / new_aspect)
-            if crop_position == "center":
-                x = (old_width - crop_w) // 2
-                y = (old_height - crop_h) // 2
-            elif crop_position == "top":
-                x = (old_width - crop_w) // 2
-                y = 0
-            elif crop_position == "bottom":
-                x = (old_width - crop_w) // 2
-                y = old_height - crop_h
-            elif crop_position == "left":
-                x = 0
-                y = (old_height - crop_h) // 2
-            elif crop_position == "right":
-                x = old_width - crop_w
-                y = (old_height - crop_h) // 2
-            out_image = out_image.narrow(-2, x, crop_w).narrow(-3, y, crop_h)
-            if mask is not None:
-                out_mask = out_mask.narrow(-1, x, crop_w).narrow(-2, y, crop_h)
+        # Preflight estimate (log-only when batching is active)
+        if per_batch != 0 and B > per_batch:
+            try:
+                bytes_per_elem = image.element_size()  # typically 4 for float32
+                est_total_bytes = B * height * width * C * bytes_per_elem
+                est_mb = est_total_bytes / (1024 * 1024)
+                msg = f"<tr><td>Resize v2</td><td>estimated output ~{est_mb:.2f} MB; batching {per_batch}/{B}</td></tr>"
+                if unique_id and PromptServer is not None:
+                    try:
+                        PromptServer.instance.send_progress_text(msg, unique_id)
+                    except:
+                        pass
+                else:
+                    print(f"[ImageResizeKJv2] estimated output ~{est_mb:.2f} MB; batching {per_batch}/{B}")
+            except:
+                pass
 
-        out_image = common_upscale(out_image.movedim(-1,1), width, height, upscale_method, crop="disabled").movedim(1,-1)
-        if mask is not None:
-            if upscale_method == "lanczos":
-                out_mask = common_upscale(out_mask.unsqueeze(1).repeat(1, 3, 1, 1), width, height, upscale_method, crop="disabled").movedim(1,-1)[:, :, :, 0]
-            else:
-                out_mask = common_upscale(out_mask.unsqueeze(1), width, height, upscale_method, crop="disabled").squeeze(1)
+        def _process_subbatch(in_image, in_mask, pad_left, pad_right, pad_top, pad_bottom):
+            # Avoid unnecessary clones; only move if needed
+            out_image = in_image if in_image.device == device else in_image.to(device)
+            out_mask = None if in_mask is None else (in_mask if in_mask.device == device else in_mask.to(device))
+
+            # Crop logic
+            if keep_proportion == "crop":
+                old_height = out_image.shape[-3]
+                old_width = out_image.shape[-2]
+                old_aspect = old_width / old_height
+                new_aspect = width / height
+                if old_aspect > new_aspect:
+                    crop_w = round(old_height * new_aspect)
+                    crop_h = old_height
+                else:
+                    crop_w = old_width
+                    crop_h = round(old_width / new_aspect)
+                if crop_position == "center":
+                    x = (old_width - crop_w) // 2
+                    y = (old_height - crop_h) // 2
+                elif crop_position == "top":
+                    x = (old_width - crop_w) // 2
+                    y = 0
+                elif crop_position == "bottom":
+                    x = (old_width - crop_w) // 2
+                    y = old_height - crop_h
+                elif crop_position == "left":
+                    x = 0
+                    y = (old_height - crop_h) // 2
+                elif crop_position == "right":
+                    x = old_width - crop_w
+                    y = (old_height - crop_h) // 2
+                out_image = out_image.narrow(-2, x, crop_w).narrow(-3, y, crop_h)
+                if out_mask is not None:
+                    out_mask = out_mask.narrow(-1, x, crop_w).narrow(-2, y, crop_h)
+
+            out_image = common_upscale(out_image.movedim(-1,1), width, height, upscale_method, crop="disabled").movedim(1,-1)
+            if out_mask is not None:
+                if upscale_method == "lanczos":
+                    out_mask = common_upscale(out_mask.unsqueeze(1).repeat(1, 3, 1, 1), width, height, upscale_method, crop="disabled").movedim(1,-1)[:, :, :, 0]
+                else:
+                    out_mask = common_upscale(out_mask.unsqueeze(1), width, height, upscale_method, crop="disabled").squeeze(1)
+
+            # Pad logic
+            if (keep_proportion.startswith("pad") or pillarbox_blur) and (pad_left > 0 or pad_right > 0 or pad_top > 0 or pad_bottom > 0):
+                padded_width = width + pad_left + pad_right
+                padded_height = height + pad_top + pad_bottom
+                if divisible_by > 1:
+                    width_remainder = padded_width % divisible_by
+                    height_remainder = padded_height % divisible_by
+                    if width_remainder > 0:
+                        extra_width = divisible_by - width_remainder
+                        pad_right += extra_width
+                    if height_remainder > 0:
+                        extra_height = divisible_by - height_remainder
+                        pad_bottom += extra_height
+
+                pad_mode = (
+                    "pillarbox_blur" if pillarbox_blur else
+                    "edge" if keep_proportion == "pad_edge" else
+                    "edge_pixel" if keep_proportion == "pad_edge_pixel" else
+                    "color"
+                )
+                out_image, out_mask = ImagePadKJ.pad(self, out_image, pad_left, pad_right, pad_top, pad_bottom, 0, pad_color, pad_mode, mask=out_mask)
 
-        # Pad logic
-        if (keep_proportion.startswith("pad") or pillarbox_blur) and (pad_left > 0 or pad_right > 0 or pad_top > 0 or pad_bottom > 0):
-            padded_width = width + pad_left + pad_right
-            padded_height = height + pad_top + pad_bottom
-            if divisible_by > 1:
-                width_remainder = padded_width % divisible_by
-                height_remainder = padded_height % divisible_by
-                if width_remainder > 0:
-                    extra_width = divisible_by - width_remainder
-                    pad_right += extra_width
-                if height_remainder > 0:
-                    extra_height = divisible_by - height_remainder
-                    pad_bottom += extra_height
-
-            pad_mode = (
-                "pillarbox_blur" if pillarbox_blur else
-                "edge" if keep_proportion == "pad_edge" else
-                "edge_pixel" if keep_proportion == "pad_edge_pixel" else
-                "color"
-            )
-            out_image, out_mask = ImagePadKJ.pad(self, out_image, pad_left, pad_right, pad_top, pad_bottom, 0, pad_color, pad_mode, mask=out_mask)
+            return out_image, out_mask
+
+        # If batching disabled (per_batch==0) or batch fits, process whole batch
+        if per_batch == 0 or B <= per_batch:
+            out_image, out_mask = _process_subbatch(image, mask, pad_left, pad_right, pad_top, pad_bottom)
+        else:
+            chunks = []
+            mask_chunks = [] if mask is not None else None
+            total_batches = (B + per_batch - 1) // per_batch
+            current_batch = 0
+            for start_idx in range(0, B, per_batch):
+                current_batch += 1
+                end_idx = min(start_idx + per_batch, B)
+                sub_img = image[start_idx:end_idx]
+                sub_mask = mask[start_idx:end_idx] if mask is not None else None
+                sub_out_img, sub_out_mask = _process_subbatch(sub_img, sub_mask, pad_left, pad_right, pad_top, pad_bottom)
+                chunks.append(sub_out_img.cpu())
+                if mask is not None:
+                    mask_chunks.append(sub_out_mask.cpu() if sub_out_mask is not None else None)
+                # Per-batch progress update
+                if unique_id and PromptServer is not None:
+                    try:
+                        PromptServer.instance.send_progress_text(
+                            f"<tr><td>Resize v2</td><td>batch {current_batch}/{total_batches} · images {end_idx}/{B}</td></tr>",
+                            unique_id
+                        )
+                    except:
+                        pass
+                else:
+                    try:
+                        print(f"[ImageResizeKJv2] batch {current_batch}/{total_batches} · images {end_idx}/{B}")
+                    except:
+                        pass
+            out_image = torch.cat(chunks, dim=0)
+            if mask is not None and any(m is not None for m in mask_chunks):
+                out_mask = torch.cat([m for m in mask_chunks if m is not None], dim=0)
+            else:
+                out_mask = None
 
         # Progress UI
         if unique_id and PromptServer is not None:

nodes/lora_nodes.py

@@ -48,7 +48,27 @@ def extract_lora(diff, key, rank, algorithm, lora_type, lowrank_iters=7, adaptiv
                 s_cum = torch.cumsum(S, dim=0)
                 min_cum_sum = adaptive_param * torch.sum(S)
                 lora_rank = torch.sum(s_cum < min_cum_sum).item()
-            print(f"{key} Extracted LoRA rank: {lora_rank}")
+            elif lora_type == "adaptive_fro":
+                S_squared = S.pow(2)
+                S_fro_sq = float(torch.sum(S_squared))
+                sum_S_squared = torch.cumsum(S_squared, dim=0) / S_fro_sq
+                lora_rank = int(torch.searchsorted(sum_S_squared, adaptive_param**2)) + 1
+                lora_rank = max(1, min(lora_rank, len(S)))
+            else:
+                pass  # Will print after capping
+            
+            # Cap adaptive rank by the specified max rank
+            lora_rank = min(lora_rank, rank)
+            
+            # Calculate and print actual fro percentage retained after capping
+            if lora_type == "adaptive_fro":
+                S_squared = S.pow(2)
+                s_fro = torch.sqrt(torch.sum(S_squared))
+                s_red_fro = torch.sqrt(torch.sum(S_squared[:lora_rank]))
+                fro_percent = float(s_red_fro / s_fro)
+                print(f"{key} Extracted LoRA rank: {lora_rank}, Frobenius retained: {fro_percent:.1%}")
+            else:
+                print(f"{key} Extracted LoRA rank: {lora_rank}")
         else:
             lora_rank = rank
 
@@ -141,13 +161,13 @@ def INPUT_TYPES(s):
                     "finetuned_model": ("MODEL",),
                     "original_model": ("MODEL",),
                     "filename_prefix": ("STRING", {"default": "loras/ComfyUI_extracted_lora"}),
-                    "rank": ("INT", {"default": 8, "min": 1, "max": 4096, "step": 1}),
-                    "lora_type": (["standard", "full", "adaptive_ratio", "adaptive_quantile", "adaptive_energy"],),
+                    "rank": ("INT", {"default": 8, "min": 1, "max": 4096, "step": 1, "tooltip": "The rank to use for standard LoRA, or maximum rank limit for adaptive methods."}),
+                    "lora_type": (["standard", "full", "adaptive_ratio", "adaptive_quantile", "adaptive_energy", "adaptive_fro"],),
                     "algorithm": (["svd_linalg", "svd_lowrank"], {"default": "svd_linalg", "tooltip": "SVD algorithm to use, svd_lowrank is faster but less accurate."}),
                     "lowrank_iters": ("INT", {"default": 7, "min": 1, "max": 100, "step": 1, "tooltip": "The number of subspace iterations for lowrank SVD algorithm."}),
                     "output_dtype": (["fp16", "bf16", "fp32"], {"default": "fp16"}),
                     "bias_diff": ("BOOLEAN", {"default": True}),
-                    "adaptive_param": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "For ratio mode, this is the ratio of the maximum singular value. For quantile mode, this is the quantile of the singular values."}),
+                    "adaptive_param": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "For ratio mode, this is the ratio of the maximum singular value. For quantile mode, this is the quantile of the singular values. For fro mode, this is the Frobenius norm retention ratio."}),
                     "clamp_quantile": ("BOOLEAN", {"default": True}),
                 },
 
@@ -520,7 +540,7 @@ def resize_lora_model(lora_sd, new_rank, save_dtype, device, dynamic_method, dyn
                 fro_retained = param_dict["fro_retained"]
                 if not np.isnan(fro_retained):
                     fro_list.append(float(fro_retained))
-                log_str = f"{block_down_name:75} | sum(S) retained: {sum_retained:.1%}, fro retained: {fro_retained:.1%}, max(S) ratio: {max_ratio:0.1f}"
+                log_str = f"{block_down_name:75} | sum(S) retained: {sum_retained:.1%}, fro retained: {fro_retained:.1%}, max(S) ratio: {max_ratio:0.1f}, new dim: {param_dict['new_rank']}"
                 tqdm.write(log_str)
                 verbose_str += log_str