import os
import re
from glob import glob
import torch
from torchvision.transforms import InterpolationMode
import torchvision.transforms.v2.functional as F
from llama_cpp import Llama
from scipy import ndimage
import numpy as np
import node_helpers
import folder_paths
import comfy
from nodes import LoraLoader
import nodes
from comfy_api.latest import io
import latent_preview
from comfy.samplers import sampler_object, KSAMPLER
from comfy_extras.nodes_custom_sampler import Noise_EmptyNoise
import json
import math
from PIL import Image
from io import BytesIO
import base64

class PadAndResize:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'image': ('IMAGE',),
                'latent': ('LATENT',),
                'mode': (['resize', 'crop', 'fit'],),
                'left': ('FLOAT', {'min': 0, 'step': .01, 'default': 0}),
                'top': ('FLOAT', {'min': 0, 'step': .01, 'default': 0}),
                'right': ('FLOAT', {'min': 0, 'step': .01, 'default': 0}),
                'bottom': ('FLOAT', {'min': 0, 'step': .01, 'default': 0}),
                }
            }
    
    RETURN_TYPES = ('IMAGE',)
    
    FUNCTION = 'run'

    CATEGORY = 'davcha'

    def run(self, image, latent, mode, left, top, right, bottom):
        image = image.permute(0,3,1,2)
        image = self.pad(image, left, top, right, bottom)
        image = self.resize(image, latent, mode)
        image = image.permute(0,2,3,1)
        return (image,)
    
    def resize(self, image, latent, mode):
        b, c, sh, sw = image.shape
        *_, th, tw = latent['samples'].shape
        th, tw = th * 8, tw * 8
        if (th, tw) == (sh, sw):
            return image
        if mode == 'resize':
            result = F.resize(image, (th, tw), InterpolationMode.NEAREST_EXACT)
        else:
            fun = max if mode == 'crop' else min
            scale = fun(torch.tensor((th/sh,tw/sw)))
            shape = tuple(int(x * scale) for x in [sh,sw])
            mask = torch.ones((1, sh, sw))
            mask = F.resize(mask, shape, InterpolationMode.NEAREST_EXACT)
            result = F.resize(image, shape, InterpolationMode.NEAREST_EXACT)
            phw = (
                (tw - int(sw*scale))//2,
                (th - int(sh*scale))//2,
                (tw - int(sw*scale)+1)//2,
                (th - int(sh*scale)+1)//2,
            )
            mask = F.pad(mask, phw, fill=0, padding_mode='constant')
            # if mask.sum() != np.prod(mask.shape[-2:]):
            #     dist = ndimage.distance_transform_edt(mask)
            #     soft_m = np.minimum(dist / 32, 1)
            #     soft_m = torch.from_numpy(soft_m).type(torch.float32)
            #     mask = soft_m
            result = F.pad(result, phw, fill=0, padding_mode='constant')
            if c != 4:
                result = torch.cat((result, mask.unsqueeze(1)), 1)
            else:
                result[:,3,:,:] *= mask
        return result

    def pad(self, image, left, top, right, bottom):
        b, c, h, w = image.shape
        left, right = [int(x * w) for x in [left, right]]
        top, bottom = [int(x * h) for x in [top, bottom]]
        mask = torch.ones(b, h, w, dtype=torch.float32)
        mask = F.pad(mask, (left, top, right, bottom), 0, padding_mode='constant')
        # if mask.sum() != np.prod(mask.shape[-2:]):
        #     dist = ndimage.distance_transform_edt(mask)
        #     soft_m = np.minimum(dist / 32, 1)
        #     soft_m = torch.from_numpy(soft_m).type(torch.float32)
        #     mask = soft_m
        image = F.pad(image, (left, top, right, bottom), 0, padding_mode='constant')
        if c != 4:
            image = torch.cat((image, mask.unsqueeze(1)), 1)
        else:
            image[:,3,...] *= mask
        return image
    
class PercentPadding:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'image': ('IMAGE',),
                'left': ('FLOAT', {'min': 0, 'step': .01, 'default': 0}),
                'top': ('FLOAT', {'min': 0, 'step': .01, 'default': 0}),
                'right': ('FLOAT', {'min': 0, 'step': .01, 'default': 0}),
                'bottom': ('FLOAT', {'min': 0, 'step': .01, 'default': 0}),
                }
            }
    
    RETURN_TYPES = ('IMAGE',)

    FUNCTION = 'run'

    CATEGORY = 'davcha'

    def run(self, image, left, top, right, bottom):
        image = image.permute(0,3,1,2)
        b, c, h, w = image.shape
        left, right = [int(x * w) for x in [left, right]]
        top, bottom = [int(x * h) for x in [top, bottom]]
        mask = torch.ones(b, h, w, dtype=torch.float32)
        mask = F.pad(mask, (left, top, right, bottom), 0, padding_mode='constant')
        dist = ndimage.distance_transform_edt(mask)
        soft_m = np.minimum(dist / 32, 1)
        soft_m = torch.from_numpy(soft_m).type(torch.float32)
        mask = soft_m
        image = F.pad(image, (left, top, right, bottom), 0.5, padding_mode='constant')
        if c != 4:
            image = torch.cat((image, mask.unsqueeze(1)), 1)
        else:
            image[:,3,...] *= mask
        image = image.permute(0,2,3,1)
        return (image, )
    
class StringScheduleHelper:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'text': ('STRING', {'multiline': True, 'dynamicPrompts': True}),
                }
            }
    
    RETURN_NAMES = ('base_prompt', 'schedule', 'has_schedule')
    RETURN_TYPES = ('STRING', 'STRING', 'BOOLEAN')

    FUNCTION = 'run'

    CATEGORY = 'davcha'

    def run(self, text):
        base = []
        sched = []
        for a,b,c in re.findall(r"(\d+):(.+)|(.+)", text):
            if a == '':
                base.append(c.strip())
            else:
                sched.append(f'"{a}": "{b.strip()}"')
        base_prompt = '\n'.join(base)
        schedule = ',\n'.join(sched)
        has_schedule = len(sched) > 0
        if not has_schedule:
            schedule = '"0": ""'
        return (base_prompt, schedule, has_schedule)

class SmartMask:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'mask': ('MASK',),
                }
            }
    
    RETURN_TYPES = ('MASK',)

    FUNCTION = 'mask'

    CATEGORY = 'davcha'

    def mask(self, mask):
        if torch.sum(mask) > 0:
            return (mask, )
        else:
            return (1 - mask, )
        
class ResizeCropFit:
    RETURN_TYPES = ('IMAGE',)
    
    FUNCTION = 'run'

    CATEGORY = 'davcha'

    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'image': ('IMAGE',),
                'latent': ('LATENT',),
                'mode': (['resize', 'crop', 'fit'],)
                }
            }
    
    def run(self, image, latent, mode):
        image = image.permute(0,3,1,2)
        b, c, sh, sw = image.shape
        _, _, th, tw = latent['samples'].shape
        th, tw = th * 8, tw * 8
        if (th, tw) == (sh, sw):
            return (image.permute(0,2,3,1), )
        if mode == 'resize':
            result = F.resize(image, (th, tw), InterpolationMode.NEAREST_EXACT)
        else:
            fun = max if mode == 'crop' else min
            scale = fun(torch.tensor((th/sh,tw/sw)))
            shape = tuple(int(x * scale) for x in [sh,sw])
            mask = torch.ones((1, sh, sw))
            mask = F.resize(mask, shape, InterpolationMode.NEAREST_EXACT)
            result = F.resize(image, shape, InterpolationMode.NEAREST_EXACT)
            phw = (
                (tw - int(sw*scale))//2,
                (th - int(sh*scale))//2,
                (tw - int(sw*scale)+1)//2,
                (th - int(sh*scale)+1)//2,
            )
            mask = F.pad(mask, phw, fill=0, padding_mode='constant')
            dist = ndimage.distance_transform_edt(mask)
            soft_m = np.minimum(dist / 32, 1)
            soft_m = torch.from_numpy(soft_m).type(torch.float32)
            mask = soft_m
            result = F.pad(result, phw, fill=0.5, padding_mode='constant')
            if c != 4:
                result = torch.cat((result, mask.unsqueeze(1)), 1)
            #     result[:,3,:,:] = 1
            else:
                result[:,3,:,:] *= mask
        result = result.permute(0,2,3,1)
        return (result, )

