from nodes import SaveImage
import json
from PIL import Image
import numpy as np
from PIL.PngImagePlugin import PngInfo
from comfy.cli_args import args # type: ignore
import folder_paths # type: ignore
from folder_paths import get_filename_list # type: ignore
import comfy
import os
import re
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import psutil
import time
from typing import Any, Dict, List, Tuple, Optional
from collections import OrderedDict, defaultdict
import threading

class ControlNetSelector:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "control_net_name": (get_filename_list("controlnet"),)
            }
        }
    
    RETURN_TYPES = (folder_paths.get_filename_list("controlnet"), )
    RETURN_NAMES = ("control_net_name",)
    FUNCTION = "get_control_net_name"
    
    CATEGORY = 'Uber Comfy'

    def get_control_net_name(self, control_net_name):
        return (control_net_name,)

class ControlNetOptionalLoader:
    @classmethod
    def INPUT_TYPES(s):
        # Extending the file list with a 'None' option for manual selection
        return {"required": { "control_net_name": (["None"] + folder_paths.get_filename_list("controlnet"), )}}

    RETURN_TYPES = ("CONTROL_NET",)
    FUNCTION = "load_controlnet"

    CATEGORY = "Uber Comfy"

    def load_controlnet(self, control_net_name):
        # Only proceed if a control_net_name is provided and it is not 'None'
        if control_net_name and control_net_name != "None":
            controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
            controlnet = comfy.controlnet.load_controlnet(controlnet_path)
            return (controlnet,)
        # Return None or skip the operation if 'None' is selected or no input is provided
        return (None,)

class DiffusersSelector:
    CATEGORY = 'Uber Comfy'
    RETURN_TYPES = (folder_paths.get_folder_paths("diffusers"), )
    RETURN_NAMES = ("model_path",)
    FUNCTION = "select_model_path"

    @classmethod
    def INPUT_TYPES(cls):
        paths = []
        for search_path in folder_paths.get_folder_paths("diffusers"):
            if os.path.exists(search_path):
                for root, subdirs, files in os.walk(search_path, followlinks=True):
                    if "model_index.json" in files:
                        paths.append(os.path.relpath(root, start=search_path))
        return {"required": {"model_path": (paths,), }}

    def select_model_path(self, model_path):
        # This function simply returns the model path that was selected
        return (model_path,)

class SaveImageJPGNoMeta(SaveImage):
    @classmethod
    def INPUT_TYPES(s):
        output = {
            "required": {
                "images": ("IMAGE",),
                "filename_prefix": ("STRING", {"default": "ComfyUI"}),
                "quality": ("INT", {"default": 100, "min": 1, "max": 100, "step": 1}),
            },
        }

        return output

    CATEGORY = "Uber Comfy"
    RETURN_TYPES = ()
    FUNCTION = "suplex_save_images"

    def suplex_save_images(
        self,
        images,
        filename_prefix="ComfyUI",
        format="jpeg",
        quality=92,
    ):
        filename_prefix += self.prefix_append
        full_output_folder, filename, counter, subfolder, filename_prefix = (
            folder_paths.get_save_image_path(
                filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]
            )
        )
        results = list()
        for batch_number, image in enumerate(images):
            i = 255.0 * image.cpu().numpy()
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
            file = f"{filename_with_batch_num}_{counter:05}_.{format}"
            results.append(
                {"filename": file, "subfolder": subfolder, "type": self.type}
            )
            counter += 1
        
        img.save(os.path.join(full_output_folder, file), quality=quality, optimize=True)
        return {"ui": {"images": results}}

class MultiInputVariableRewrite:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": ("STRING", {"multiline": True}),
            },
            "optional": {
                "a": ("STRING", {"forceInput": True}),
                "b": ("STRING", {"forceInput": True}),
                "c": ("STRING", {"forceInput": True}),
                "d": ("STRING", {"forceInput": True}),
                "e": ("STRING", {"forceInput": True}),
                # ... add more up to 'z' if needed
            }
        }
   
    CATEGORY = "Uber Comfy"
    FUNCTION = "multicombinetext"
    RETURN_NAMES = ("TEXT",)
    RETURN_TYPES = ("STRING",)

    def multicombinetext(self, text="", **kwargs):
        for key, value in kwargs.items():
            if value:
                text = text.replace(f"{{{key}}}", value)
        return (text,)
    
