From 5449ab7cf077d7ba6febb88825d8dc7c9c15eee0 Mon Sep 17 00:00:00 2001
From: nkchocoai <nkai3359@gmail.com>
Date: Sun, 18 Feb 2024 00:22:29 +0900
Subject: [PATCH] Support Clip skip

---
 py/capture.py                   |  5 ++++-
 py/defs/captures.py             | 13 ++++++++++++-
 py/defs/ext/efficiency_nodes.py |  7 ++++++-
 py/defs/formatters.py           |  4 ++++
 py/defs/meta.py                 |  1 +
 py/trace.py                     |  2 --
 6 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/py/capture.py b/py/capture.py
index e1e6da1..4956950 100644
--- a/py/capture.py
+++ b/py/capture.py
@@ -55,6 +55,7 @@ class Capture:
 
         update_pnginfo_dict(MetaField.POSITIVE_PROMPT, "Positive prompt")
         update_pnginfo_dict(MetaField.NEGATIVE_PROMPT, "Negative prompt")
+
         update_pnginfo_dict(MetaField.STEPS, "Steps")
 
         sampler_names = inputs.get(MetaField.SAMPLER_NAME, [])
@@ -70,12 +71,14 @@ class Capture:
         update_pnginfo_dict(MetaField.CFG, "CFG Scale")
         update_pnginfo_dict(MetaField.SEED, "Seed")
 
+        update_pnginfo_dict(MetaField.CLIP_SKIP, "Clip skip")
+
         image_widths = inputs.get(MetaField.IMAGE_WIDTH, [])
         image_heights = inputs.get(MetaField.IMAGE_HEIGHT, [])
         if len(image_widths) > 0 and len(image_heights) > 0:
             pnginfo_dict["Size"] = f"{image_widths[0][1]}x{image_heights[0][1]}"
 
-        update_pnginfo_dict(MetaField.MODEL_NAME, "Model name")
+        update_pnginfo_dict(MetaField.MODEL_NAME, "Model")
         update_pnginfo_dict(MetaField.MODEL_HASH, "Model hash")
 
         return pnginfo_dict
diff --git a/py/defs/captures.py b/py/defs/captures.py
index 41cf3c3..0cd0e0c 100644
--- a/py/defs/captures.py
+++ b/py/defs/captures.py
@@ -1,6 +1,11 @@
 from .meta import MetaField
 from .validators import is_positive_prompt, is_negative_prompt
-from .formatters import calc_model_hash, get_scaled_width, get_scaled_height
+from .formatters import (
+    calc_model_hash,
+    convert_skip_clip,
+    get_scaled_width,
+    get_scaled_height,
+)
 
 
 CAPTURE_FIELD_LIST = {
@@ -8,6 +13,12 @@ CAPTURE_FIELD_LIST = {
         MetaField.MODEL_NAME: {"field_name": "ckpt_name"},
         MetaField.MODEL_HASH: {"field_name": "ckpt_name", "format": calc_model_hash},
     },
+    "CLIPSetLastLayer": {
+        MetaField.CLIP_SKIP: {
+            "field_name": "stop_at_clip_layer",
+            "format": convert_skip_clip,
+        },
+    },
     "EmptyLatentImage": {
         MetaField.IMAGE_WIDTH: {"field_name": "width"},
         MetaField.IMAGE_HEIGHT: {"field_name": "height"},
diff --git a/py/defs/ext/efficiency_nodes.py b/py/defs/ext/efficiency_nodes.py
index 6de3302..c91f30e 100644
--- a/py/defs/ext/efficiency_nodes.py
+++ b/py/defs/ext/efficiency_nodes.py
@@ -1,6 +1,6 @@
 # https://github.com/jags111/efficiency-nodes-comfyui
 from ..meta import MetaField
-from ..formatters import calc_model_hash
+from ..formatters import calc_model_hash, convert_skip_clip
 
 SAMPLERS = {
     "KSampler (Efficient)": {
@@ -21,6 +21,7 @@ CAPTURE_FIELD_LIST = {
     "Efficient Loader": {
         MetaField.MODEL_NAME: {"field_name": "ckpt_name"},
         MetaField.MODEL_HASH: {"field_name": "ckpt_name", "format": calc_model_hash},
+        MetaField.CLIP_SKIP: {"field_name": "clip_skip", "format": convert_skip_clip},
         MetaField.POSITIVE_PROMPT: {"field_name": "positive"},
         MetaField.NEGATIVE_PROMPT: {"field_name": "negative"},
         MetaField.IMAGE_WIDTH: {"field_name": "empty_latent_width"},
@@ -32,6 +33,10 @@ CAPTURE_FIELD_LIST = {
             "field_name": "base_ckpt_name",
             "format": calc_model_hash,
         },
+        MetaField.CLIP_SKIP: {
+            "field_name": "base_clip_skip",
+            "format": convert_skip_clip,
+        },
         MetaField.POSITIVE_PROMPT: {"field_name": "positive"},
         MetaField.NEGATIVE_PROMPT: {"field_name": "negative"},
         MetaField.IMAGE_WIDTH: {"field_name": "empty_latent_width"},
diff --git a/py/defs/formatters.py b/py/defs/formatters.py
index d85cc57..4a183b4 100644
--- a/py/defs/formatters.py
+++ b/py/defs/formatters.py
@@ -21,6 +21,10 @@ def calc_model_hash(model_name, input_data):
     return model_hash
 
 
+def convert_skip_clip(stop_at_clip_layer, input_data):
+    return stop_at_clip_layer * -1
+
+
 def get_scaled_width(scaled_by, input_data):
     samples = input_data["samples"][0]["samples"]
     return round(samples.shape[3] * scaled_by * 8)
diff --git a/py/defs/meta.py b/py/defs/meta.py
index 250ee38..b18e2c9 100644
--- a/py/defs/meta.py
+++ b/py/defs/meta.py
@@ -6,6 +6,7 @@ class MetaField(IntEnum):
     MODEL_HASH = 1
     POSITIVE_PROMPT = 10
     NEGATIVE_PROMPT = 11
+    CLIP_SKIP = 12
     SEED = 20
     STEPS = 21
     CFG = 22
diff --git a/py/trace.py b/py/trace.py
index 1da6cc6..83f5550 100644
--- a/py/trace.py
+++ b/py/trace.py
@@ -1,5 +1,3 @@
-import json
-
 from collections import deque
 
 from .defs.samplers import SAMPLERS