Imagine/imagine-v5/vyro_workflows/nodes/loaders.py

from functools import lru_cache
import json
import os
# from turtle import speed

import folder_paths
from folder_paths import folder_names_and_paths, supported_pt_extensions
import nodes
import spacy
import torch

def update_paths():
        model_path = folder_paths.models_dir
        folder_names_and_paths["vyro_configs"] = ([os.path.join(os.path.dirname(os.path.abspath(__file__)), "../configs")], ['.json'])
        folder_names_and_paths["spacy"] = ([(os.path.join(model_path, "spacy"))], supported_pt_extensions)
        folder_names_and_paths["interposers"] = ([(os.path.join(model_path, "interposers"))], supported_pt_extensions)
        folder_names_and_paths['oneflow_graphs'] = ([(os.path.join(model_path, "oneflow_graphs"))], (""))
        
update_paths()

class VyroConfigLoader:
    def __init__(self):
        pass
    
    @classmethod
    def INPUT_TYPES(s):   
        
        paths = []
        update_paths()
        spacy_paths = folder_paths.get_folder_paths("spacy")
        print(f"[VyroConfigLoader] spacy paths: {spacy_paths}")  # Debug print

        for search_path in spacy_paths:
            if os.path.exists(search_path):
                print(f"[VyroConfigLoader] Found spacy path: {search_path}") # Debug print
                for root, subdir, files in os.walk(search_path, followlinks=True):
                    if "config.cfg" in files:
                        rel_path = os.path.relpath(root, start=search_path)
                        paths.append(rel_path)
                        print(f"[VyroConfigLoader] Added classifier path: {rel_path}") # Debug print

        print(f"[VyroConfigLoader] Available classifier paths: {paths}")  # Debug print
        return {
            "required": {
                "config_path": (folder_paths.get_filename_list("vyro_configs"),),
                "classifier_path": (paths,),
            }
        }
    
    RETURN_TYPES = ("LIST","DICT","DICT","TRANSFORMER","LIST")
    RETURN_NAMES = ("styles","prompt_tree","model_config","classifier","unweighted_styles")
    FUNCTION = "load_config"
    CATEGORY = "Vyro/Loaders"
    
    # def load_config(self, config_path, classifier_path):
    #     config_path = folder_paths.get_full_path("vyro_configs", config_path)
    #     classifier_path = os.path.join(folder_names_and_paths["spacy"][0][0],classifier_path)
    #     # classifier = pipeline("zero-shot-classification", model=classifier_path,device='cuda:0',local_files_only=True,use_safetensors=True)
    #     spacy.prefer_gpu(gpu_id=0)
    #     classifier = spacy.load(classifier_path)
    #     # classifier = pipeline("zero-shot-classification", model=classifier_path, config=f'{classifier_path}.json', device=0)
    #     with open(config_path, 'r') as json_file:
    #         try:
    #             config = json.load(json_file)
    #         except json.JSONDecodeError as json_e:
    #             print(f"[VyroConfigLoader] Error loading {config_path}:", json_e)
    #             return None
    #         unweighted_styles = []
    #         for style in config['styles']:
    #             if ':' in style:
    #                 style = style.split(':')[0]
    #             unweighted_styles.append(style)
    #         return (config['styles'], config['prompt_tree'], config['model_config'], classifier, unweighted_styles)

    def load_config(self, config_path, classifier_path):
        config_path = folder_paths.get_full_path("vyro_configs", config_path)
        classifier_path = os.path.join(folder_names_and_paths["spacy"][0][0], classifier_path)
        
        # Load the configuration file first to get styles
        with open(config_path, 'r') as json_file:
            try:
                config = json.load(json_file)
            except json.JSONDecodeError as json_e:
                print(f"[VyroConfigLoader] Error loading {config_path}:", json_e)
                return None
        
        # Extract styles from config
        styles_list = []
        unweighted_styles = []
        for style in config['styles']:
            if ':' in style:
                style_name = style.split(':')[0]
            else:
                style_name = style
            unweighted_styles.append(style_name)
            styles_list.append(style_name.replace(' ', '_'))
        
        # Load the base spaCy model
        spacy.prefer_gpu(gpu_id=0)
        classifier = spacy.load(classifier_path)
        
        # Remove existing textcat components if any
        if "textcat" in classifier.pipe_names:
            classifier.remove_pipe("textcat")
        if "textcat_multilabel" in classifier.pipe_names:
            classifier.remove_pipe("textcat_multilabel")
        
        # Create a custom style matcher component instead of textcat_multilabel
        @classifier.component("style_matcher")
        def style_matcher(doc):
            # Initialize scores dictionary
            doc.cats = {}
            
            # Simple rule-based matching
            text_lower = doc.text.lower()
            
            # Set a base score for all styles
            for style in styles_list:
                style_human = style.replace('_', ' ').lower()
                # Default low score
                base_score = 0.1
                
                # Check for exact matches
                if style_human in text_lower:
                    base_score = 0.9
                # Check for partial matches
                elif any(word in text_lower for word in style_human.split()):
                    base_score = 0.5
                    
                doc.cats[style] = base_score
            
            # If no strong matches found, set the first style as default with medium score
            if not any(score > 0.5 for score in doc.cats.values()) and styles_list:
                doc.cats[styles_list[0]] = 0.6
                
            return doc
        
        # Add the component to the pipeline
        if "style_matcher" not in classifier.pipe_names:
            classifier.add_pipe("style_matcher")
        
        print(f"[VyroConfigLoader] Successfully configured style_matcher with {len(styles_list)} labels")
        print(f"[VyroConfigLoader] Pipeline: {classifier.pipe_names}")
        
        return (config['styles'], config['prompt_tree'], config['model_config'], classifier, unweighted_styles)