class SoftErosion:
    RETURN_TYPES = ('MASK',)
    FUNCTION = 'run'
    CATEGORY = 'davcha'

    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'mask': ('MASK',),
                'size': ('INT', {'min': 0, 'default': 5}),
            }
        }
    
    def run(self, mask, size):
        if size == 0:
            return (mask, )
        soft_mask = []
        for m in mask:
            dist = ndimage.distance_transform_edt(m)
            soft_m = np.minimum(dist / size, 1)
            soft_m = torch.from_numpy(soft_m)
            soft_mask.append(soft_m)
        soft_mask = torch.from_numpy(np.array(soft_mask)).type(torch.FloatTensor)
        return (soft_mask, )
    
class ApplyMask:
    RETURN_TYPES = ('IMAGE',)
    FUNCTION = 'run'
    CATEGORY = 'davcha'

    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'image': ('IMAGE',),
                'mask': ('MASK',),
            }
        }
    
    def run(self, image, mask):
        image = image.permute(0,3,1,2)
        _, _, h, w = image.shape
        mask = F.resize(mask, (h, w), InterpolationMode.BICUBIC)
        image = image.permute(0,2,3,1)[:,:,:,:3]
        image = image * mask[...,None]
        return (image, )
    
class DStack:
    RETURN_TYPES = ('IMAGE',)
    FUNCTION = 'run'
    CATEGORY = 'davcha'

    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'image1': ('IMAGE',),
                'image2': ('IMAGE',),
            }
        }
    
    def run(self, image1, image2):
        stack = torch.dstack((image1, image2))
        return (stack, )

ratio_merge = ("FLOAT", {"default": 1.0, "min": -5.0, "max": 6.0, "step": 0.001})

class DavchaConditioningConcat:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "conditioning_to": ("CONDITIONING",),
            "conditioning_from": ("CONDITIONING",),
            }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "concat"

    CATEGORY = "conditioning"

    def concat(self, conditioning_to, conditioning_from):
        out = []
        
        if len(conditioning_from) < 1: return (conditioning_to,)
        if len(conditioning_to) < 1: return (conditioning_from,)

        cond_from = conditioning_from[0][0].to(conditioning_to[0][0].device)

        for i in range(len(conditioning_to)):
            t1 = conditioning_to[i][0]
            mb = max(t1.shape[0], cond_from.shape[0])
            tw = torch.cat((t1.repeat(mb // t1.shape[0],1,1), cond_from.repeat(mb // cond_from.shape[0],1,1)), 1)
            n = [tw, conditioning_to[i][1].copy()]
            out.append(n)

        return (out, )

class DavchaModelMergeSimple:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "model1": ("MODEL",),
                              "model2": ("MODEL",),
                              "ratio": ratio_merge,
                              }}
    RETURN_TYPES = ("MODEL",)
    FUNCTION = "merge"

    CATEGORY = "advanced/model_merging"

    def merge(self, model1, model2, ratio):
        m = model1.clone()
        kp = model2.get_key_patches("diffusion_model.")
        for k in kp:
            m.add_patches({k: kp[k]}, 1.0 - ratio, ratio)
        return (m, )
    
class DavchaCLIPMergeSimple:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "clip1": ("CLIP",),
                              "clip2": ("CLIP",),
                              "ratio": ratio_merge,
                              }}
    RETURN_TYPES = ("CLIP",)
    FUNCTION = "merge"

    CATEGORY = "advanced/model_merging"

    def merge(self, clip1, clip2, ratio):
        m = clip1.clone()
        kp = clip2.get_key_patches()
        for k in kp:
            if k.endswith(".position_ids") or k.endswith(".logit_scale"):
                continue
            m.add_patches({k: kp[k]}, 1.0 - ratio, ratio)
        return (m, )

import comfy_extras.nodes_model_merging

class DavchaModelMergeSD1(comfy_extras.nodes_model_merging.ModelMergeBlocks):
    CATEGORY = "advanced/model_merging/model_specific"
    @classmethod
    def INPUT_TYPES(s):
        arg_dict = { "model1": ("MODEL",),
                              "model2": ("MODEL",)}

        argument = ratio_merge

        arg_dict["time_embed."] = argument
        arg_dict["label_emb."] = argument

        for i in range(12):
            arg_dict["input_blocks.{}.".format(i)] = argument

        for i in range(3):
            arg_dict["middle_block.{}.".format(i)] = argument

        for i in range(12):
            arg_dict["output_blocks.{}.".format(i)] = argument

        arg_dict["out."] = argument

        return {"required": arg_dict}


class DavchaModelMergeSDXL(comfy_extras.nodes_model_merging.ModelMergeBlocks):
    CATEGORY = "advanced/model_merging/model_specific"

    @classmethod
    def INPUT_TYPES(s):
        arg_dict = { "model1": ("MODEL",),
                              "model2": ("MODEL",)}

        argument = ratio_merge

        arg_dict["time_embed."] = argument
        arg_dict["label_emb."] = argument

        for i in range(9):
            arg_dict["input_blocks.{}".format(i)] = argument

        for i in range(3):
            arg_dict["middle_block.{}".format(i)] = argument

        for i in range(9):
            arg_dict["output_blocks.{}".format(i)] = argument

        arg_dict["out."] = argument

        return {"required": arg_dict}
    
class ConditioningCompress:
    CATEGORY = "davcha"

    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "conditioning": ("CONDITIONING",),
            "max_tokens": ("INT", {"default": 75, "min": 1, "step": 1})
            }}
    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "run"
    
    
    def run(self, conditioning, max_tokens):
        def reduc(o):
            u,s,v = torch.svd(o)
            return (u[0,:max_tokens,:max_tokens]@s[0,:max_tokens].diag()@v[0,:,:max_tokens].T)[None]
        
        out = [[c[0].clone(), c[1]] for c in conditioning]
        for o in out:
            o[0] = reduc(o[0])
        return (out, )
    
################ DavchaCLIPTextEncode

import regex
from itertools import product

def parse(text):
    def _inner(text, start=0, end=1):
        matches = regex.findall(r'(\[((?:(?:embedding:)?[^\[\]])*?):((?:(?:embedding:)?[^\[\]])*?):(\d*\.?\d+)\])', text, regex.DOTALL)
        if len(matches) == 0:
            return [(text, (start, end))]
        x = [m[1:3] for m in matches]
        prod = list(product(*x))
        results = {}
        for p in prod:
            if p not in results: results[p] = []
            for f, b, a, t in matches:
                t = float(t)
                if b in p:
                    results[p].append((f, b, (start, t)))
                if a in p:
                    results[p].append((f, a, (t, end)))
        final = []
        for items in results.values():
            txt = text
            st, en = start, end
            for f, r, (s, e) in items:
                txt = txt.replace(f, r)
                st = max(st, s)
                en = min(en, e)
            if st <= en:
                for prompt, (a, b) in _inner(txt, st, en):
                    final.append((prompt, (a, b)))
        return final
    candidates = _inner(text)
    dic = {}
    for prompt, (s, e) in candidates:
        if prompt not in dic:
            dic[prompt] = (s, e)
        else:
            dic[prompt] = (min(s, dic[prompt][0]), max(e, dic[prompt][1]))
    return [(prompt, (s,e)) for prompt, (s, e) in dic.items()]

