199 lines
5.0 KiB
Python
199 lines
5.0 KiB
Python
from comfy.utils import load_torch_file
|
|
from comfy.sd import load_lora_for_models
|
|
from copy import deepcopy
|
|
from jsonschema import validate
|
|
from . import schema
|
|
from .utils import clean_name, apply_schema_defaults, unclean_name
|
|
import folder_paths
|
|
import glob
|
|
import logging
|
|
import os
|
|
import os.path
|
|
import random
|
|
import re
|
|
import json
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DarkJSONGetLoRASettings(object):
|
|
"""Gets LoRA Settings based on the LoRA Name."""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"json_in": (
|
|
"JSON",
|
|
{"ForceInput": True},
|
|
),
|
|
"lora_name": (
|
|
"STRING",
|
|
{"ForceInput": True},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("JSON", "float", "float")
|
|
RETURN_NAMES = (
|
|
"json_out",
|
|
"default_model_weight",
|
|
"default_clip_weight",
|
|
)
|
|
FUNCTION = "action"
|
|
|
|
CATEGORY = "DarkPrompt"
|
|
|
|
def action(
|
|
self,
|
|
json_in,
|
|
lora_name,
|
|
):
|
|
|
|
if lora_name:
|
|
lora_name = clean_name(lora_name)
|
|
print("Looking up %s" % (lora_name))
|
|
|
|
try:
|
|
logger.info(
|
|
"Found LoRA settings for %s: %s" % (lora_name, json_in[lora_name])
|
|
)
|
|
return (
|
|
json_in[lora_name],
|
|
json_in[lora_name].get("default_model_weight", 0.8),
|
|
json_in[lora_name].get("default_clip_weight", 0),
|
|
)
|
|
except KeyError:
|
|
logger.info("Settings do not exist for LoRA: %s" % (lora_name))
|
|
return (
|
|
apply_schema_defaults({lora_name: {}}, schema.lora_settings_schema),
|
|
json_in[lora_name].get("default_model_weight", 0.8),
|
|
json_in[lora_name].get("default_clip_weight", 0),
|
|
)
|
|
else:
|
|
return (
|
|
{},
|
|
0.8,
|
|
0,
|
|
)
|
|
|
|
|
|
class DarkJSONGetLoRACheckpointSettings(object):
|
|
"""Gets LoRA Settings based on a checkpoint name."""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"json_in": (
|
|
"JSON",
|
|
{"ForceInput": True},
|
|
),
|
|
"checkpoint": (
|
|
folder_paths.get_filename_list("checkpoints"),
|
|
{
|
|
"forceInput": True,
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("JSON",)
|
|
RETURN_NAMES = ("json_out",)
|
|
FUNCTION = "action"
|
|
|
|
CATEGORY = "DarkPrompt"
|
|
|
|
def action(self, json_in, checkpoint):
|
|
checkpoint_name = clean_name(checkpoint)
|
|
print("json_in: %s" % (json_in))
|
|
|
|
if json_in:
|
|
try:
|
|
return (
|
|
apply_schema_defaults(
|
|
json_in[checkpoint_name],
|
|
schema.lora_settings_checkpoint_schema,
|
|
),
|
|
)
|
|
except KeyError as e:
|
|
print(
|
|
"Data does not exist for checkpoint %s: %s" % (checkpoint_name, e)
|
|
)
|
|
return (
|
|
apply_schema_defaults({}, schema.lora_settings_checkpoint_schema),
|
|
)
|
|
else:
|
|
return {}
|
|
|
|
|
|
class DarkLoraStackFromJSON(object):
|
|
"""
|
|
Takes a checkpoint name and a LoRA JSON object and creates a LoRA stack with the proper attributes
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"json_in": (
|
|
"JSON",
|
|
{
|
|
"forceInput": True,
|
|
},
|
|
),
|
|
"lora_name": (
|
|
"STRING",
|
|
{
|
|
"forceInput": True,
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"LORA_STACK",
|
|
"STRING",
|
|
)
|
|
RETURN_NAMES = ("LORA_STACK", "AS_STRING")
|
|
FUNCTION = "action"
|
|
|
|
CATEGORY = "DarkPrompt"
|
|
|
|
def action(self, json_in, lora_name):
|
|
lora_stack = []
|
|
logger.info(json_in)
|
|
|
|
if json_in and lora_name:
|
|
lora_stack.append(
|
|
(
|
|
unclean_name(lora_name),
|
|
float(json_in.get("model_weight", 0.8)),
|
|
float(json_in.get("clip_weight", 0)),
|
|
)
|
|
)
|
|
as_string = "<lora:%s:%s:%s>" % (
|
|
clean_name(lora_name),
|
|
json_in.get("model_weight", 0.8),
|
|
json_in.get("clip_weight", 0),
|
|
)
|
|
|
|
return (
|
|
lora_stack,
|
|
as_string,
|
|
)
|
|
else:
|
|
return (
|
|
[],
|
|
"",
|
|
)
|