class VyroModelLoader:
    def __init__(self):
        self.chkp_loader = nodes.CheckpointLoaderSimple()
        self.lora_loader = nodes.LoraLoader()
        self.tree = None
        self.config = None
        print("\n\nInitializing VyroModelLoader")
        
    @lru_cache(maxsize=6)
    def get_model(self, cfg):
        base = self.config['configs'][cfg]['base']
        refiner = self.config['configs'][cfg]['refiner']
        loras = self.config['configs'][cfg]['loras']
        tonemap_multiplier = self.config['configs'][cfg]['tonemap']
        base_model, base_clip, _ = self.chkp_loader.load_checkpoint(base)
        refiner_model, refiner_clip, _ = self.chkp_loader.load_checkpoint(refiner)
        for lora in loras:
            base_model, base_clip = self.lora_loader.load_lora(base_model, base_clip, lora['name'], lora['unet'], lora['clip'])
            
        def sampler_tonemap_reinhard(args):
            cond = args["cond"]
            uncond = args["uncond"]
            cond_scale = args["cond_scale"]
            noise_pred = (cond - uncond)
            noise_pred_vector_magnitude = (torch.linalg.vector_norm(noise_pred, dim=(1)) + 0.0000000001)[:,None]
            noise_pred /= noise_pred_vector_magnitude

            mean = torch.mean(noise_pred_vector_magnitude, dim=(1,2,3), keepdim=True)
            std = torch.std(noise_pred_vector_magnitude, dim=(1,2,3), keepdim=True)

            top = (std * 3 + mean) * tonemap_multiplier

            #reinhard
            noise_pred_vector_magnitude *= (1.0 / top)
            new_magnitude = noise_pred_vector_magnitude / (noise_pred_vector_magnitude + 1.0)
            new_magnitude *= top

            return uncond + noise_pred * new_magnitude * cond_scale

        base_model = base_model.clone()
        base_model.set_model_sampler_cfg_function(sampler_tonemap_reinhard)
        
        refiner_model = refiner_model.clone()
        refiner_model.set_model_sampler_cfg_function(sampler_tonemap_reinhard)
        
        return (base_model, refiner_model, base_clip, refiner_clip)
        
        
    @classmethod
    def INPUT_TYPES(s):   
        return {
            "required": {
                "style": ("STYLE",),
                "prompt_tree": ("DICT",),
                "model_config": ("DICT",),
            }
        }
    
    RETURN_TYPES = ("MODEL","MODEL","CLIP","CLIP")
    RETURN_NAMES = ("base_model","refiner_model","base_clip","refiner_clip")
    FUNCTION = "load"
    CATEGORY = "Vyro/Loaders"
    
    def load(self, style, prompt_tree, model_config):
        print("\n\nExecuting VyroModelLoader load function...")
        if prompt_tree is None:
            raise ValueError("Prompt tree is None")

        if style is None or style == "qr":
            print("⛔ Style is qr changing to default qr models")
            self.tree = prompt_tree
            self.config = model_config
            cfg = 'default_qr'
            return self.get_model(cfg)
            #raise ValueError("Style is None")


        if style is None or style == "":
            print("⛔ Style is none using Default config")
            self.tree = prompt_tree
            self.config = model_config
            cfg = 'default'
            return self.get_model(cfg)
            #raise ValueError("Style is None")
        
        if style not in prompt_tree.keys():
            raise ValueError("Style not in prompt tree")
        
        node = prompt_tree[style]
        if 'model_config' not in node:
            cfg = 'default'
        else:
            cfg = node['model_config']
            
        self.tree = prompt_tree
        self.config = model_config
        
        return self.get_model(cfg)
    
    
    