class DavchaCLIPTextEncode:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "text": ("STRING", {"multiline": True, "dynamicPrompts": True, "tooltip": "The text to be encoded."}), 
                "clip": ("CLIP", {"tooltip": "The CLIP model used for encoding the text."})
            }
        }
    RETURN_TYPES = ("CONDITIONING",)
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "run"

    CATEGORY = "davcha"
    
    def run(self, clip, text):
        if text.strip() == "":
            tokens = clip.tokenize("")
            output = clip.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
            cond = output.pop("cond")
            return ([[[cond, output]]], )
        text = text.split('||')
        result = []
        for txt in text:
            texts = re.split(r"\b(AREA\(\s*(?:\d*\.?\d+\s*){4,5}\))", txt)
            
            areas = []
            current_area = None

            for item in texts:
                if item.startswith('AREA'):
                    current_area = item
                else:
                    prompt = parse(item)
                    if current_area is None:
                        areas.append((prompt, None))
                    else:
                        values = re.findall(r"(\d*\.?\d+)", current_area)
                        x, y, w, h, s = [float(i) for i in values + [1.0] * (5 - len(values))]
                        areas.append((prompt, (x, y, w, h, s)))
                    current_area = None
            
            for schedule, area in areas:
                for prompt, (start, end) in schedule:
                    print(f'{area} {start}-{end}: {prompt}')

            cs = []
            for prompts, area in areas:
                for txt, (lb, ub) in prompts:
                    if len(txt) > 0:
                        tokens = clip.tokenize(txt)
                        output = clip.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
                        cond = output.pop("cond")
                        c = [[cond, output]]
                        c = node_helpers.conditioning_set_values(c, {"start_percent": lb,
                                                                    "end_percent": ub})
                        if area is not None:
                            x, y, w, h, s = area
                            c = node_helpers.conditioning_set_values(c, {"area": ("percentage", h, w, y, x),
                                                                        "strength": s,
                                                                        "set_area_to_bounds": False})
                        cs += c
            result.append(cs)
        return (result, )

from pathlib import Path
from webp import WebPData, WebPAnimDecoderOptions, WebPAnimDecoder, mimread

def video_reader(file, start=0, count=0, target_fps=0):
    if Path(file).suffix.lower() == '.webp':
        return _webp_reader(file, start, count, target_fps)
    else:
        return _ffmpeg_reader(file, start, count, target_fps)
    
def _webp_reader(file, start, count, target_fps):
    with open(file, 'rb') as f:
        webp_data = WebPData.from_buffer(f.read())
        dec_opts = WebPAnimDecoderOptions.new(use_threads=True)
        dec = WebPAnimDecoder.new(webp_data, dec_opts)
        eps = 1e-7
        
        frames_data = list(dec.frames())

    frames = [arr[:,:,:3] for arr, _ in frames_data]
    fps = 1000 * len(frames_data) / frames_data[-1][1]
    if target_fps > 0:
        if fps > target_fps:
            frames = [frame[:,:,:3] for frame in mimread(file, fps=target_fps)]
    else:
        target_fps = fps
    
    if count > 0:
        frames = frames[int(start):int(start+count)]
    else:
        frames = frames[int(start):]
    
    frames = torch.from_numpy(np.array(frames)).float() / 255.0

    return frames, target_fps

import cv2
import numpy as np

def _ffmpeg_reader(file, start, count, target_fps):
    cap = cv2.VideoCapture(file)
    
    # Get video properties
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    if target_fps == 0:
        target_fps = fps
    start_frame = int(min(max(0, start), frame_count - 1) * fps / target_fps)
    if count == 0:
        count = int(frame_count * target_fps / fps)
    count = min(int(frame_count * target_fps / fps - start), count)
    
    frames = np.empty((count, frame_height, frame_width, 3), dtype=np.uint8)
    for i in range(count):
        cap.set(cv2.CAP_PROP_POS_FRAMES, int(start_frame + i * fps / target_fps))
        ret, frame = cap.read()
        if not ret:
            break
    
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames[i] = frame_rgb
    
    cap.release()
    
    frames = torch.from_numpy(frames).float() / 255.0
    
    return frames, target_fps

class DavchaLoadVideo:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "path": ("STRING", {"multiline": False}),
                "start": ("INT", {"default": 0, "min": 0, "step": 1}),
                "count": ("INT", {"default": 0, "min": 0, "step": 1}),
                "target_fps": ("FLOAT", {"default": 0, "min": 0, "step": 0.01}),
            }
        }
    RETURN_TYPES = ("IMAGE", "INT", "FLOAT")
    RETURN_NAMES = ("frames", "frame_count", "fps")
    FUNCTION = "run"

    CATEGORY = "davcha"
    
    def run(self, path, start, count, target_fps):
        fullpath = os.path.join(folder_paths.get_input_directory(), path)
        frames, fps = video_reader(fullpath, start, count, target_fps)
        return (frames, len(frames), fps)
    
MAX_RESOLUTION = 16384

class AnyType(str):
    def __eq__(self, _) -> bool:
        return True
    def __ne__(self, __value: object) -> bool:
        return False
    
class DavchaEmptyLatentImage:
    def __init__(self):
        self.device = comfy.model_management.intermediate_device()

    @classmethod
    def INPUT_TYPES(s):        
        return {
            "required": { 
                "width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The width of the latent images in pixels."}),
                "height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The height of the latent images in pixels."}),
                "option": ([AnyType("*")], ),
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "tooltip": "The number of latent images in the batch."})
            }
        }
    RETURN_TYPES = ("LATENT", "FLOAT", "INT")
    RETURN_NAMES = ("latent", "upscale_factor", "batch_size")
    OUTPUT_TOOLTIPS = ("The empty latent image batch.",)
    FUNCTION = "generate"

    CATEGORY = "davcha"
    DESCRIPTION = "Create a new batch of empty latent images to be denoised via sampling."

    def generate(self, width, height, option, batch_size=1):
        w, h, _ = re.findall(r"(\d+)x(\d+): (\d+\.?\d*)", option)[0]
        w, h = int(w), int(h)
        upscale_factor = width / w
        latent = torch.zeros([batch_size, 4, h // 8, w // 8], device=self.device)
        return ({"samples":latent}, upscale_factor, batch_size)

class DavchaMaskImage:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "image": ("IMAGE",),
                              "mask": ("MASK",),
                              }}
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "run"

    CATEGORY = "davcha"

    def run(self, image, mask):
        b, w, h, c = image.shape
        mask = F.resize(mask, (w, h), InterpolationMode.BICUBIC)
        image = image.permute(0, 3, 1, 2)
        if c != 4:
            image = torch.cat((image, torch.ones((b, 1, w, h), dtype=torch.float32)), 1)# mask.unsqueeze(1)), 1)
        image = image * mask.unsqueeze(1)
        image = image.permute(0, 2, 3, 1)
        return (image, )

class DavchaPop:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "items": (AnyType("*"),),
                              }}
    INPUT_IS_LIST = (True, )
    OUTPUT_IS_LIST = (False, True)
    RETURN_TYPES = (AnyType("*"), AnyType("*"))
    FUNCTION = "run"

    CATEGORY = "davcha"

    def run(self, items):
        return (items[0], items[1:])

gguf_folder = os.path.join(folder_paths.models_dir, "llm_gguf")
if os.path.isdir(gguf_folder):
    gguf_files = [file for file in os.listdir(gguf_folder) if file.endswith('.gguf')]
else:
    gguf_files = []
    
class OldDavchaLoadLLM:
    @classmethod
    def INPUT_TYPES(s):
        return {'required':{
            'modelname': (gguf_files,),
        }}
    RETURN_NAMES = ('model',)
    RETURN_TYPES = ('DavchaLLModel',)
    FUNCTION = "run"

    CATEGORY = "davcha"
    
    def run(self, modelname):
        model_path = os.path.join(gguf_folder, modelname)
        model = Llama(
            model_path=model_path,
            n_gpu_layers=-1,
            verbose=False,
            n_ctx=2048,
        )
        return (model,)        

