custom_nodes/ComfyUI-Easy-Use/py/libs/xyplot.py

import os, torch
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont
from .utils import easySave, get_sd_version
from .adv_encode import advanced_encode
from .controlnet import easyControlnet
from .log import log_node_warn
from ..modules.layer_diffuse import LayerDiffuse
from ..config import RESOURCES_DIR
from nodes import CLIPTextEncode
import pprint
try:
    from comfy_extras.nodes_flux import FluxGuidance
except:
    FluxGuidance = None

class easyXYPlot():

    def __init__(self, xyPlotData, save_prefix, image_output, prompt, extra_pnginfo, my_unique_id, sampler, easyCache):
        self.x_node_type, self.x_type = sampler.safe_split(xyPlotData.get("x_axis"), ': ')
        self.y_node_type, self.y_type = sampler.safe_split(xyPlotData.get("y_axis"), ': ')
        self.x_values = xyPlotData.get("x_vals") if self.x_type != "None" else []
        self.y_values = xyPlotData.get("y_vals") if self.y_type != "None" else []
        self.custom_font = xyPlotData.get("custom_font")

        self.grid_spacing = xyPlotData.get("grid_spacing")
        self.latent_id = 0
        self.output_individuals = xyPlotData.get("output_individuals")

        self.x_label, self.y_label = [], []
        self.max_width, self.max_height = 0, 0
        self.latents_plot = []
        self.image_list = []

        self.num_cols = len(self.x_values) if len(self.x_values) > 0 else 1
        self.num_rows = len(self.y_values) if len(self.y_values) > 0 else 1

        self.total = self.num_cols * self.num_rows
        self.num = 0

        self.save_prefix = save_prefix
        self.image_output = image_output
        self.prompt = prompt
        self.extra_pnginfo = extra_pnginfo
        self.my_unique_id = my_unique_id

        self.sampler = sampler
        self.easyCache = easyCache

    # Helper Functions
    @staticmethod
    def define_variable(plot_image_vars, value_type, value, index):

        plot_image_vars[value_type] = value
        if value_type in ["seed", "Seeds++ Batch"]:
            value_label = f"seed: {value}"
        else:
            value_label = f"{value_type}: {value}"

        if "ControlNet" in value_type:
            value_label = f"ControlNet {index + 1}"

        if value_type in ['Lora', 'Checkpoint']:
            arr = value.split(',')
            model_name = os.path.basename(os.path.splitext(arr[0])[0])
            trigger_words = ' ' + arr[3] if value_type == 'Lora' and len(arr) > 3 and len(arr[3]) > 2 else ''
            lora_weight = float(arr[1]) if value_type == 'Lora' and len(arr) > 1 else 0
            lora_weight_desc = f" w:{lora_weight:.2f}" if value_type == 'Lora' and lora_weight != 1.0 else ''
            value_label = f"{model_name[:25]}{lora_weight_desc}{trigger_words}"

        if value_type in ["ModelMergeBlocks"]:
            if ":" in value:
                line = value.split(':')
                value_label = f"{line[0]}"
            elif len(value) > 16:
                value_label = f"ModelMergeBlocks {index + 1}"
            else:
                value_label = f"MMB: {value}"

        if value_type in ["Pos Condition"]:
            value_label = f"pos cond {index + 1}" if index>0 else f"pos cond"
        if value_type in ["Neg Condition"]:
            value_label = f"neg cond {index + 1}" if index>0 else f"neg cond"

        if value_type in ["Positive Prompt S/R"]:
            value_label = f"pos prompt {index + 1}" if index>0 else f"pos prompt"
        if value_type in ["Negative Prompt S/R"]:
            value_label = f"neg prompt {index + 1}" if index>0 else f"neg prompt"

        if value_type in ["steps", "cfg", "denoise", "clip_skip",
                          "lora_model_strength", "lora_clip_strength"]:
            value_label = f"{value_type}: {value}"

        if value_type == "positive":
            value_label = f"pos prompt {index + 1}"
        elif value_type == "negative":
            value_label = f"neg prompt {index + 1}"

        return plot_image_vars, value_label

    @staticmethod
    def get_font(font_size, font_path=None):
        if font_path is None:
            font_path = str(Path(os.path.join(RESOURCES_DIR, 'OpenSans-Medium.ttf')))
        return ImageFont.truetype(font_path, font_size)

    @staticmethod
    def update_label(label, value, num_items):
        if len(label) < num_items:
            return [*label, value]
        return label

    @staticmethod
    def rearrange_tensors(latent, num_cols, num_rows):
        new_latent = []
        for i in range(num_rows):
            for j in range(num_cols):
                index = j * num_rows + i
                new_latent.append(latent[index])
        return new_latent

    def calculate_background_dimensions(self):
        border_size = int((self.max_width // 8) * 1.5) if self.y_type != "None" or self.x_type != "None" else 0

        bg_width = self.num_cols * (self.max_width + self.grid_spacing) - self.grid_spacing + border_size * (
                    self.y_type != "None")
        bg_height = self.num_rows * (self.max_height + self.grid_spacing) - self.grid_spacing + border_size * (
                    self.x_type != "None")
        
        # Add space at the bottom of the image for common informaiton about the image
        bg_height = bg_height + (border_size*2)
#        print(f"Grid Size: width = {bg_width} height = {bg_height} border_size = {border_size}")

        x_offset_initial = border_size if self.y_type != "None" else 0
        y_offset = border_size if self.x_type != "None" else 0

        return bg_width, bg_height, x_offset_initial, y_offset


    def adjust_font_size(self, text, initial_font_size, label_width):
        font = self.get_font(initial_font_size, self.custom_font)
        text_width = font.getbbox(text)
#        pprint.pp(f"Initial font size: {initial_font_size}, text: {text}, text_width: {text_width}")
        if text_width and text_width[2]:
            text_width = text_width[2]

        scaling_factor = 0.9
        if text_width > (label_width * scaling_factor):
#            print(f"Adjusting font size from {initial_font_size} to fit text width {text_width} into label width {label_width} scaling_factor {scaling_factor}")
            return int(initial_font_size * (label_width / text_width) * scaling_factor)
        else:
            return initial_font_size

    def textsize(self, d, text, font):
        _, _, width, height = d.textbbox((0, 0), text=text, font=font)
        return width, height

    def create_label(self, img, text, initial_font_size, is_x_label=True, max_font_size=70, min_font_size=10, label_width=0, label_height=0):

        # if the label_width is specified, leave it along.  Otherwise do the old logic.
        if label_width == 0:          
            label_width = img.width if is_x_label else img.height

        text_lines = text.split('\n')
        longest_line = max(text_lines, key=len)
                
        # Adjust font size
        font_size = self.adjust_font_size(longest_line, initial_font_size, label_width)
        font_size = min(max_font_size, font_size)  # Ensure font isn't too large
        font_size = max(min_font_size, font_size)  # Ensure font isn't too small

        if label_height == 0:
            label_height = int(font_size * 1.5) if is_x_label else font_size

        label_bg = Image.new('RGBA', (label_width, label_height), color=(255, 255, 255, 0))
        d = ImageDraw.Draw(label_bg)

        font = self.get_font(font_size, self.custom_font)

        # Check if text will fit, if not insert ellipsis and reduce text
        if self.textsize(d, text, font=font)[0] > label_width:
            while self.textsize(d, text + '...', font=font)[0] > label_width and len(text) > 0:
                text = text[:-1]
            text = text + '...'

        # Compute text width and height for multi-line text
  
        text_widths, text_heights = zip(*[self.textsize(d, line, font=font) for line in text_lines])
        max_text_width = max(text_widths)
        total_text_height = sum(text_heights)

        # Compute position for each line of text
        lines_positions = []
        current_y = 0
        for line, line_width, line_height in zip(text_lines, text_widths, text_heights):
            text_x = (label_width - line_width) // 2
            text_y = current_y + (label_height - total_text_height) // 2
            current_y += line_height
            lines_positions.append((line, (text_x, text_y)))

        # Draw each line of text
        for line, (text_x, text_y) in lines_positions:
            d.text((text_x, text_y), line, fill='black', font=font)

        return label_bg

    def sample_plot_image(self, plot_image_vars, samples, preview_latent, latents_plot, image_list, disable_noise,
                          start_step, last_step, force_full_denoise, x_value=None, y_value=None):
        model, clip, vae, positive, negative, seed, steps, cfg = None, None, None, None, None, None, None, None
        sampler_name, scheduler, denoise = None, None, None

        a1111_prompt_style = plot_image_vars['a1111_prompt_style'] if "a1111_prompt_style" in plot_image_vars else False
        clip = clip if clip is not None else plot_image_vars["clip"]
        steps = plot_image_vars['steps'] if "steps" in plot_image_vars else 1

        sd_version = get_sd_version(plot_image_vars['model'])          
        # 高级用法
        if plot_image_vars["x_node_type"] == "advanced" or plot_image_vars["y_node_type"] == "advanced":
            if self.x_type == "Seeds++ Batch" or self.y_type == "Seeds++ Batch":
                seed = int(x_value) if self.x_type == "Seeds++ Batch" else int(y_value)
            if self.x_type == "Steps" or self.y_type == "Steps":
                steps = int(x_value) if self.x_type == "Steps" else int(y_value)
            if self.x_type == "StartStep" or self.y_type == "StartStep":
                start_step = int(x_value) if self.x_type == "StartStep" else int(y_value)
            if self.x_type == "EndStep" or self.y_type == "EndStep":
                last_step = int(x_value) if self.x_type == "EndStep" else int(y_value)
            if self.x_type == "CFG Scale" or self.y_type == "CFG Scale":
                cfg = float(x_value) if self.x_type == "CFG Scale" else float(y_value)
            if self.x_type == "Sampler" or self.y_type == "Sampler":
                sampler_name = x_value if self.x_type == "Sampler" else y_value
            if self.x_type == "Scheduler" or self.y_type == "Scheduler":
                scheduler = x_value if self.x_type == "Scheduler" else y_value
            if self.x_type == "Sampler&Scheduler" or self.y_type == "Sampler&Scheduler":
                arr = x_value.split(',') if self.x_type == "Sampler&Scheduler" else y_value.split(',')
                if arr[0] and arr[0]!= 'None':
                    sampler_name = arr[0]
                if arr[1] and arr[1]!= 'None':
                    scheduler = arr[1]
            if self.x_type == "Denoise" or self.y_type == "Denoise":
                denoise = float(x_value) if self.x_type == "Denoise" else float(y_value)
            if self.x_type == "Pos Condition" or self.y_type == "Pos Condition":
                positive = plot_image_vars['positive_cond_stack'][int(x_value)] if self.x_type == "Pos Condition" else plot_image_vars['positive_cond_stack'][int(y_value)]
            if self.x_type == "Neg Condition" or self.y_type == "Neg Condition":
                negative = plot_image_vars['negative_cond_stack'][int(x_value)] if self.x_type == "Neg Condition" else plot_image_vars['negative_cond_stack'][int(y_value)]
            # 模型叠加
            if self.x_type == "ModelMergeBlocks" or self.y_type == "ModelMergeBlocks":
                ckpt_name_1, ckpt_name_2 = plot_image_vars['models']
                model1, clip1, vae1, clip_vision = self.easyCache.load_checkpoint(ckpt_name_1)
                model2, clip2, vae2, clip_vision = self.easyCache.load_checkpoint(ckpt_name_2)
                xy_values = x_value if self.x_type == "ModelMergeBlocks" else y_value
                if ":" in xy_values:
                    xy_line = xy_values.split(':')
                    xy_values = xy_line[1]

                xy_arrs = xy_values.split(',')
                # ModelMergeBlocks
                if len(xy_arrs) == 3:
                    input, middle, out = xy_arrs
                    kwargs = {
                        "input": input,
                        "middle": middle,
                        "out": out
                    }
                elif len(xy_arrs) == 30:
                    kwargs = {}
                    kwargs["time_embed."] = xy_arrs[0]
                    kwargs["label_emb."] = xy_arrs[1]

                    for i in range(12):
                        kwargs["input_blocks.{}.".format(i)] = xy_arrs[2+i]

                    for i in range(3):
                        kwargs["middle_block.{}.".format(i)] = xy_arrs[14+i]

                    for i in range(12):
                        kwargs["output_blocks.{}.".format(i)] = xy_arrs[17+i]

                    kwargs["out."] = xy_arrs[29]
                else:
                    raise Exception("ModelMergeBlocks weight length error")
                default_ratio = next(iter(kwargs.values()))

                m = model1.clone()
                kp = model2.get_key_patches("diffusion_model.")

                for k in kp:
                    ratio = float(default_ratio)
                    k_unet = k[len("diffusion_model."):]

                    last_arg_size = 0
                    for arg in kwargs:
                        if k_unet.startswith(arg) and last_arg_size < len(arg):
                            ratio = float(kwargs[arg])
                            last_arg_size = len(arg)

                    m.add_patches({k: kp[k]}, 1.0 - ratio, ratio)

                vae_use = plot_image_vars['vae_use']

                clip = clip2 if vae_use == 'Use Model 2' else clip1
                if vae_use == 'Use Model 2':
                    vae = vae2
                elif vae_use == 'Use Model 1':
                    vae = vae1
                else:
                    vae = self.easyCache.load_vae(vae_use)
                model = m

                # 如果存在lora_stack叠加lora
                optional_lora_stack = plot_image_vars['lora_stack']
                if optional_lora_stack is not None and optional_lora_stack != []:
                    for lora in optional_lora_stack:
                        model, clip = self.easyCache.load_lora(lora)

                # 处理clip
                clip = clip.clone()
                if plot_image_vars['clip_skip'] != 0:
                    clip.clip_layer(plot_image_vars['clip_skip'])

            # CheckPoint
            if self.x_type == "Checkpoint" or self.y_type == "Checkpoint":
                xy_values = x_value if self.x_type == "Checkpoint" else y_value
                ckpt_name, clip_skip, vae_name = xy_values.split(",")
                ckpt_name = ckpt_name.replace('*', ',')
                vae_name = vae_name.replace('*', ',')
                model, clip, vae, clip_vision = self.easyCache.load_checkpoint(ckpt_name)
                if vae_name != 'None':
                    vae = self.easyCache.load_vae(vae_name)

                # 如果存在lora_stack叠加lora
                optional_lora_stack = plot_image_vars['lora_stack']
                if optional_lora_stack is not None and optional_lora_stack != []:
                    for lora in optional_lora_stack:
                        lora['model'] = model
                        lora['clip'] = clip
                        model, clip = self.easyCache.load_lora(lora)

                # 处理clip
                clip = clip.clone()
                if clip_skip != 'None':
                    clip.clip_layer(int(clip_skip))
                    positive = plot_image_vars['positive']
                    negative = plot_image_vars['negative']
                    a1111_prompt_style = plot_image_vars['a1111_prompt_style']
                    steps = plot_image_vars['steps']
                    clip = clip if clip is not None else plot_image_vars["clip"]
                    positive = advanced_encode(clip, positive,
                                               plot_image_vars['positive_token_normalization'],
                                               plot_image_vars['positive_weight_interpretation'],
                                               w_max=1.0,
                                               apply_to_pooled="enable",
                                               a1111_prompt_style=a1111_prompt_style, steps=steps)

                    negative = advanced_encode(clip, negative,
                                               plot_image_vars['negative_token_normalization'],
                                               plot_image_vars['negative_weight_interpretation'],
                                               w_max=1.0,
                                               apply_to_pooled="enable",
                                               a1111_prompt_style=a1111_prompt_style, steps=steps)
                    if "positive_cond" in plot_image_vars:
                        positive = positive + plot_image_vars["positive_cond"]
                    if "negative_cond" in plot_image_vars:
                        negative = negative + plot_image_vars["negative_cond"]

            # Lora
            if self.x_type == "Lora" or self.y_type == "Lora":
#                print(f"Lora: {x_value} {y_value}")
                model = model if model is not None else plot_image_vars["model"]
                clip = clip if clip is not None else plot_image_vars["clip"]
                
                # Build lora_stack from both X and Y axes if both are LoRA types
                lora_stack = []
                
                # Add X axis LoRA if present
                if self.x_type == "Lora":
                    lora_name, lora_model_strength, lora_clip_strength, _ = x_value.split(",")
                    lora_stack.append({"lora_name": lora_name, "model": model, "clip": clip, "model_strength": float(lora_model_strength), "clip_strength": float(lora_clip_strength)})
                
                # Add Y axis LoRA if present
                if self.y_type == "Lora":
                    lora_name, lora_model_strength, lora_clip_strength, _ = y_value.split(",")
                    lora_stack.append({"lora_name": lora_name, "model": model, "clip": clip, "model_strength": float(lora_model_strength), "clip_strength": float(lora_clip_strength)})
                
#                print(f"new_lora_stack: {new_lora_stack}")

                
                if 'lora_stack' in plot_image_vars:
                    lora_stack = lora_stack + plot_image_vars['lora_stack']
                
                if lora_stack is not None and lora_stack != []:
                    for lora in lora_stack:
                        # Each generation of the model, must use the reference to previously created model / clip objects.
                        lora['model'] = model
                        lora['clip'] = clip
                        model, clip = self.easyCache.load_lora(lora)

            # 提示词
            if "Positive" in self.x_type or "Positive" in self.y_type:
                if self.x_type == 'Positive Prompt S/R' or self.y_type == 'Positive Prompt S/R':
                    positive = x_value if self.x_type == "Positive Prompt S/R" else y_value

                if sd_version == 'flux':
                    positive, = CLIPTextEncode().encode(clip, positive)
                else:
                    positive = advanced_encode(clip, positive,
                                                plot_image_vars['positive_token_normalization'],
                                                plot_image_vars['positive_weight_interpretation'],
                                                w_max=1.0,
                                                apply_to_pooled="enable", a1111_prompt_style=a1111_prompt_style, steps=steps)

                # if "positive_cond" in plot_image_vars:
                #     positive = positive + plot_image_vars["positive_cond"]

            if "Negative" in self.x_type or "Negative" in self.y_type:
                if self.x_type == 'Negative Prompt S/R' or self.y_type == 'Negative Prompt S/R':
                    negative = x_value if self.x_type == "Negative Prompt S/R" else y_value

                if sd_version == 'flux':
                    negative, = CLIPTextEncode().encode(clip, negative)
                else:
                    negative = advanced_encode(clip, negative,
                                                plot_image_vars['negative_token_normalization'],
                                                plot_image_vars['negative_weight_interpretation'],
                                                w_max=1.0,
                                                apply_to_pooled="enable", a1111_prompt_style=a1111_prompt_style, steps=steps)
                # if "negative_cond" in plot_image_vars:
                #     negative = negative + plot_image_vars["negative_cond"]

            # ControlNet
            if "ControlNet" in self.x_type or "ControlNet" in self.y_type:
                cnet = plot_image_vars["cnet"] if "cnet" in plot_image_vars else None
                positive = plot_image_vars["positive_cond"] if "positive" in plot_image_vars else None
                negative = plot_image_vars["negative_cond"] if "negative" in plot_image_vars else None
                if cnet:
                    index = x_value if "ControlNet" in self.x_type else y_value
                    controlnet = cnet[index]
                    for index, item in enumerate(controlnet):
                        control_net_name = item[0]
                        image = item[1]
                        strength = item[2]
                        start_percent = item[3]
                        end_percent = item[4]
                        provided_control_net = item[5] if len(item) > 5 else None
                        positive, negative = easyControlnet().apply(control_net_name, image, positive, negative, strength, start_percent, end_percent, provided_control_net, 1)
            # Flux guidance
            if self.x_type == "Flux Guidance" or self.y_type == "Flux Guidance":
                positive = plot_image_vars["positive_cond"] if "positive" in plot_image_vars else None
                flux_guidance = float(x_value) if self.x_type == "Flux Guidance" else float(y_value)
                positive, = FluxGuidance().append(positive, flux_guidance)

        # 简单用法
        if plot_image_vars["x_node_type"] == "loader" or plot_image_vars["y_node_type"] == "loader":
            if self.x_type == 'ckpt_name' or self.y_type == 'ckpt_name':
                ckpt_name = x_value if self.x_type == "ckpt_name" else y_value
                model, clip, vae, clip_vision = self.easyCache.load_checkpoint(ckpt_name)

            if self.x_type == 'lora_name' or self.y_type == 'lora_name':
                model, clip, vae, clip_vision = self.easyCache.load_checkpoint(plot_image_vars['ckpt_name'])
                lora_name = x_value if self.x_type == "lora_name" else y_value
                lora = {"lora_name": lora_name, "model": model, "clip": clip, "model_strength": 1, "clip_strength": 1}
                model, clip = self.easyCache.load_lora(lora)

            if self.x_type == 'lora_model_strength' or self.y_type == 'lora_model_strength':
                model, clip, vae, clip_vision = self.easyCache.load_checkpoint(plot_image_vars['ckpt_name'])
                lora_model_strength = float(x_value) if self.x_type == "lora_model_strength" else float(y_value)
                lora = {"lora_name": plot_image_vars['lora_name'], "model": model, "clip": clip, "model_strength": lora_model_strength, "clip_strength": plot_image_vars['lora_clip_strength']}
                model, clip = self.easyCache.load_lora(lora)

            if self.x_type == 'lora_clip_strength' or self.y_type == 'lora_clip_strength':
                model, clip, vae, clip_vision = self.easyCache.load_checkpoint(plot_image_vars['ckpt_name'])
                lora_clip_strength = float(x_value) if self.x_type == "lora_clip_strength" else float(y_value)
                lora = {"lora_name": plot_image_vars['lora_name'], "model": model, "clip": clip, "model_strength": plot_image_vars['lora_model_strength'], "clip_strength": lora_clip_strength}
                model, clip = self.easyCache.load_lora(lora)

            # Check for custom VAE
            if self.x_type == 'vae_name' or self.y_type == 'vae_name':
                vae_name = x_value if self.x_type == "vae_name" else y_value
                vae = self.easyCache.load_vae(vae_name)

            # CLIP skip
            if not clip:
                raise Exception("No CLIP found")
            clip = clip.clone()
            clip.clip_layer(plot_image_vars['clip_skip'])

            if sd_version == 'flux':
                positive, = CLIPTextEncode().encode(clip, positive)
            else:
                positive = advanced_encode(clip, plot_image_vars['positive'],
                                                            plot_image_vars['positive_token_normalization'],
                                                            plot_image_vars['positive_weight_interpretation'], w_max=1.0,
                                                            apply_to_pooled="enable",a1111_prompt_style=a1111_prompt_style, steps=steps)

            if sd_version == 'flux':
                negative, = CLIPTextEncode().encode(clip, negative)
            else:
                negative = advanced_encode(clip, plot_image_vars['negative'],
                                                            plot_image_vars['negative_token_normalization'],
                                                            plot_image_vars['negative_weight_interpretation'], w_max=1.0,
                                                            apply_to_pooled="enable", a1111_prompt_style=a1111_prompt_style, steps=steps)


        model = model if model is not None else plot_image_vars["model"]
        vae = vae if vae is not None else plot_image_vars["vae"]
        positive = positive if positive is not None else plot_image_vars["positive_cond"]
        negative = negative if negative is not None else plot_image_vars["negative_cond"]

        seed = seed if seed is not None else plot_image_vars["seed"]
        steps = steps if steps is not None else plot_image_vars["steps"]
        cfg = cfg if cfg is not None else plot_image_vars["cfg"]
        sampler_name = sampler_name if sampler_name is not None else plot_image_vars["sampler_name"]
        scheduler = scheduler if scheduler is not None else plot_image_vars["scheduler"]
        denoise = denoise if denoise is not None else plot_image_vars["denoise"]

        noise_device = plot_image_vars["noise_device"] if "noise_device" in plot_image_vars else 'cpu'

        # LayerDiffuse
        layer_diffusion_method = plot_image_vars["layer_diffusion_method"] if "layer_diffusion_method" in plot_image_vars else None
        empty_samples = plot_image_vars["empty_samples"] if "empty_samples" in plot_image_vars else None

        if layer_diffusion_method:
            samp_blend_samples = plot_image_vars["blend_samples"] if "blend_samples" in plot_image_vars else None
            additional_cond = plot_image_vars["layer_diffusion_cond"] if "layer_diffusion_cond" in plot_image_vars else None

            images = plot_image_vars["images"].movedim(-1, 1) if "images" in plot_image_vars else None
            weight = plot_image_vars['layer_diffusion_weight'] if 'layer_diffusion_weight' in plot_image_vars else 1.0
            model, positive, negative = LayerDiffuse().apply_layer_diffusion(model, layer_diffusion_method, weight, samples,
                                                                                  samp_blend_samples, positive,
                                                                                  negative, images, additional_cond)

        samples = empty_samples if layer_diffusion_method is not None and empty_samples is not None else samples
        # Sample
        samples = self.sampler.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, samples,
                                          denoise=denoise, disable_noise=disable_noise, preview_latent=preview_latent,
                                          start_step=start_step, last_step=last_step,
                                          force_full_denoise=force_full_denoise, noise_device=noise_device)

        # Decode images and store
        latent = samples["samples"]

        # Add the latent tensor to the tensors list
        latents_plot.append(latent)

        # Decode the image
        image = vae.decode(latent).cpu()

        if self.output_individuals in [True, "True"]:
            easySave(image, self.save_prefix, self.image_output)

        # Convert the image from tensor to PIL Image and add it to the list
        pil_image = self.sampler.tensor2pil(image)
        image_list.append(pil_image)

        # Update max dimensions
        self.max_width = max(self.max_width, pil_image.width)
        self.max_height = max(self.max_height, pil_image.height)

        # Return the touched variables
        return image_list, self.max_width, self.max_height, latents_plot

    # Process Functions
    def validate_xy_plot(self):
        if self.x_type == 'None' and self.y_type == 'None':
            log_node_warn(f'#{self.my_unique_id}','No Valid Plot Types - Reverting to default sampling...')
            return False
        else:
            return True

    def get_latent(self, samples):
        # Extract the 'samples' tensor from the dictionary
        latent_image_tensor = samples["samples"]

        # Split the tensor into individual image tensors
        image_tensors = torch.split(latent_image_tensor, 1, dim=0)

        # Create a list of dictionaries containing the individual image tensors
        latent_list = [{'samples': image} for image in image_tensors]

        # Set latent only to the first latent of batch
        if self.latent_id >= len(latent_list):
            log_node_warn(f'#{self.my_unique_id}',f'The selected latent_id ({self.latent_id}) is out of range.')
            log_node_warn(f'#{self.my_unique_id}', f'Automatically setting the latent_id to the last image in the list (index: {len(latent_list) - 1}).')

            self.latent_id = len(latent_list) - 1

        return latent_list[self.latent_id]

    def get_labels_and_sample(self, plot_image_vars, latent_image, preview_latent, start_step, last_step,
                              force_full_denoise, disable_noise):
        # Handle X-only variation (Y is "None")
        if self.y_type == 'None':
            for x_index, x_value in enumerate(self.x_values):
                plot_image_vars, x_value_label = self.define_variable(plot_image_vars, self.x_type, x_value, x_index)
                self.x_label = self.update_label(self.x_label, x_value_label, len(self.x_values))
                
                self.image_list, self.max_width, self.max_height, self.latents_plot = self.sample_plot_image(
                    plot_image_vars, latent_image, preview_latent, self.latents_plot, self.image_list,
                    disable_noise, start_step, last_step, force_full_denoise, x_value)
                self.num += 1
        # Handle Y-only variation (X is "None")
        elif self.x_type == 'None':
            for y_index, y_value in enumerate(self.y_values):
                plot_image_vars, y_value_label = self.define_variable(plot_image_vars, self.y_type, y_value, y_index)
                self.y_label = self.update_label(self.y_label, y_value_label, len(self.y_values))
                
                self.image_list, self.max_width, self.max_height, self.latents_plot = self.sample_plot_image(
                    plot_image_vars, latent_image, preview_latent, self.latents_plot, self.image_list,
                    disable_noise, start_step, last_step, force_full_denoise, y_value=y_value)
                self.num += 1
        # Handle both X and Y variation
        else:
            for x_index, x_value in enumerate(self.x_values):
                plot_image_vars, x_value_label = self.define_variable(plot_image_vars, self.x_type, x_value, x_index)
                self.x_label = self.update_label(self.x_label, x_value_label, len(self.x_values))
                
                for y_index, y_value in enumerate(self.y_values):
                    plot_image_vars, y_value_label = self.define_variable(plot_image_vars, self.y_type, y_value, y_index)
                    self.y_label = self.update_label(self.y_label, y_value_label, len(self.y_values))
                    
                    self.image_list, self.max_width, self.max_height, self.latents_plot = self.sample_plot_image(
                        plot_image_vars, latent_image, preview_latent, self.latents_plot, self.image_list,
                        disable_noise, start_step, last_step, force_full_denoise, x_value, y_value)
                    self.num += 1

        # Rearrange latent array to match preview image grid
        self.latents_plot = self.rearrange_tensors(self.latents_plot, self.num_cols, self.num_rows)

        # Concatenate the tensors along the first dimension (dim=0)
        self.latents_plot = torch.cat(self.latents_plot, dim=0)

        return self.latents_plot

    def plot_images_and_labels(self, plot_image_vars):
                    
        bg_width, bg_height, x_offset_initial, y_offset = self.calculate_background_dimensions()

        background = Image.new('RGBA', (int(bg_width), int(bg_height)), color=(255, 255, 255, 255))

        output_image = []
        for row_index in range(self.num_rows):
            x_offset = x_offset_initial

            for col_index in range(self.num_cols):
                index = col_index * self.num_rows + row_index
                img = self.image_list[index]
                output_image.append(self.sampler.pil2tensor(img))
                background.paste(img, (x_offset, y_offset))

                # Handle X label
                if row_index == 0 and self.x_type != "None":
                    label_bg = self.create_label(img, self.x_label[col_index], int(48 * img.width / 512))
                    label_y = (y_offset - label_bg.height) // 2
                    background.alpha_composite(label_bg, (x_offset, label_y))

                # Handle Y label
                if col_index == 0 and self.y_type != "None":
                    label_bg = self.create_label(img, self.y_label[row_index], int(48 * img.height / 512), False)
                    label_bg = label_bg.rotate(90, expand=True)

                    label_x = (x_offset - label_bg.width) // 2
                    label_y = y_offset + (img.height - label_bg.height) // 2
                    background.alpha_composite(label_bg, (label_x, label_y))

                x_offset += img.width + self.grid_spacing

            y_offset += img.height + self.grid_spacing

        # lookup used models in the image
        common_label = ""
        # Update to add a function to do the heavy lifting. Parameters are plot_image_vars name, label to use, names of the axis, 

        # pprint.pp(plot_image_vars)

        # We don't process LORAs here because there can be multiple of them.
        labels = [
            {"id": "ckpt_name", "id_desc": "ckpt", "axis_type" : "Checkpoint"},
            {"id": "vae_name", "id_desc": '', "axis_type" : "vae_name"},
            {"id": "sampler_name", "id_desc": "sampler", "axis_type" : "Sampler"},
            {"id": "scheduler", "id_desc": '', "axis_type" : "Scheduler"},
            {"id": "steps", "id_desc": '', "axis_type" : "Steps"},
            {"id": "Flux Guidance", "id_desc": 'guidance', "axis_type" : "Flux Guidance"},     
            {"id": "seed", "id_desc": '', "axis_type" : "Seeds++ Batch"}
        ]
        
        for item in labels:
            # Only add the label if it's not one of the axis
            # print(f"Checking item: {item['id']} axis_type {item['axis_type']} x_type: {self.x_type} y_type: {self.y_type}")
            if self.x_type != item['axis_type'] and self.y_type != item['axis_type']:
                common_label += self.add_common_label(item['id'], plot_image_vars, item['id_desc'])
        common_label += f"\n"
                
        if plot_image_vars['lora_stack'] is not None and plot_image_vars['lora_stack'] != []:
#            print(f"lora_stack: {plot_image_vars['lora_stack']}")
            for lora in plot_image_vars['lora_stack']:

                lora_name = lora['lora_name']
                lora_weight = lora['model_strength']
                if lora_name is not None and len(lora_name) > 0 and lora_weight > 0:
                    common_label += f"LORA: {lora_name} weight: {lora_weight:.2f} \n"
                    
        common_label = common_label.strip()
        
        if len(common_label) > 0:
            label_height = background.height - y_offset
            label_bg = self.create_label(background, common_label, int(48 * background.width / 512), label_width=background.width, label_height=label_height)
            label_x = (background.width - label_bg.width) // 2
            label_y = y_offset
#            print(f"Adding common label: {common_label} x = {label_x} y = {label_y}")
            background.alpha_composite(label_bg, (label_x, label_y))

        return (self.sampler.pil2tensor(background), output_image)

    def add_common_label(self, tag, plot_image_vars, description = ''):
        label = ''
        if description == '': description = tag
        if tag in plot_image_vars and plot_image_vars[tag] is not None and plot_image_vars[tag] != 'None':
            label += f"{description}: {plot_image_vars[tag]} "
#        print(f"add_common_label: {tag} description: {description} label: {label}" )
        return label