class VyroOneflowModelLoader:
    def __init__(self):
        self.chkp_loader = nodes.CheckpointLoaderSimple()
        self.lora_loader = nodes.LoraLoader()
        self.tree = None
        self.config = None
        print("\n\nInitializing VyroModelLoader")
        
    @lru_cache(maxsize=6)
    def get_model(self, cfg):
        base = self.config['configs'][cfg]['base']
        refiner = self.config['configs'][cfg]['refiner']
        loras = self.config['configs'][cfg]['loras']
        tonemap_multiplier = self.config['configs'][cfg]['tonemap']
        base_model, base_clip, _ = self.chkp_loader.load_checkpoint(base)
        refiner_model, refiner_clip, _ = self.chkp_loader.load_checkpoint(refiner)
        for lora in loras:
            base_model, base_clip = self.lora_loader.load_lora(base_model, base_clip, lora['name'], lora['unet'], lora['clip'])
            
        def sampler_tonemap_reinhard(args):
            cond = args["cond"]
            uncond = args["uncond"]
            cond_scale = args["cond_scale"]
            noise_pred = (cond - uncond)
            noise_pred_vector_magnitude = (torch.linalg.vector_norm(noise_pred, dim=(1)) + 0.0000000001)[:,None]
            noise_pred /= noise_pred_vector_magnitude

            mean = torch.mean(noise_pred_vector_magnitude, dim=(1,2,3), keepdim=True)
            std = torch.std(noise_pred_vector_magnitude, dim=(1,2,3), keepdim=True)

            top = (std * 3 + mean) * tonemap_multiplier

            #reinhard
            noise_pred_vector_magnitude *= (1.0 / top)
            new_magnitude = noise_pred_vector_magnitude / (noise_pred_vector_magnitude + 1.0)
            new_magnitude *= top

            return uncond + noise_pred * new_magnitude * cond_scale

        base_model = base_model.clone()
        base_model.set_model_sampler_cfg_function(sampler_tonemap_reinhard)
                
        refiner_model = refiner_model.clone()
        refiner_model.set_model_sampler_cfg_function(sampler_tonemap_reinhard)
        
        return (base_model, refiner_model, base_clip, refiner_clip)
        
        
    @classmethod
    def INPUT_TYPES(s):   
        return {
            "required": {
                "style": ("STYLE",),
                "prompt_tree": ("DICT",),
                "model_config": ("DICT",),
            }
        }
    
    RETURN_TYPES = ("MODEL","MODEL","CLIP","CLIP")
    RETURN_NAMES = ("base_model","refiner_model","base_clip","refiner_clip")
    FUNCTION = "load"
    CATEGORY = "Vyro/Loaders"
    
    def load(self, style, prompt_tree, model_config):
        print("\n\nExecuting VyroOneflowModelLoader load function...")
        
        if prompt_tree is None:
            raise ValueError("Prompt tree is None")

        if style is None or style == "qr":
            print("⛔ Style is qr changing to default qr models")
            self.tree = prompt_tree
            self.config = model_config
            cfg = 'default_qr'
            return self.get_model(cfg)
            #raise ValueError("Style is None")


        if style is None or style == "":
            print("⛔ Style is none using Default config")
            self.tree = prompt_tree
            self.config = model_config
            cfg = 'default'
            return self.get_model(cfg)
            #raise ValueError("Style is None")
        
        if style not in prompt_tree.keys():
            raise ValueError("Style not in prompt tree")
        
        node = prompt_tree[style]
        if 'model_config' not in node:
            cfg = 'default'
        else:
            cfg = node['model_config']
            
        self.tree = prompt_tree
        self.config = model_config
        
        return self.get_model(cfg)     