class DavchaLLMAdvanced:
    @classmethod
    def INPUT_TYPES(s):
        return {'required':{
            'model': ('DavchaLLModel',),
            'seed': ('INT', {'default': 0, 'min': 0, 'max': 0xffffffffffffffff}),
            'system': ('STRING', {'multiline': True, 'dynamicPrompts': True}),
            'text': ('STRING', {'multiline': True, 'dynamicPrompts': True}),
            'max_tokens': ('INT', {'default': 512, 'min': 1, 'max': 8192}),
            'temperature': ('FLOAT', {'default': 1.0, 'min': 0, 'max': 1.0, 'step': 0.1}),
            'top_p': ('FLOAT', {'default': 0.9, 'min': 0, 'max': 1.0, 'step': 0.1}),
            'top_k': ('INT', {'default': 50, 'min': 0, 'max': 100}),
            'repeat_penalty': ('FLOAT', {'default': 1.2, 'min': 0, 'max': 5.0, 'step': 0.1}),
        }}
    
    RETURN_NAMES = ('text',)
    RETURN_TYPES = ('STRING',)
    FUNCTION = "run"

    CATEGORY = "davcha"
    
    def run(self, model, seed, system, text, max_tokens, temperature, top_p, top_k, repeat_penalty):
        generate_kwargs = {
            'max_tokens': max_tokens,
            'temperature': temperature,
            'top_p': top_p,
            'top_k': top_k,
            'repeat_penalty': repeat_penalty
        }
        msgs = []
        if len(system.strip()) > 0:
            msgs += [{'role': 'system', 'content': system}]
        msgs += [{"role": "user", "content": text}]
        model.set_seed(seed)
        llm_result = model.create_chat_completion(msgs, **generate_kwargs)
        return (llm_result['choices'][0]['message']['content'].strip(),)

class OldDavchaLLM:
    @classmethod
    def INPUT_TYPES(s):
        return {'required':{
            'model': ('DavchaLLModel',),
            'seed': ('INT', {'default': 0, 'min': 0, 'max': 0xffffffffffffffff}),
            'text': ('STRING', {'multiline': True, 'dynamicPrompts': True}),
        }}
    
    RETURN_NAMES = ('text',)
    RETURN_TYPES = ('STRING',)
    FUNCTION = "run"

    CATEGORY = "davcha"
    
    def run(self, model, seed, text):
        generate_kwargs = {
            'max_tokens': 512,
            'temperature': 1.0,
            'top_p': 0.9,
            'top_k': 50,
            'repeat_penalty': 1.2
        }
        msgs = []
        msgs += [{"role": "user", "content": text}]
        model.set_seed(seed)
        llm_result = model.create_chat_completion(msgs, **generate_kwargs)
        return (llm_result['choices'][0]['message']['content'].strip(),)

def viterbi_distance(a, b):
    """Compute the Viterbi distance between two sequences."""
    n = len(a)
    m = len(b)
    dp = [[0] * (m + 1) for _ in range(n + 1)]
    for i in range(n + 1):
        dp[i][0] = i
    for j in range(m + 1):
        dp[0][j] = j
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            if a[i - 1] == b[j - 1]:
                dp[i][j] = dp[i - 1][j - 1]
            else:
                dp[i][j] = min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1]) + 1
    return dp[n][m]

def viterbi_diff(a, b):
    """Compute the Viterbi diff between two sequences."""
    n = len(a)
    m = len(b)
    dp = [[0] * (m + 1) for _ in range(n + 1)]
    for i in range(n + 1):
        dp[i][0] = i
    for j in range(m + 1):
        dp[0][j] = j
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            if a[i - 1] == b[j - 1]:
                dp[i][j] = dp[i - 1][j - 1]
            else:
                dp[i][j] = min(dp[i - 1][j], dp[i][j - 1], dp[i - 1][j - 1]) + 1
    i, j = n, m
    diff = []
    while i > 0 or j > 0:
        if i > 0 and j > 0 and a[i - 1] == b[j - 1]:
            diff.append((' ', a[i - 1], b[j - 1]))
            i -= 1
            j -= 1
        elif j > 0 and (i == 0 or dp[i][j] == dp[i][j - 1] + 1):
            diff.append(('+', '', b[j - 1]))
            j -= 1
        elif i > 0 and (j == 0 or dp[i][j] == dp[i - 1][j] + 1):
            diff.append(('-', a[i - 1], ''))
            i -= 1
        else:
            diff.append(('*', a[i - 1], b[j - 1]))
            i -= 1
            j -= 1
    diff.reverse()
    return diff

def get_highlow(lorapath, model):
    if not model.endswith('.safetensors'):
        model = f'{model}.safetensors'
    x = glob(os.path.join(lorapath, '**', model), recursive=True)[0]
    x = os.path.relpath(x, lorapath)
    x = os.path.normpath(x)
    files = [os.path.relpath(f, lorapath) for f in glob(os.path.join(lorapath, os.path.dirname(x), '*.safetensors'))]
    files.sort(key=lambda f: viterbi_distance(x, f))
    other = files[1]
    diff = viterbi_diff(x,other)
    x_tag = ''.join(d[1] for d in diff if d[0]=='*').lower()
    high, low = (x,other) if 'h' in x_tag else (other,x)
    return (high, low)
    
class DavchaWan22LoraTagLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {'required':{
            'high': ('MODEL',),
            'low': ('MODEL',),
            'txt': ('STRING', {'multiline': True, 'dynamicPrompts': True}),
        }}
    RETURN_NAMES = ('high','low', 'txt')
    RETURN_TYPES = ('MODEL','MODEL', 'STRING')
    FUNCTION = "run"
    CATEGORY = "davcha"
        
    def run(self, high, low, txt):
        loraspath = folder_paths.get_folder_paths('loras')[0]
        m = re.findall(r'<lora:([^:]+):([^>]+)>', txt)
        for model, weight in m:
            lora_high, lora_low = get_highlow(loraspath, model)
            high, _ = LoraLoader().load_lora(high, None, lora_high, float(weight), float(weight))
            low, _ = LoraLoader().load_lora(low, None, lora_low, float(weight), float(weight))
        txt = re.sub(r'<lora:[^:]+:[^>]+>', '', txt)
        return (high, low, txt)

class DavchaWan22LoraTagParser:
    @classmethod
    def INPUT_TYPES(s):
        return {'required':{
            'txt': ('STRING', {'multiline': True, 'dynamicPrompts': True}),
        }}
    RETURN_NAMES = ('high','low', 'txt')
    RETURN_TYPES = ('STRING','STRING', 'STRING')
    FUNCTION = "run"
    CATEGORY = "davcha"
        
    def run(self, txt):
        loraspath = folder_paths.get_folder_paths('loras')[0]
        m = re.findall(r'<lora:([^:]+):([^>]+)>', txt)
        highs, lows = [], []
        for model, weight in m:
            lora_high, lora_low = get_highlow(loraspath, model)
            highs.append(f'<lora:{lora_high}:{weight}>')
            lows.append(f'<lora:{lora_low}:{weight}>')
        highs = ''.join(highs)
        lows = ''.join(lows)
        txt = re.sub(r'<lora:[^:]+:[^>]+>', '', txt)
        return (highs, lows, txt)

class DavchaTextEncodeQwenImageEditPlus(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="DavchaTextEncodeQwenImageEditPlus",
            category="advanced/conditioning",
            inputs=[
                io.Clip.Input("clip"),
                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
                io.Vae.Input("vae", optional=True),
                io.Image.Input("image1", optional=True),
                io.Image.Input("image2", optional=True),
                io.Image.Input("image3", optional=True),
            ],
            outputs=[
                io.Conditioning.Output(),
            ],
        )

    @classmethod
    def execute(cls, clip, prompt, vae=None, image1=None, image2=None, image3=None) -> io.NodeOutput:
        ref_latents = []
        images = [image1, image2, image3]
        images_vl = []
        llama_template = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
        image_prompt = ""

        for i, image in enumerate(images):
            if image is not None:
                samples = image.movedim(-1, 1)
                total = int(384 * 384)

                scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
                width = round(samples.shape[3] * scale_by)
                height = round(samples.shape[2] * scale_by)

                s = comfy.utils.common_upscale(samples, width, height, "area", "disabled")
                images_vl.append(s.movedim(1, -1))
                if vae is not None:
                    ref_latents.append(vae.encode(samples.movedim(1, -1)[:, :, :, :3]))

                image_prompt += "Picture {}: <|vision_start|><|image_pad|><|vision_end|>".format(i + 1)

        tokens = clip.tokenize(image_prompt + prompt, images=images_vl, llama_template=llama_template)
        conditioning = clip.encode_from_tokens_scheduled(tokens)
        if len(ref_latents) > 0:
            conditioning = node_helpers.conditioning_set_values(conditioning, {"reference_latents": ref_latents}, append=True)
        return io.NodeOutput(conditioning)