class TextRegexOperations:
    @classmethod
    def INPUT_TYPES(cls):
        inputs = {
            "required": {
                "text": ("STRING", {"multiline": True, "forceInput": True}),
                "num_operations": ("INT", {"default": 1, "min": 1, "max": 20, "step": 1}),
            },
            "optional": {}
        }
        
        # Create inputs in interleaved order (pattern_1, replacement_1, multiline_1, pattern_2, etc.)
        for i in range(1, 21):
            inputs["optional"][f"pattern_{i}"] = ("STRING", {"multiline": True})
            inputs["optional"][f"replacement_{i}"] = ("STRING", {"multiline": True})
            inputs["optional"][f"use_multiline_{i}"] = ("BOOLEAN", {"default": True})
        
        return inputs
    
    RETURN_TYPES = ("STRING",)
    FUNCTION = "process_text"
    CATEGORY = "Uber Comfy"

    def process_text(self, text, num_operations, **kwargs):
        processed_text = text
        
        for i in range(1, num_operations + 1):
            pattern = kwargs.get(f"pattern_{i}", "")
            replacement = kwargs.get(f"replacement_{i}", "")
            use_multiline = kwargs.get(f"use_multiline_{i}", True)
            
            if pattern:
                try:
                    flags = re.MULTILINE if use_multiline else 0
                    processed_text = re.sub(pattern, replacement, processed_text, flags=flags)
                except re.error as e:
                    print(f"Regex error in operation {i}: {str(e)}")
                    print(f"Pattern: {pattern}")
        
        return (processed_text,)
    
class VideoSegmentCalculator:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "duration": ("FLOAT", {
                    "default": 30.0, 
                    "min": 1.0, 
                    "max": 3600.0, 
                    "step": 1.0, 
                    "tooltip": "Duration of each segment in seconds"
                }),
                "frame_rate": ("FLOAT", {
                    "default": 25.0, 
                    "min": 1.0, 
                    "max": 120.0, 
                    "step": 0.1,
                    "tooltip": "Frame rate of the video"
                }),
                "index": ("INT", {
                    "default": 0, 
                    "min": 0, 
                    "max": 1000, 
                    "step": 1,
                    "tooltip": "Current segment index (0-based)"
                }),
            },
            "optional": {
                "overlap_frames": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 10,
                    "step": 1,
                    "tooltip": "Number of frames to overlap between segments"
                }),
                "precise_timing": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Use precise decimal timing for audio trimming"
                }),
            }
        }
    
    RETURN_TYPES = ("INT", "INT", "FLOAT", "FLOAT", "FLOAT")
    RETURN_NAMES = ("frame_load_cap", "skip_first_frames", "force_rate", "start_time", "end_time")
    FUNCTION = "calculate_segment"
    CATEGORY = "Uber Comfy"
    
    def calculate_segment(self, duration, frame_rate, index, overlap_frames=0, precise_timing=True):
        # Calculate the exact frames for the given duration
        exact_frames = duration * frame_rate
        
        # Calculate the number of frames in each segment (using ceiling to prevent gaps)
        frames_per_segment = math.ceil(exact_frames)
        
        # Calculate skip_first_frames based on index, with optional overlap
        if index == 0:
            skip_first_frames = 0
        else:
            skip_first_frames = index * frames_per_segment - overlap_frames
            # Ensure we don't go negative
            skip_first_frames = max(0, skip_first_frames)
        
        # Calculate start and end times based on exact frame positions
        start_time = skip_first_frames / frame_rate
        end_time = (skip_first_frames + frames_per_segment) / frame_rate
        
        # Round to 2 decimal places for audio timing if needed
        if not precise_timing:
            start_time = round(start_time, 2)
            end_time = round(end_time, 2)
        
        return (frames_per_segment, skip_first_frames, frame_rate, start_time, end_time)
    

class ModelSimilarityNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {
            "base_model":   ("MODEL",),
            "target_model": ("MODEL",)}
        }

    RETURN_TYPES  = ("STRING",)
    RETURN_NAMES  = ("similarity_report",)
    FUNCTION      = "compare"
    CATEGORY      = "Uber Comfy"

    # ---------- helpers ----------
    @staticmethod
    def _cross(to_q, to_k, to_v, rnd):
        h, e = to_q.shape
        lq = nn.Linear(h, e, bias=False); lq.weight.copy_(to_q)
        lk = nn.Linear(h, e, bias=False); lk.weight.copy_(to_k)
        lv = nn.Linear(h, e, bias=False); lv.weight.copy_(to_v)
        return torch.einsum(
            "ik,jk->ik",
            torch.softmax(torch.einsum("ij,kj->ik", lq(rnd), lk(rnd)), dim=-1),
            lv(rnd)
        )

    @classmethod
    def _unwrap(cls, obj):
        if isinstance(obj, dict):               return obj
        if hasattr(obj, "get_weights"):         return obj.get_weights()
        if hasattr(obj, "state_dict"):          return obj.state_dict()
        for a in ("model", "original_model"):
            if hasattr(obj, a):                 return cls._unwrap(getattr(obj, a))
        raise TypeError("Cannot unwrap MODEL object")

    # ---------- main ----------
    def compare(self, base_model, target_model):
        b_sd = self._unwrap(base_model)
        t_sd = self._unwrap(target_model)
        torch.manual_seed(114514)

        rnd, b_attn, sims, common = {}, {}, [], []

        # prefixes to scan: input, middle, output
        scan = [
            ("diffusion_model.input_blocks",  5),   # indices 0-4
            ("diffusion_model.middle_block", 1),    # index   0
            ("diffusion_model.output_blocks",11),   # indices 0-10
        ]

        # discover layers present in both models
        for prefix, max_idx in scan:
            for i in range(max_idx):
                key = f"{prefix}.{i}.1.transformer_blocks.0.attn1.to_q.weight"
                if key in b_sd and key in t_sd:
                    common.append((prefix, i))

        if not common:
            return ("No matching attention layers found in both models.",)

        # compute attention outputs for base model
        for p, i in common:
            q = b_sd[f"{p}.{i}.1.transformer_blocks.0.attn1.to_q.weight"]
            k = b_sd[f"{p}.{i}.1.transformer_blocks.0.attn1.to_k.weight"]
            v = b_sd[f"{p}.{i}.1.transformer_blocks.0.attn1.to_v.weight"]
            h, e = q.shape
            key = f"{p}.{i}"
            rnd[key]   = torch.randn(e, h)
            b_attn[key] = self._cross(q, k, v, rnd[key])

        # compare with target model
        for p, i in common:
            key = f"{p}.{i}"
            t_attn = self._cross(
                t_sd[f"{p}.{i}.1.transformer_blocks.0.attn1.to_q.weight"],
                t_sd[f"{p}.{i}.1.transformer_blocks.0.attn1.to_k.weight"],
                t_sd[f"{p}.{i}.1.transformer_blocks.0.attn1.to_v.weight"],
                rnd[key]
            )
            sims.append(torch.mean(torch.cosine_similarity(b_attn[key], t_attn)))

        score = torch.mean(torch.stack(sims)) * 100
        return (f"Similarity: {score:.2f}%  (compared {len(common)} blocks)",)
    
class ModelWeightDumperNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {
            "model": ("MODEL",),
            "show_shapes": ("BOOLEAN", {"default": True}),
            "filter_prefix": ("STRING", {"default": ""}),
        }}

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("weight_info",)
    FUNCTION = "dump_weights"
    CATEGORY = "Uber Comfy"

    @classmethod
    def _unwrap(cls, obj):
        if isinstance(obj, dict):               return obj
        if hasattr(obj, "get_weights"):         return obj.get_weights()
        if hasattr(obj, "state_dict"):          return obj.state_dict()
        for a in ("model", "original_model"):
            if hasattr(obj, a):                 return cls._unwrap(getattr(obj, a))
        raise TypeError("Cannot unwrap MODEL object")

    def dump_weights(self, model, show_shapes=True, filter_prefix=""):
        sd = self._unwrap(model)
        
        output_lines = []
        output_lines.append(f"Total weights: {len(sd)}\n")
        output_lines.append("="*80 + "\n")
        
        # Filter keys if prefix is provided
        keys = sorted(sd.keys())
        if filter_prefix:
            keys = [k for k in keys if k.startswith(filter_prefix)]
            output_lines.append(f"Filtered by prefix: '{filter_prefix}'\n")
            output_lines.append(f"Matching weights: {len(keys)}\n")
            output_lines.append("="*80 + "\n")
        
        for key in keys:
            weight = sd[key]
            if show_shapes:
                shape_str = f" → {tuple(weight.shape)}" if hasattr(weight, 'shape') else ""
                dtype_str = f" ({weight.dtype})" if hasattr(weight, 'dtype') else ""
                output_lines.append(f"{key}{shape_str}{dtype_str}\n")
            else:
                output_lines.append(f"{key}\n")
        
        return ("".join(output_lines),)
    