class VyroOneFlowBaseModelLoader:
    def __init__(self) -> None:
        self.chkp_loader = nodes.CheckpointLoaderSimple()
        print("INITIALIZING ONEFLOW BASE MODEL LOADER")
    
    @lru_cache(maxsize=6)
    def get_base_model(self, base_model):
        print(f"\n\nLoading Base Model: {base_model}\n\n")
        base_model, base_clip, _ = self.chkp_loader.load_checkpoint(base_model)
        tonemap_multiplier = 1.0
        
        def sampler_tonemap_reinhard(args):
            cond = args["cond"]
            uncond = args["uncond"]
            cond_scale = args["cond_scale"]
            noise_pred = (cond - uncond)
            noise_pred_vector_magnitude = (torch.linalg.vector_norm(noise_pred, dim=(1)) + 0.0000000001)[:,None]
            noise_pred /= noise_pred_vector_magnitude

            mean = torch.mean(noise_pred_vector_magnitude, dim=(1,2,3), keepdim=True)
            std = torch.std(noise_pred_vector_magnitude, dim=(1,2,3), keepdim=True)

            top = (std * 3 + mean) * tonemap_multiplier

            #reinhard
            noise_pred_vector_magnitude *= (1.0 / top)
            new_magnitude = noise_pred_vector_magnitude / (noise_pred_vector_magnitude + 1.0)
            new_magnitude *= top

            return uncond + noise_pred * new_magnitude * cond_scale
        
        
        base_model = base_model.clone()
        base_model.set_model_sampler_cfg_function(sampler_tonemap_reinhard)
    
        return (base_model, base_clip)
        

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "base_model": (folder_paths.get_filename_list("checkpoints"), ),
            }
        }
    
    RETURN_TYPES = ("MODEL", "CLIP",)
    RETURN_NAMES = ("base_model", "base_clip",)
    FUNCTION = "load"
    CATERGORY = "Vyro/Loaders/Oneflow"
    
    def load(self, base_model):
        return self.get_base_model(base_model)


class VyroLoraLoader:
    
    def __init__(self) -> None:
        self.lora_loader = nodes.LoraLoader()
        self.tree = None
        self.config = None
        print("\n\nInitializing Vyro LORA Loader\n\n")
    
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "base_model": ("MODEL",),
                "base_clip": ("CLIP",),
                "style": ("STYLE",),
                "prompt_tree": ("DICT",),
                "model_config": ("DICT",),
            }
        }
    RETURN_TYPES = ("MODEL", "CLIP")
    RETURN_NAMES = ("base_model", "base_clip")
    FUNCTION = "load_loras"
    CATEGORY = "Vyro/Loaders/Lora"
    
    def load_loras(self, base_model, base_clip, style, prompt_tree, model_config):
        print("\n\nExecuting VyroLORA load function...")
        if prompt_tree is None:
            raise ValueError("Prompt tree is None")

        if style is None or style == "qr":
            print("⛔ Style is qr changing to default qr models")
            self.tree = prompt_tree
            self.config = model_config
            cfg = 'default_qr'
            return self.get_model(cfg)
            #raise ValueError("Style is None")


        if style is None or style == "":
            print("⛔ Style is none using Default config")
            self.tree = prompt_tree
            self.config = model_config
            cfg = 'default'
            return self.get_model(cfg)
            #raise ValueError("Style is None")
        
        if style not in prompt_tree.keys():
            raise ValueError("Style not in prompt tree")
        
        node = prompt_tree[style]
        if 'model_config' not in node:
            cfg = 'default'
        else:
            cfg = node['model_config']
            
        self.tree = prompt_tree
        self.config = model_config
        
        loras = self.config['configs'][cfg]['loras']
        for lora in loras:
            print(f"\nLoading Lora: {lora['name']}\n")
            base_model, base_clip = self.lora_loader.load_lora(base_model, base_clip, lora['name'], lora['unet'], lora['clip'])
        
        return base_model, base_clip        