class DavchaScheduledTextEncoderQwenImageEditPlus(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="DavchaScheduledTextEncoderQwenImageEditPlus",
            category="advanced/conditioning",
            inputs=[
                io.Clip.Input("clip"),
                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
                io.Vae.Input("vae", optional=True),
                io.Image.Input("image1", optional=True),
                io.Image.Input("image2", optional=True),
                io.Image.Input("image3", optional=True),
            ],
            outputs=[
                io.Conditioning.Output(),
            ],
        )
    
    @classmethod
    def execute(cls, clip, prompt, vae=None, image1=None, image2=None, image3=None) -> io.NodeOutput:
        images = [x for x in [image1, image2, image3] if x is not None]
        
        # 1. PARSE PROMPT & SCHEDULES
        # ===========================
        llama_template = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
        
        # Regex Explanation:
        # img(\d+)       -> Matches "img1", "img2"
        # ( ... )        -> Group 2: Captures the parameters string
        # (?:\s+ ... )*  -> Matches sequence of parameters separated by spaces
        # Inner Choices:
        #   novl|noref                     -> Flags
        #   \d*\.?\d+-\d*\.?\d+    -> Time Range (e.g. 0.5-1.0)
        #   \d*\.?\d+(?:,\d*\.?\d+)? -> Strength (e.g. 0.8 or 0.8,0.5)
        re_block = r'img(\d+)((?:\s+(?:novl|noref|\d*\.?\d+-\d*\.?\d+|\d*\.?\d+(?:,\d*\.?\d+)?))*)'
        
        image_schedules_raw = regex.findall(re_block, prompt)
        text_prompt_base = regex.sub(re_block, '', prompt).strip()

        image_schedules = []
        
        for i_str, params_str in image_schedules_raw:
            idx = int(i_str) - 1 # Convert to 0-based index
            
            # CRITICAL: Ignore schedule if image input is missing
            if idx < 0 or idx >= len(images) or images[idx] is None:
                continue

            # Defaults
            start, end = 0.0, 1.0
            ms, us = 1.0, 1.0
            novl = False
            noref = False

            # Parse parameters (Flexible Order)
            tokens = params_str.strip().split()
            for token in tokens:
                if token == 'novl':
                    novl = True
                elif token == 'noref':
                    noref = True
                elif '-' in token:
                    # Time Range: 0.5-1.0
                    try:
                        s_str, e_str = token.split('-')
                        start = float(s_str)
                        end = float(e_str)
                    except ValueError: pass
                else:
                    # Strength: 0.8 or 0.8,0.5
                    try:
                        if ',' in token:
                            ms_str, us_str = token.split(',')
                            ms = float(ms_str)
                            us = float(us_str)
                        else:
                            ms = float(token)
                            us = float(token)
                    except ValueError: pass

            image_schedules.append({
                "index": idx,
                "start": start,
                "end": end,
                "ms": ms,
                "us": us,
                "noref": noref,
                "novl": novl,
            })

        # Add default schedules for connected images that weren't mentioned in prompt
        scheduled_indices = {s['index'] for s in image_schedules}
        for i, img in enumerate(images):
            if img is not None and i not in scheduled_indices:
                image_schedules.append({
                    "index": i, "start": 0.0, "end": 1.0, "ms": 1.0, "us": 1.0,
                    "noref": False, "novl": False,
                })

        prompt_schedules = parse(text_prompt_base)

        # 2. PRE-PROCESS IMAGES
        # =====================
        preprocessed_images = [None] * len(images)
        for i, image in enumerate(images):
            if image is None: 
                continue

            # Handle Mask/RGB
            if image.shape[-1] == 4:
                mask = image[:,:,:,3:4].permute(0, 3, 1, 2)
                image_rgb = image[:,:,:,:3]
            else:
                mask = torch.zeros((1, 1, image.shape[1], image.shape[2]), device=image.device)
                image_rgb = image
            
            # Prepare Vision Tokens (Resize)
            samples = image_rgb.permute(0, 3, 1, 2)
            total = int(384 * 384)
            scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
            width = round(samples.shape[3] * scale_by)
            height = round(samples.shape[2] * scale_by)
            s_resized = comfy.utils.common_upscale(samples, width, height, "area", "disabled")
            
            entry = {"image_vl": s_resized.permute(0, 2, 3, 1)}
            
            # Prepare Latents (VAE)
            if vae is not None:
                lat = vae.encode(image_rgb)
                *_, H, W = lat.shape
                m = torch.nn.functional.interpolate(mask, size=(H, W), mode='bilinear', align_corners=False)
                entry["base_latent"] = lat
                entry["latent_mask"] = m
            
            preprocessed_images[i] = entry

        # 3. CALCULATE INTERVALS
        # ======================
        timestamps = {0.0, 1.0}
        for _, (s, e) in prompt_schedules:
            timestamps.add(s); timestamps.add(e)
        for sched in image_schedules:
            timestamps.add(sched["start"]); timestamps.add(sched["end"])
            
        sorted_timestamps = sorted(list(timestamps))
        intervals = []
        for i in range(len(sorted_timestamps) - 1):
            if sorted_timestamps[i] < sorted_timestamps[i+1]:
                intervals.append((sorted_timestamps[i], sorted_timestamps[i+1]))

        # 4. BUILD CONDITIONING
        # =====================
        final_conditioning = []
        
        for start, end in intervals:
            midpoint = (start + end) / 2.0
            
            # Identify active text
            active_prompt_text = ""
            for txt, (p_start, p_end) in prompt_schedules:
                if p_start <= midpoint < p_end:
                    active_prompt_text = txt
                    break
            
            # Identify active images
            active_images_data = [
                s for s in image_schedules if s["start"] <= midpoint < s["end"]
            ]
            
            # Optimization: Skip empty intervals
            if not active_prompt_text.strip() and not active_images_data:
                continue

            active_images_data.sort(key=lambda x: x['index'])

            image_placeholders = []
            interval_images_vl = []
            interval_ref_latents = []

            for data in active_images_data:
                idx = data['index']
                entry = preprocessed_images[idx]
                if entry is None: continue # Safety check
                
                # Vision Token Assembly
                if not data['novl']:
                    image_placeholders.append(f"Picture {idx + 1}: <|vision_start|><|image_pad|><|vision_end|>")
                    interval_images_vl.append(entry["image_vl"])
                
                # Reference Latent Assembly
                if not data['noref'] and "base_latent" in entry:
                    lat = entry["base_latent"]
                    m = entry["latent_mask"]
                    # Apply Mask Strengths
                    scaled_lat = lat * (1 - m) * data["ms"] + lat * m * data["us"]
                    interval_ref_latents.append(scaled_lat)

            # Combine Prompts
            full_prompt_parts = []
            if image_placeholders:
                full_prompt_parts.append(" ".join(image_placeholders))
            if active_prompt_text:
                full_prompt_parts.append(active_prompt_text.strip())
            
            full_prompt = " ".join(full_prompt_parts)
            
            # Encode
            tokens = clip.tokenize(full_prompt, images=interval_images_vl, llama_template=llama_template)
            conditioning_chunk = clip.encode_from_tokens_scheduled(tokens)
            
            # Apply Reference Latents
            if interval_ref_latents:
                conditioning_chunk = node_helpers.conditioning_set_values(
                    conditioning_chunk, {"reference_latents": interval_ref_latents}, append=True
                )
            
            # Set Time Bounds
            conditioning_chunk = node_helpers.conditioning_set_values(
                conditioning_chunk, {"start_percent": start, "end_percent": end}
            )
            
            final_conditioning.extend(conditioning_chunk)

        return io.NodeOutput(final_conditioning)

class DavchaQwenImageEditLoraTagLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {'required':{
            'model': ('MODEL',),
            'mode': (['sequential', 'stack'],),
            'txt': ('STRING', {'multiline': True, 'dynamicPrompts': True}),
        }}
        
    RETURN_NAMES = ('model', 'txt')
    RETURN_TYPES = ('MODEL', 'STRING')
    FUNCTION = "run"
    CATEGORY = "davcha"

    def run(self, model, mode, txt):
        loraspath = folder_paths.get_folder_paths('loras')[0]
        
        if mode == 'sequential':
            NunchakuQwenImageLoraLoader = nodes.NODE_CLASS_MAPPINGS['NunchakuQwenImageLoraLoader']
            loader = NunchakuQwenImageLoraLoader()
        else: # stack
            NunchakuQwenImageLoraStack = nodes.NODE_CLASS_MAPPINGS['NunchakuQwenImageLoraStackV2']
            loader = NunchakuQwenImageLoraStack()

        ms = re.findall(r'<lora:([^:]+):([^>]+)>', txt)
        rtxt = re.sub(r'<lora:[^:]+:[^>]+>', '', txt)

        mod = model.clone()
        kwargs = {}
        for i, (lora, strength) in enumerate(ms):
            strength = float(strength)
            if not lora.endswith('.safetensors'):
                lora = f'{lora}.safetensors'
            x = glob(os.path.join(loraspath, '**', lora), recursive=True)[0]
            x = os.path.relpath(x, loraspath)
            x = os.path.normpath(x)
            if mode == 'sequential':
                mod,*_ = loader.load_lora(mod, x, strength)
            else: # stack
                kwargs.update({f'lora_name_{i+1}': x, f'lora_strength_{i+1}': strength})
        if mode == 'stack':
            mod, *_ = loader.load_lora_stack(mod, len(ms), **kwargs)
        
        return (mod, rtxt)

class DavchaRescaleSigmas(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="DavchaRescaleSigmas",
            category="advanced/sampling",
            inputs=[
                io.Sigmas.Input("sigmas"),
                io.String.Input("scales", multiline=False, dynamic_prompts=False),
            ],
            outputs=[
                io.Sigmas.Output(),
            ],
        )

    @classmethod
    def execute(cls, sigmas, scales) -> io.NodeOutput:
        if scales.strip() == "":
            return io.NodeOutput(sigmas)
        
        scale_factors = [float(s) for s in scales.split(',')]
        rescaled_sigmas = sigmas.clone()
        for i, factor in enumerate(scale_factors):
            if i < len(rescaled_sigmas):
                rescaled_sigmas[i] = rescaled_sigmas[i] * factor
                
        return io.NodeOutput(rescaled_sigmas)

class DavchaSamplerCustomAdvanced(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="DavchaSamplerCustomAdvanced",
            category="sampling/custom_sampling",
            inputs=[
                io.Noise.Input("noise"),
                io.Guider.Input("guider"),
                io.String.Input("samplers_list", multiline=False, dynamic_prompts=False),
                io.String.Input("sigmas_list", multiline=False, dynamic_prompts=False),
                io.Latent.Input("latent_image"),
            ],
            outputs=[
                io.Latent.Output(display_name="output"),
                io.Latent.Output(display_name="denoised_output"),
            ]
        )

    @classmethod
    def execute(cls, noise, guider, samplers_list, sigmas_list, latent_image) -> io.NodeOutput:
        samplers_names = [s.strip() for s in samplers_list.split(',')]
        sigmas_l = json.loads(sigmas_list)
        sigmas_l = [torch.tensor(s, device=comfy.model_management.intermediate_device()) for s in sigmas_l]
        latent = latent_image
        latent_image = latent["samples"]
        latent = latent.copy()
        latent_image = comfy.sample.fix_empty_latent_channels(guider.model_patcher, latent_image, latent.get("downscale_ratio_spacial", None))
        latent["samples"] = latent_image

        noise_mask = None
        if "noise_mask" in latent:
            noise_mask = latent["noise_mask"]

        x0_output = {}
        sigmas_len = sum([len(s) - 1 for s in sigmas_l])
        callback = latent_preview.prepare_callback(guider.model_patcher, sigmas_len, x0_output)
        disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
        curr_noise = noise.generate_noise(latent)
        for sampler_name, sigmas in zip(samplers_names, sigmas_l):
            sampler = sampler_object(sampler_name)            
            samples = guider.sample(curr_noise, latent_image, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise.seed)
            latent_image = samples
            curr_noise = Noise_EmptyNoise().generate_noise({'samples': latent_image})
        samples = samples.to(comfy.model_management.intermediate_device())

        out = latent.copy()
        out.pop("downscale_ratio_spacial", None)
        out["samples"] = samples
        if "x0" in x0_output:
            x0_out = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu())
            if samples.is_nested:
                latent_shapes = [x.shape for x in samples.unbind()]
                x0_out = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0_out, latent_shapes))
            out_denoised = latent.copy()
            out_denoised["samples"] = x0_out
        else:
            out_denoised = out
        return io.NodeOutput(out, out_denoised)

    sample = execute

class DavchaSamplerCustomAdvanced2(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="DavchaSamplerCustomAdvanced2",
            category="sampling/custom_sampling",
            inputs=[
                io.Noise.Input("noise"),
                io.Guider.Input("guider"),
                io.String.Input("samplers", multiline=False, dynamic_prompts=False),
                io.Sigmas.Input("sigmas"),
                io.String.Input("splits", multiline=False, dynamic_prompts=False),
                io.Latent.Input("latent_image"),
            ],
            outputs=[
                io.Latent.Output(display_name="output"),
                io.Latent.Output(display_name="denoised_output"),
            ]
        )

    @classmethod
    def execute(cls, noise, guider, samplers, sigmas, splits, latent_image) -> io.NodeOutput:
        samplers_names = [s.strip() for s in samplers.split(',')]
        splits_list = [int(s.strip()) for s in splits.split(',')]
        splits_l = np.array([0] + splits_list + [len(sigmas)])
        sigmas_l = [sigmas[start:end] for start, end in list(zip(splits_l, splits_l[1:]+1))]
        latent = latent_image
        latent_image = latent["samples"]
        latent = latent.copy()
        latent_image = comfy.sample.fix_empty_latent_channels(guider.model_patcher, latent_image, latent.get("downscale_ratio_spacial", None))
        latent["samples"] = latent_image

        noise_mask = None
        if "noise_mask" in latent:
            noise_mask = latent["noise_mask"]

        x0_output = {}
        sigmas_len = sum([len(s) - 1 for s in sigmas_l])
        callback = latent_preview.prepare_callback(guider.model_patcher, sigmas_len, x0_output)
        disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
        curr_noise = noise.generate_noise(latent)
        for sampler_name, sigmas in zip(samplers_names, sigmas_l):
            sampler = sampler_object(sampler_name)            
            samples = guider.sample(curr_noise, latent_image, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise.seed)
            latent_image = samples
            curr_noise = Noise_EmptyNoise().generate_noise({'samples': latent_image})
        samples = samples.to(comfy.model_management.intermediate_device())

        out = latent.copy()
        out.pop("downscale_ratio_spacial", None)
        out["samples"] = samples
        if "x0" in x0_output:
            x0_out = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu())
            if samples.is_nested:
                latent_shapes = [x.shape for x in samples.unbind()]
                x0_out = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0_out, latent_shapes))
            out_denoised = latent.copy()
            out_denoised["samples"] = x0_out
        else:
            out_denoised = out
        return io.NodeOutput(out, out_denoised)

    sample = execute