class RunwareResolutionCalculator:
    """ComfyUI node for intelligent resolution calculation with model-specific presets."""
    
    PRESETS = {
        "Nano Banana 2": [
            # Square formats (1:1)
            (1024, 1024), (2048, 2048), (4096, 4096),
            # Landscape 3:2 ratio
            (1200, 896), (2400, 1792), (4800, 3584),
            # Portrait 2:3 ratio
            (896, 1200), (1792, 2400), (3584, 4800),
            # Landscape 5:4 ratio
            (1152, 928), (2304, 1856), (4608, 3712),
            # Portrait 4:5 ratio
            (928, 1152), (1856, 2304), (3712, 4608),
            # Landscape 3:2 wider
            (1264, 848), (2528, 1696), (5056, 3392), (5096, 3392),
            # Portrait 2:3 taller
            (848, 1264), (1696, 2528), (3392, 5056), (3392, 5096),
            # Ultra-wide landscape
            (1376, 768), (2752, 1536), (5504, 3072),
            # Ultra-tall portrait
            (768, 1376), (1536, 2752), (3072, 5504),
            # Cinematic ultra-wide
            (1548, 672), (1584, 672), (3168, 1344), (6336, 2688),
        ],
    }
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "full_image": ("IMAGE",),
                "model_preset": (list(cls.PRESETS.keys()), {
                    "default": "Nano Banana 2"
                }),
            },
            "optional": {
                "cropped_image": ("IMAGE",),
                "mask": ("MASK",),
            }
        }
    
    RETURN_TYPES = ("INT", "INT")
    RETURN_NAMES = ("width", "height")
    FUNCTION = "calculate"
    CATEGORY = "Uber Comfy"
    
    def calculate(self, full_image, model_preset, cropped_image=None, mask=None):
        """Main calculation entry point."""
        _, full_h, full_w, _ = full_image.shape
        resolutions = self.PRESETS.get(model_preset, self.PRESETS["Nano Banana 2"])
        
        if cropped_image is None or mask is None:
            return self._direct_resolution_selection(full_w, full_h, resolutions, model_preset)
        
        return self._mask_constrained_selection(full_image, mask, resolutions, model_preset)
    
    def _direct_resolution_selection(self, img_w, img_h, resolutions, preset_name):
        """Select resolution using multi-tier filtering and ranking."""
        print(f"🎯 {preset_name} Optimizer (Direct Resolution Mode)")
        print(f"   Source dimensions: {img_w}×{img_h}")
        
        source_ratio = img_w / img_h
        source_area = img_w * img_h
        
        # Tier 1: Spatial filtering
        spatial_candidates = [
            (w, h) for w, h in resolutions 
            if w <= img_w and h <= img_h
        ]
        
        if not spatial_candidates:
            print(f"   ⚠️ No candidates within bounds, using 1024x1024")
            return (1024, 1024)
        
        # Tier 2: Aspect ratio bucketing
        def get_aspect_bucket(ratio):
            if 0.95 <= ratio <= 1.05:
                return "square"
            elif ratio > 1.5:
                return "wide"
            elif ratio < 0.67:
                return "tall"
            else:
                return "standard"
        
        source_bucket = get_aspect_bucket(source_ratio)
        
        # Tier 3: Multi-criteria scoring
        scored_candidates = []
        
        for res_w, res_h in spatial_candidates:
            res_ratio = res_w / res_h
            res_area = res_w * res_h
            res_bucket = get_aspect_bucket(res_ratio)
            
            log_ratio_distance = abs(np.log2(source_ratio) - np.log2(res_ratio))
            aspect_score = 1.0 / (1.0 + log_ratio_distance)
            
            utilization = res_area / source_area
            utilization_score = utilization
            
            bucket_bonus = 1.2 if res_bucket == source_bucket else 1.0
            
            quality_tier = min(1.0, np.log2(res_area / 1048576) / 2)
            
            # Weighted product
            composite = (
                (aspect_score ** 2.8) * 
                (utilization_score ** 1.5) * 
                bucket_bonus * 
                (1.0 + quality_tier * 0.3)
            )
            
            scored_candidates.append({
                'resolution': (res_w, res_h),
                'score': composite,
                'aspect_score': aspect_score,
                'utilization': utilization
            })
        
        # Maximize score
        best = max(scored_candidates, key=lambda x: x['score'])
        
        print(f"   ✓ Selected: {best['resolution'][0]}×{best['resolution'][1]} "
              f"(aspect: {best['aspect_score']:.3f}, util: {best['utilization']:.3f})")
        return best['resolution']
    
    def _mask_constrained_selection(self, full_image, mask_tensor, resolutions, preset_name):
        """Select resolution using bi-directional fitting analysis."""
        print(f"🎯 {preset_name} Optimizer (Mask-Based Mode)")
        
        _, full_h, full_w, _ = full_image.shape
        
        mask_data = mask_tensor.cpu().numpy()[0]
        active_pixels = mask_data > 0.5
        
        y_coords, x_coords = np.where(active_pixels)
        
        if len(y_coords) == 0 or len(x_coords) == 0:
            return (1024, 1024)
        
        roi_x1, roi_x2 = x_coords.min(), x_coords.max()
        roi_y1, roi_y2 = y_coords.min(), y_coords.max()
        roi_width = roi_x2 - roi_x1 + 1
        roi_height = roi_y2 - roi_y1 + 1
        
        print(f"   ROI bounds: {roi_width}×{roi_height}")
        
        roi_ratio = roi_width / roi_height
        roi_area = roi_width * roi_height
        
        scored_candidates = []
        
        for res_w, res_h in resolutions:
            target_ratio = res_w / res_h
            
            # Three fitting strategies
            fit_a_w = roi_width
            fit_a_h = int(roi_width / target_ratio)
            
            fit_b_h = roi_height
            fit_b_w = int(roi_height * target_ratio)
            
            scale_to_ratio = (target_ratio / roi_ratio) ** 0.5
            fit_c_w = int(roi_width * scale_to_ratio)
            fit_c_h = int(roi_height / scale_to_ratio)
            
            viable_fits = []
            
            for fit_w, fit_h, strategy in [
                (fit_a_w, fit_a_h, 'width_anchor'),
                (fit_b_w, fit_b_h, 'height_anchor'),
                (fit_c_w, fit_c_h, 'proportional')
            ]:
                if fit_w <= full_w and fit_h <= full_h:
                    expansion_w = fit_w - roi_width
                    expansion_h = fit_h - roi_height
                    total_expansion = expansion_w + expansion_h
                    
                    viable_fits.append({
                        'width': fit_w,
                        'height': fit_h,
                        'expansion': total_expansion,
                        'strategy': strategy
                    })
            
            if not viable_fits:
                continue
            
            best_fit = min(viable_fits, key=lambda x: x['expansion'])
            
            fitted_area = best_fit['width'] * best_fit['height']
            target_area = res_w * res_h
            
            expansion_ratio = fitted_area / roi_area
            expansion_efficiency = 1.0 / expansion_ratio
            
            scale_match = min(fitted_area, target_area) / max(fitted_area, target_area)
            
            strategy_weight = 1.1 if best_fit['strategy'] == 'proportional' else 1.0
            
            # Harmonic mean
            harmonic_mean = 3.0 / (
                (1.0 / expansion_efficiency) + 
                (1.0 / scale_match) + 
                (1.0 / strategy_weight)
            )
            
            scored_candidates.append({
                'resolution': (res_w, res_h),
                'score': harmonic_mean,
                'fitted': (best_fit['width'], best_fit['height']),
                'expansion': expansion_ratio,
                'strategy': best_fit['strategy']
            })
        
        if not scored_candidates:
            return (1024, 1024)
        
        best = max(scored_candidates, key=lambda x: x['score'])
        
        print(f"   Fitted area: {best['fitted'][0]}×{best['fitted'][1]} ({best['strategy']})")
        print(f"   ✓ Selected: {best['resolution'][0]}×{best['resolution'][1]} "
              f"(expansion: {best['expansion']:.2f}x)")
        return best['resolution']

