From 5ad93e8bac28007ccd3afec8a299c2eac0aeb845 Mon Sep 17 00:00:00 2001
From: Leslie Perjes <gitarfogas@gmail.com>
Date: Fri, 22 Dec 2023 22:02:55 +0100
Subject: [PATCH] V1.3 - lora and lyco version displayed

---
 Nodes/Visuals.py             | 11 +++++++++++
 front_end/primere_visuals.js | 35 ++++++++++++++++++++++++++++-------
 2 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/Nodes/Visuals.py b/Nodes/Visuals.py
index ffe405b..f34174c 100644
--- a/Nodes/Visuals.py
+++ b/Nodes/Visuals.py
@@ -57,9 +57,12 @@ class PrimereVisualLORA:
     CATEGORY = TREE_VISUALS
     LORASCOUNT = 6
 
+    lora_versions = utility.get_category_from_cache('lora_version')
+
     @classmethod
     def INPUT_TYPES(cls):
         LoraList = folder_paths.get_filename_list("loras")
+
         return {
             "required": {
                 "model": ("MODEL",),
@@ -107,6 +110,9 @@ class PrimereVisualLORA:
                 "lora_keywords_num": ("INT", {"default": 1, "min": 1, "max": 50, "step": 1}),
                 "lora_keyword_weight": ("FLOAT", {"default": 1.0, "min": 0, "max": 10.0, "step": 0.1}),
             },
+            "hidden": {
+                "cached_lora": (cls.lora_versions,),
+            }
         }
 
     def visual_lora_stacker(self, model, clip, use_only_model_weight, use_lora_keyword, lora_keyword_placement, lora_keyword_selection, lora_keywords_num, lora_keyword_weight, stack_version = 'Any', model_version = "BaseModel_1024", **kwargs):
@@ -479,6 +485,8 @@ class PrimereVisualLYCORIS:
     CATEGORY = TREE_VISUALS
     LYCOSCOUNT = 6
 
+    lyco_versions = utility.get_category_from_cache('lycoris_version')
+
     @classmethod
     def INPUT_TYPES(cls):
         LYCO_DIR = os.path.join(comfy_dir, 'models', 'lycoris')
@@ -533,6 +541,9 @@ class PrimereVisualLYCORIS:
                 "lycoris_keywords_num": ("INT", {"default": 1, "min": 1, "max": 50, "step": 1}),
                 "lycoris_keyword_weight": ("FLOAT", {"default": 1.0, "min": 0, "max": 10.0, "step": 0.1}),
             },
+            "hidden": {
+                "cached_lyco": (cls.lyco_versions,),
+            }
         }
 
     def primere_visual_lycoris_stacker(self, model, clip, use_only_model_weight, use_lycoris_keyword, lycoris_keyword_placement, lycoris_keyword_selection, lycoris_keywords_num, lycoris_keyword_weight, stack_version = 'Any', model_version = "BaseModel_1024", **kwargs):
diff --git a/front_end/primere_visuals.js b/front_end/primere_visuals.js
index b8b6eca..3d65bc7 100644
--- a/front_end/primere_visuals.js
+++ b/front_end/primere_visuals.js
@@ -2,10 +2,12 @@ import { app } from "/scripts/app.js";
 
 const realPath = "extensions/Primere";
 const validClasses = ['PrimereVisualCKPT', 'PrimereVisualLORA', 'PrimereVisualEmbedding', 'PrimereVisualHypernetwork', 'PrimereVisualStyle', 'PrimereVisualLYCORIS'];
+const versionClasses = ['PrimereVisualCKPT', 'PrimereVisualLORA', 'PrimereVisualLYCORIS']
 let lastDirObject = {};
 let currentClass = false;
-let hiddenWidgets = false;
+let hiddenWidgets = {};
 let checkpointVersions = false;
+const SDXLVersionTags = ['SDXL_2048', 5120]
 
 function createCardElement(checkpoint, container, SelectedModel, ModelType) {
     let checkpoint_new = checkpoint.replaceAll('\\', '/');
@@ -21,11 +23,13 @@ function createCardElement(checkpoint, container, SelectedModel, ModelType) {
 
     let pathLastIndex = finalName.lastIndexOf('/');
     let ckptName = finalName.substring(pathLastIndex + 1);
-    let versionWidget = '<div class="ckpt-version unknown-ckpt" title="Unknow checkpoint version. Use before the version cached."></div>';
+    let versionWidget = ''; //'<div class="ckpt-version unknown-ckpt" title="Unknow checkpoint version. Use before the version cached."></div>';
 
+    //console.log(checkpointVersions)
     if (checkpointVersions !== false) {
         if (checkpointVersions.hasOwnProperty(ckptName) === true) {
-            if (checkpointVersions[ckptName] === 'SDXL_2048') {
+            //if (checkpointVersions[ckptName] === 'SDXL_2048') {
+            if (SDXLVersionTags.includes(checkpointVersions[ckptName])) {
                 versionWidget = '<div class="ckpt-version sdxl-ckpt" title="SDXL checkpoint. Select right version of additional networks."></div>';
             } else {
                 versionWidget = '<div class="ckpt-version sd1-ckpt" title="SD1.5 checkpoint. Select right version of additional networks."></div>';
@@ -326,8 +330,8 @@ app.registerExtension({
                 modaltitle = 'Select checkpoint';
                 nodematch = '^base_model';
                 isnumeric_end = false;
-                if (hiddenWidgets != false) {
-                    $.each(hiddenWidgets.hidden, function(index, value) {
+                if (hiddenWidgets.hasOwnProperty('PrimereVisualCKPT') === true) {
+                    $.each(hiddenWidgets['PrimereVisualCKPT'], function(index, value) {
                         if (index == 'cached_model') {
                             checkpointVersions = value[0];
                         }
@@ -340,6 +344,13 @@ app.registerExtension({
                 modaltitle = 'Select LoRA';
                 nodematch = '^lora_';
                 isnumeric_end = true;
+                if (hiddenWidgets.hasOwnProperty('PrimereVisualLORA') === true) {
+                    $.each(hiddenWidgets['PrimereVisualLORA'], function(index, value) {
+                        if (index == 'cached_lora') {
+                            checkpointVersions = value[0];
+                        }
+                    });
+                }
             }
 
             if (node.type == 'PrimereVisualEmbedding') {
@@ -368,6 +379,13 @@ app.registerExtension({
                 modaltitle = 'Select LYCORIS';
                 nodematch = '^lycoris_';
                 isnumeric_end = true;
+                if (hiddenWidgets.hasOwnProperty('PrimereVisualLYCORIS') === true) {
+                    $.each(hiddenWidgets['PrimereVisualLYCORIS'], function(index, value) {
+                        if (index == 'cached_lyco') {
+                            checkpointVersions = value[0];
+                        }
+                    });
+                }
             }
 
             if (event.type != LiteGraph.pointerevents_method + "down") {
@@ -455,8 +473,11 @@ app.registerExtension({
     },
 
     async beforeRegisterNodeDef(nodeType, nodeData, app) {
-        if (nodeData.name === "PrimereVisualCKPT") {
-            hiddenWidgets = nodeData.input;
+        if (versionClasses.includes(nodeData.name)) {
+            if (nodeData.input.hasOwnProperty('hidden') === true) {
+                hiddenWidgets[nodeData.name] = nodeData.input.hidden;
+            }
+            //console.log(hiddenWidgets);
         }
     },
 });