class DavchaQwenVL3Loader(io.ComfyNode):
    gguf_folder = os.path.join(folder_paths.models_dir, "LLM", "GGUF")
    
    @classmethod
    def define_schema(cls):
        path = os.path.join(cls.gguf_folder, "**", "*.gguf")
        files = list(filter(lambda x: "mmproj" not in x, glob(path, recursive=True)))
        
        options = []
        for file in files:
            llm = Llama(file, vocab_only=True, verbose=False)
            arch = llm.metadata.get("general.architecture", "llama")
            max_n_ctx = int(llm.metadata.get(f"{arch}.context_length"))
            max_n_gpu_layers = int(llm.metadata.get(f"{arch}.block_count"))
            
            inputs = [
                io.Int.Input("n_ctx", min=512, max=max_n_ctx, default=4096),
                io.Int.Input("n_gpu_layers", min=-1, max=max_n_gpu_layers, default=-1),
            ]
                
            name = os.path.relpath(file, cls.gguf_folder)
            option = io.DynamicCombo.Option(name, inputs)
            options.append(option)
        
        return io.Schema(
            node_id="DavchaQwenVL3Loader",
            category="davcha/llm",
            inputs=[
                io.DynamicCombo.Input("model", options=options),
                io.Int.Input("n_batch", min=64, max=32768, default=512),
                io.Int.Input("pool_size", min=1048576, max=10485760, default=4194304, step=524288),
            ],
            outputs=[
                io.Custom("DavchaQwenVL3Model").Output(),
            ]
        )
    
    @classmethod
    def execute(cls, model, n_batch, pool_size):
        from llama_cpp import Llama
            
        n_ctx = model.get("n_ctx", 4096)
        n_gpu_layers = model.get("n_gpu_layers", -1)
        model = model.get("model", None)
        if model is None:
            raise FileExistsError("model")
        model = os.path.join(cls.gguf_folder, model)
        
        mmproj = next(iter(glob(os.path.join(os.path.dirname(model), "*mmproj*"))), None)
        if mmproj:
            from llama_cpp.llama_chat_format import Qwen3VLChatHandler
            mmproj = Qwen3VLChatHandler(clip_model_path=mmproj, force_reasoning=False)

        llm = Llama(
            model_path=model,
            chat_handler=mmproj,
            n_ctx=n_ctx,
            n_gpu_layers=n_gpu_layers,
            swa_full=True,
            n_batch=n_batch,
            pool_size=pool_size,
            verbose=False
        )
        
        return io.NodeOutput(llm)

class DavchaQwenVL3(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="DavchaQwenVL3",
            category="davcha/llm",
            inputs=[
                io.Custom("DavchaQwenVL3Model").Input("llm"),
                io.String.Input("system", multiline=True, dynamic_prompts=False),
                io.String.Input("prompt", multiline=True, dynamic_prompts=False),
                io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                io.Int.Input("max_tokens", min=1, default=512),
                io.Float.Input("temperature", min=0.0, max=2.0, default=0.9),
                io.Float.Input("top_p", min=0.0, max=1.0, default=0.95),
                io.Int.Input("top_k", min=0, max=100, default=40),
                io.Float.Input("repeat_penalty", min=0.5, max=2.0, default=1.2),
                io.String.Input("response_format", multiline=True, dynamic_prompts=False, default=""),
                io.Image.Input("images", optional=True),
            ],
            outputs=[
                io.String.Output()
            ]
        )
    
    @classmethod
    def execute(cls, llm, system, prompt, seed, max_tokens, temperature, top_p, top_k, repeat_penalty, response_format, images=None):
        if images is not None:
            content = []
            
            for image in images:
                array = (image*255).clamp(0, 255).byte().cpu().numpy()
                pil_img = Image.fromarray(array, mode="RGB")
                buf = BytesIO()
                pil_img.save(buf, format="PNG")
                b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
                content.append({'type': 'image_url', 'image_url': {'url': f"data:image/png;base64,{b64}"}})
            
            content.append({'type': 'text', 'text': prompt})

            messages = [
                {"role": "system", "content": system},
                {"role": "user", "content": content}
            ]
        else:
            messages = [
                {"role": "system", "content": system},
                {"role": "user", "content": prompt}
            ]
            
        response_format = json.loads(response_format) if response_format.strip() != "" else None
        
        result = llm.create_chat_completion(
            seed=seed,
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            repeat_penalty=repeat_penalty,
            response_format=response_format,
            stop=["<|im_end|>", "<|im_start|>", "<im_end>", "<im_start>", "</im_end>", "</im_start>"]
        )
        
        return io.NodeOutput(result['choices'][0]['message']['content'])

class OldDavchaPromptEnricher(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OldDavchaPromptEnricher",
            category="davcha/llm",
            inputs=[
                io.Custom("DavchaQwenVL3Model").Input("llm"),
                io.String.Input("system", multiline=True, dynamic_prompts=False),
                io.String.Input("prompt", multiline=True, dynamic_prompts=False),
                io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
                io.Int.Input("max_tokens", min=1, default=512),
                io.Float.Input("temperature", min=0.0, max=2.0, default=0.9, step=0.01),
                io.Float.Input("top_p", min=0.0, max=1.0, default=0.95, step=0.01),
                io.Int.Input("top_k", min=0, max=100, default=40),
                io.Float.Input("repeat_penalty", min=0.5, max=2.0, default=1.2, step=0.01),
                io.Boolean.Input("force_json_output", default=True),
                io.Image.Input("images", optional=True),
            ],
            outputs=[
                io.String.Output()
            ]
        )
    
    @classmethod
    def execute(cls, llm, system, prompt, seed, max_tokens, temperature, top_p, top_k, repeat_penalty, force_json_output, images=None):
        if not re.search(r'\{([^}]+)\}', prompt):
            return io.NodeOutput(prompt)
        
        m = re.findall(r'\{([^}]+)\}', prompt)
        keys = '\n'.join([f'{x}:' for x in m])

        p = f"""PROMPT:\n{prompt}\n\nKEYS:\n{keys}"""
        
        if images is not None:
            content = []
            
            for image in images:
                array = (image*255).clamp(0, 255).byte().cpu().numpy()
                pil_img = Image.fromarray(array, mode="RGB")
                buf = BytesIO()
                pil_img.save(buf, format="PNG")
                b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
                content.append({'type': 'image_url', 'image_url': {'url': f"data:image/png;base64,{b64}"}})
            
            content.append({'type': 'text', 'text': p})

            messages = [
                {"role": "system", "content": system},
                {"role": "user", "content": content}
            ]
        else:
            messages = [
                {"role": "system", "content": system},
                {"role": "user", "content": p}
            ]

        if force_json_output:
            response_format = {
                "type": "json_object",
                "schema": {
                    "type": "object",
                    "properties": {key:{"type":"string"} for key in m},
                    "required": m
                }
            }
        else:
            response_format = None
        
        result = llm.create_chat_completion(
            seed=seed,
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            repeat_penalty=repeat_penalty,
            response_format=response_format,
            stop=["<|im_end|>", "<|im_start|>", "<im_end>", "<im_start>", "</im_end>", "</im_start>"]
        )
        groups = result['choices'][0]['message']['content']
        import json_repair
        p = prompt
        for k, v in json_repair.loads(groups).items():
            p = p.replace(f"{{{k}}}", "\n".join([str(x) for x in v]) if isinstance(v, list) else str(v))
        
        return io.NodeOutput(p)
    