class AdaptiveImageScaler:
    """Intelligent image scaling with optional ML-based upscaling and dimension constraints."""
    
    INTERPOLATION_MODES = [
        "lanczos", "bicubic", "bilinear", "area", "nearest-exact"
    ]
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "target_width": ("INT",),
                "target_height": ("INT",),
                "interpolation": (cls.INTERPOLATION_MODES,),
                "dimension_alignment": ("INT", {
                    "default": 0, "min": 0, "max": 256, "step": 1
                }),
            },
            "optional": {
                "upscale_model": ("UPSCALE_MODEL",),
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "process"
    CATEGORY = "Uber Comfy"
    
    MODEL_ENGAGEMENT_THRESHOLD = 1.08
    TILE_REDUCTION_FACTOR = 0.75
    MIN_SAFE_TILE_SIZE = 128
    
    def process(self, image, target_width, target_height, interpolation, 
                dimension_alignment, upscale_model=None):
        """Main processing pipeline."""
        import comfy.utils
        from comfy import model_management
        
        batch_size, source_h, source_w, channels = image.shape
        
        if dimension_alignment > 0:
            target_width = (target_width // dimension_alignment) * dimension_alignment
            target_height = (target_height // dimension_alignment) * dimension_alignment
        
        scale_w = target_width / source_w
        scale_h = target_height / source_h
        geometric_mean_scale = (scale_w * scale_h) ** 0.5
        
        print(f"🖼️  Adaptive Scaler: {source_w}×{source_h}×{channels} → {target_width}×{target_height}")
        print(f"    Scale: {geometric_mean_scale:.3f}x (w:{scale_w:.2f}, h:{scale_h:.2f}) | {interpolation}")
        
        use_ml_model = (
            geometric_mean_scale > self.MODEL_ENGAGEMENT_THRESHOLD and 
            upscale_model is not None
        )
        
        if use_ml_model:
            return self._ml_assisted_pipeline(
                image, target_width, target_height, interpolation, channels,
                upscale_model, comfy, model_management
            )
        else:
            return self._standard_interpolation_pipeline(
                image, target_width, target_height, interpolation, 
                geometric_mean_scale, comfy
            )
    
    def _ml_assisted_pipeline(self, image, target_w, target_h, interpolation, 
                              channels, model, comfy, model_mgmt):
        """ML-based upscaling with synchronized alpha processing."""
        print(f"    ⚡ ML upscaling engaged")
        
        has_alpha = (channels == 4)
        
        if has_alpha:
            print(f"    💎 Alpha channel: synchronized processing")
            rgb_channels = image[:, :, :, :3]
            alpha_channel = image[:, :, :, 3:4]
        else:
            rgb_channels = image
            alpha_channel = None
        
        device = model_mgmt.get_torch_device()
        model.to(device)
        
        rgb_gpu = rgb_channels.movedim(-1, -3).to(device)
        
        upscaled_rgb = self._progressive_tiled_upscale(
            rgb_gpu, model, comfy, model_mgmt
        )
        
        model.cpu()
        
        upscaled_rgb = torch.clamp(upscaled_rgb.movedim(-3, -1), 0.0, 1.0)
        
        upscaled_rgb = upscaled_rgb.movedim(-1, 1)
        upscaled_rgb = comfy.utils.common_upscale(
            upscaled_rgb, target_w, target_h, interpolation, "disabled"
        )
        upscaled_rgb = upscaled_rgb.movedim(1, -1)
        
        if has_alpha:
            # Bicubic for quality
            alpha_scaled = self._scale_alpha_channel(
                alpha_channel, target_w, target_h, use_bicubic=True
            )
            result = torch.cat([upscaled_rgb, alpha_scaled], dim=-1)
            print(f"    💎 Alpha reconstructed with bicubic interpolation")
        else:
            result = upscaled_rgb
        
        return (result,)
    
    def _progressive_tiled_upscale(self, tensor, model, comfy, model_mgmt):
        """Progressive tile reduction (0.75x vs 0.5x)."""
        current_tile = 512
        overlap = 32
        
        while True:
            try:
                steps = tensor.shape[0] * comfy.utils.get_tiled_scale_steps(
                    tensor.shape[3], tensor.shape[2],
                    tile_x=current_tile, tile_y=current_tile, overlap=overlap
                )
                
                pbar = comfy.utils.ProgressBar(steps)
                
                result = comfy.utils.tiled_scale(
                    tensor,
                    lambda t: model(t),
                    tile_x=current_tile, tile_y=current_tile,
                    overlap=overlap,
                    upscale_amount=model.scale,
                    pbar=pbar
                )
                
                return result
                
            except model_mgmt.OOM_EXCEPTION as e:
                new_tile = int(current_tile * self.TILE_REDUCTION_FACTOR)
                
                if new_tile < self.MIN_SAFE_TILE_SIZE:
                    print(f"    ❌ Tile size below safe threshold ({self.MIN_SAFE_TILE_SIZE}px)")
                    raise e
                
                print(f"    ⚠️  Memory limit reached, reducing tile: {current_tile}→{new_tile}px")
                current_tile = new_tile
    
    def _scale_alpha_channel(self, alpha, target_w, target_h, use_bicubic=True):
        """Scale alpha using squeeze/unsqueeze (different from permute)."""
        alpha_reshaped = alpha.squeeze(-1).unsqueeze(1)
        
        mode = 'bicubic' if use_bicubic else 'bilinear'
        
        alpha_scaled = F.interpolate(
            alpha_reshaped,
            size=(target_h, target_w),
            mode=mode,
            align_corners=False if mode == 'bilinear' else None
        )
        
        alpha_result = alpha_scaled.squeeze(1).unsqueeze(-1)
        
        return alpha_result
    
    def _standard_interpolation_pipeline(self, image, target_w, target_h, 
                                          method, scale, comfy):
        """Direct interpolation with integrated channels."""
        direction = "⬇️ downscale" if scale < 1.0 else "⬆️ upscale"
        print(f"    {direction} via {method} interpolation")
        
        tensor_bchw = image.movedim(-1, 1)
        scaled_bchw = comfy.utils.common_upscale(
            tensor_bchw, target_w, target_h, method, "disabled"
        )
        result = scaled_bchw.movedim(1, -1)
        
        return (result,)

NODE_CLASS_MAPPINGS = {
    "ControlNet Selector": ControlNetSelector,
    "ControlNetOptionalLoader": ControlNetOptionalLoader,
    "DiffusersSelector": DiffusersSelector,
    "SaveImageJPGNoMeta": SaveImageJPGNoMeta,
    "MultiInputVariableRewrite": MultiInputVariableRewrite,
    "TextRegexOperations": TextRegexOperations,
    "VideoSegmentCalculator": VideoSegmentCalculator,
    "ModelSimilarityNode": ModelSimilarityNode,
    "ModelWeightDumperNode": ModelWeightDumperNode,
    "RunwareResolutionCalculator": RunwareResolutionCalculator,
    "AdaptiveImageScaler": AdaptiveImageScaler,

}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ControlNet Selector": "ControlNet Selector",
    "ControlNetOptionalLoader": "Load Optional ControlNet Model",
    "DiffusersSelector": "Diffusers Selector",
    "SaveImageJPGNoMeta": "Save Image JPG No Meta",
    "MultiInputVariableRewrite": "Multi Input Variable Rewrite",
    "TextRegexOperations": "Text Regex Operations",
    "VideoSegmentCalculator": "Video Segment Calculator",
    "ModelSimilarityNode": "Model Similarity Node",
    "ModelWeightDumperNode": "Model Weight Dumper",
    "RunwareResolutionCalculator": "Runware Resolution Calculator",
    "AdaptiveImageScaler": "Adaptive Image Scaler",
}