class VyroOneFlowRefinerModelLoader:
    def __init__(self) -> None:
        self.chkp_loader = nodes.CheckpointLoaderSimple()
        self.tree = None
        self.config = None
        print("INITIALIZING ONEFLOW REFINER MODEL LOADER")
    
    @lru_cache(maxsize=6)
    def get_refiner_model(self, cfg):
        refiner = self.config['configs'][cfg]['refiner']
        
        print(f"\n\nLoading Refiner Model: {refiner}\n\n")
        
        tonemap_multiplier = self.config['configs'][cfg]['tonemap']
        
        refiner_model, refiner_clip, _ = self.chkp_loader.load_checkpoint(refiner)
        
        def sampler_tonemap_reinhard(args):
            cond = args["cond"]
            uncond = args["uncond"]
            cond_scale = args["cond_scale"]
            noise_pred = (cond - uncond)
            noise_pred_vector_magnitude = (torch.linalg.vector_norm(noise_pred, dim=(1)) + 0.0000000001)[:,None]
            noise_pred /= noise_pred_vector_magnitude

            mean = torch.mean(noise_pred_vector_magnitude, dim=(1,2,3), keepdim=True)
            std = torch.std(noise_pred_vector_magnitude, dim=(1,2,3), keepdim=True)

            top = (std * 3 + mean) * tonemap_multiplier

            #reinhard
            noise_pred_vector_magnitude *= (1.0 / top)
            new_magnitude = noise_pred_vector_magnitude / (noise_pred_vector_magnitude + 1.0)
            new_magnitude *= top

            return uncond + noise_pred * new_magnitude * cond_scale

        refiner_model = refiner_model.clone()
        refiner_model.set_model_sampler_cfg_function(sampler_tonemap_reinhard)
                
        return (refiner_model, refiner_clip)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "style": ("STYLE",),
                "prompt_tree": ("DICT",),
                "model_config": ("DICT",)
            }
        }
    
    RETURN_TYPES = ("MODEL", "CLIP")
    RETURN_NAMES = ("refiner_model", "refiner_clip")
    FUNCTION = "load"
    CATERGORY = "Vyro/Loaders/Oneflow"
    
    def load(self, style, prompt_tree, model_config):
        if prompt_tree is None:
            raise ValueError("Prompt Tree is None")

        if style is None or style == "qr":
            print("⛔ Style is qr changing to default qr models")
            self.tree = prompt_tree
            self.config = model_config
            cfg = "default_qr"
            return self.get_model(cfg)

        if style is None or style == "":
            print("⛔ Style is none using Default config")
            self.tree = prompt_tree
            self.config = model_config
            cfg = 'default'
            return self.get_model(cfg)
            #raise ValueError("Style is None")
        
        if style not in prompt_tree.keys():
            raise ValueError("Style not in prompt tree")
        
        node = prompt_tree[style]
        if 'model_config' not in node:
            cfg = 'default'
        else:
            cfg = node['model_config']
            
        self.tree = prompt_tree
        self.config = model_config
        
        return self.get_refiner_model(cfg)
        

        

NODE_CLASS_MAPPINGS = {
    "Vyro Config Loader": VyroConfigLoader,
    "Vyro Model Loader": VyroModelLoader,
    "Vyro OneFlow Model Loader": VyroOneflowModelLoader,
    "Vyro Oneflow Base Model Loader": VyroOneFlowBaseModelLoader,
    "Vyro Oneflow Refiner Model Loader": VyroOneFlowRefinerModelLoader,
    "Vyro LoRa Loader": VyroLoraLoader,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "VyroConfigLoader": "Vyro Config Loader",
    "VyroModelLoader": "Vyro Model Loader",
    "VyroOneFlowModelLoader": "Vyro Oneflow Model Loader",
    "VyroOneFlowBaseModelLoader": "Vyro Oneflow Base Model Loader",
    "VyroOneflowRefinerModelLoader": "Vyro Oneflow Refiner Model Loader",
    "VyroLoraLoader": "Vyro LoRa Loader"
}