class DavchaImageStack(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        template = io.Autogrow.TemplatePrefix(input=io.Image.Input("image", optional=True), prefix="image", min=1, max=10)
        return io.Schema(
            node_id="DavchaImageStack",
            category="image",
            inputs=[
                io.Autogrow.Input("autogrow", template=template, optional=True),
            ],
            outputs=[
                io.Image.Output()
            ]
        )
    
    @classmethod
    def execute(cls, autogrow) -> io.NodeOutput: 
        images = [v for k, v in sorted(autogrow.items()) if v is not None]
        return io.NodeOutput(images)

class DavchaScheduledReferenceLatent(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="DavchaScheduledReferenceLatent",
            category="latent",
            inputs=[
                io.Conditioning.Input("cond"),
                io.Latent.Input("latent", optional=True),
                io.Float.Input("start", min=0.0, max=1.0, default=0.0, step=0.01),
                io.Float.Input("end", min=0.0, max=1.0, default=1.0, step=0.01),
            ],
            outputs=[
                io.Conditioning.Output()
            ]
        )

    @classmethod
    def execute(cls, cond, latent, start, end) -> io.NodeOutput:
        ts = {0.0,1.0,start,end}
        for x in cond:
            if "start_percent" in x[1]:
                ts.add(x[1]['start_percent'])
                ts.add(x[1]['end_percent'])
            else:
                x[1]["start_percent"] = 0
                x[1]["end_percent"] = 1
        ts = sorted(list(ts))

        c = []

        for s, e in zip(ts, ts[1:]):
            midpoint = (s + e) / 2.0

            for x in cond:
                if x[1]['start_percent'] <= midpoint < x[1]['end_percent']:
                    active_cond = [x]
                break

            if start <= midpoint < end:
                active_cond = node_helpers.conditioning_set_values(active_cond, {"reference_latents": [latent["samples"]]}, append=True)

            c.extend(node_helpers.conditioning_set_values(active_cond, {"start_percent": s, "end_percent": e}))

        return io.NodeOutput(c)

class DavchaScaleImageToTotalPixelsMax(io.ComfyNode):
    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
    crop_methods = ["disabled", "center"]
    
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="DavchaScaleImageToTotalPixelsMax",
            display_name="Scale Image to Total Pixels",
            category="image/upscaling",
            inputs=[
                io.Combo.Input("upscale_method", options=cls.upscale_methods),
                io.Float.Input("megapixels", default=1.0, min=0.01, max=16.0, step=0.01),
                io.Int.Input("resolution_steps", default=1, min=1, max=256, advanced=True),
                io.Image.Input("image", optional=True),
            ],
            outputs=[
                io.Image.Output(),
            ],
        )

    @classmethod
    def execute(cls, upscale_method, megapixels, resolution_steps, image=None) -> io.NodeOutput:
        if image is None:
            return io.NodeOutput(None)
        samples = image.movedim(-1,1)
        total = megapixels * 1024 * 1024

        scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
        width = round(samples.shape[3] * scale_by / resolution_steps) * resolution_steps
        height = round(samples.shape[2] * scale_by / resolution_steps) * resolution_steps

        s = comfy.utils.common_upscale(samples, int(width), int(height), upscale_method, "disabled")
        s = s.movedim(1,-1)
        return io.NodeOutput(s)

NODE_CLASS_MAPPINGS = {
    'DavchaScheduledReferenceLatent': DavchaScheduledReferenceLatent,
    'DavchaScaleImageToTotalPixelsMax': DavchaScaleImageToTotalPixelsMax,
    'DavchaImageStack': DavchaImageStack,
    'OldDavchaPromptEnricher': OldDavchaPromptEnricher,
    'DavchaQwenVL3Loader': DavchaQwenVL3Loader,
    'DavchaQwenVL3': DavchaQwenVL3,
    'DavchaSamplerCustomAdvanced': DavchaSamplerCustomAdvanced,
    'DavchaSamplerCustomAdvanced2': DavchaSamplerCustomAdvanced2,
    'DavchaRescaleSigmas': DavchaRescaleSigmas,
    'DavchaQwenImageEditLoraTagLoader': DavchaQwenImageEditLoraTagLoader,
    'DavchaScheduledTextEncoderQwenImageEditPlus': DavchaScheduledTextEncoderQwenImageEditPlus,
    'DavchaTextEncodeQwenImageEditPlus': DavchaTextEncodeQwenImageEditPlus,
    'DavchaWan22LoraTagLoader': DavchaWan22LoraTagLoader,
    'DavchaWan22LoraTagParser': DavchaWan22LoraTagParser,
    'OldDavchaLoadLLM': OldDavchaLoadLLM,
    'OldDavchaLLM': OldDavchaLLM,
    'DavchaLLMAdvanced': DavchaLLMAdvanced,
    'PadAndResize': PadAndResize,
    'SmartMask': SmartMask,
    'ResizeCropFit': ResizeCropFit,
    'PercentPadding': PercentPadding,
    'SoftErosion': SoftErosion,
    'ApplyMask': ApplyMask,
    'StringScheduleHelper': StringScheduleHelper,
    'DStack': DStack,
    'DavchaConditioningConcat': DavchaConditioningConcat,
    'DavchaModelMergeSimple' : DavchaModelMergeSimple,
    'DavchaCLIPMergeSimple' : DavchaCLIPMergeSimple,
    'DavchaModelMergeSD1' : DavchaModelMergeSD1,
    'DavchaModelMergeSDXL' : DavchaModelMergeSDXL,
    'ConditioningCompress': ConditioningCompress,
    'DavchaCLIPTextEncode': DavchaCLIPTextEncode,
    'DavchaLoadVideo': DavchaLoadVideo,
    'DavchaEmptyLatentImage': DavchaEmptyLatentImage,
    'DavchaMaskImage': DavchaMaskImage,
    'DavchaPop': DavchaPop,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    'DavchaScheduledReferenceLatent': 'Scheduled Reference Latent',
    'DavchaScaleImageToTotalPixelsMax': 'Scale Image To Total Pixels Max',
    'DavchaImageStack': 'Davcha Image Stack',
    'OldDavchaPromptEnricher': 'Davcha Prompt Enricher',
    'DavchaQwenVL3Loader': 'DavchaQwenVL3Loader',
    'DavchaQwenVL3': 'DavchaQwenVL3',
    'DavchaSamplerCustomAdvanced': 'Custom Sampler Advanced',
    'DavchaSamplerCustomAdvanced2': 'Custom Sampler Advanced 2',
    'DavchaRescaleSigmas': 'Rescale Sigmas',
    'DavchaQwenImageEditLoraTagLoader': 'Qwen Image Edit Lora Tag Loader',
    'DavchaScheduledTextEncoderQwenImageEditPlus': 'Scheduled Text Encoder Qwen Image Edit Plus',
    'DavchaTextEncodeQwenImageEditPlus': 'Text Encode Qwen Image Edit Plus',
    'DavchaWan22LoraTagLoader': 'Wan22 Lora Tag Loader',
    'DavchaWan22LoraTagParser': 'Wan22 Lora Tag Parser',
    'OldDavchaLoadLLM': 'OldDavchaLoadLLM',
    'OldDavchaLLM': 'OldDavchaLLM',
    'DavchaLLMAdvanced': 'DavchaLLMAdvanced',
    'PadAndResize': 'PadAndResize',
    'SmartMask': 'SmartMask',
    'ResizeCropFit': 'Resize, Crop or Fit',
    'PercentPadding': 'Percent Padding',
    'SoftErosion': 'SoftErosion',
    'ApplyMask': 'ApplyMask',
    'StringScheduleHelper': 'StringScheduleHelper',
    'DStack': 'DStack',
    'DavchaConditioningConcat': 'DavchaConditioningConcat',
    'DavchaModelMergeSimple': 'DavchaModelMergeSimple',
    'DavchaCLIPMergeSimple': 'DavchaCLIPMergeSimple',
    'DavchaModelMergeSD1': 'DavchaModelMergeSD1',
    'DavchaModelMergeSDXL': 'DavchaModelMergeSDXL',
    'ConditioningCompress': 'ConditioningCompress',
    'DavchaCLIPTextEncode': 'CLIP Text Encode (Davcha)',
    'DavchaLoadVideo': 'Load Video (Davcha)',
    'DavchaEmptyLatentImage': 'Empty Latent Image (Davcha)',
    'DavchaMaskImage': 'Mask Image (Davcha)',
    'DavchaPop': 'Pop',
}