323 lines
11 KiB
Python
323 lines
11 KiB
Python
from pathlib import Path
|
|
from .utils.darkdata import DarkData
|
|
from .prompts import strip_comments_from_lines
|
|
import folder_paths
|
|
import glob
|
|
import logging
|
|
import os
|
|
import os.path
|
|
import random
|
|
import re
|
|
import json
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DarkCheckpointSwitcher(object):
|
|
"""
|
|
Takes checkpoints as an input and allows you to manually switch between them
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"checkpoint1": (
|
|
folder_paths.get_filename_list("checkpoints"),
|
|
{
|
|
"forceInput": True,
|
|
},
|
|
),
|
|
"checkpoint2": (
|
|
folder_paths.get_filename_list("checkpoints"),
|
|
{
|
|
"forceInput": True,
|
|
},
|
|
),
|
|
"use_checkpoint": (
|
|
"INT",
|
|
{
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 2,
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (folder_paths.get_filename_list("checkpoints"),)
|
|
RETURN_NAMES = ("ckpt_name",)
|
|
FUNCTION = "get_switched_checkpoint"
|
|
|
|
CATEGORY = "DarkPrompt"
|
|
|
|
def get_switched_checkpoint(self, checkpoint1, checkpoint2, use_checkpoint):
|
|
if use_checkpoint == 1:
|
|
return (checkpoint1,)
|
|
else:
|
|
return (checkpoint2,)
|
|
|
|
|
|
class DarkCheckpointRandomizer(object):
|
|
"""
|
|
Handles randomizing checkpoints
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xFFFFFFFFFFFFFFFF,
|
|
"forceInput": True,
|
|
},
|
|
),
|
|
"use_for_iterations": ("INT", {"default": 10, "min": 1}),
|
|
"checkpoint_names": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (folder_paths.get_filename_list("checkpoints"),)
|
|
RETURN_NAMES = ("ckpt_name",)
|
|
FUNCTION = "get_checkpoint"
|
|
|
|
CATEGORY = "DarkPrompt"
|
|
|
|
def get_checkpoint(self, seed, use_for_iterations, checkpoint_names):
|
|
# In order to avoid random crashes, we need to validate that every file
|
|
# exists. This prevents someone from adding 'doesnotexist.safetensors'
|
|
# in to the random list queueing up a bunch of jobs and coming back the
|
|
# next morning to find that at some point we tried to use
|
|
# doesnotexist.safetensors and everything crashed.
|
|
checkpoints = []
|
|
checkpoint = None
|
|
checkpoint_lines = strip_comments_from_lines(checkpoint_names.splitlines())
|
|
|
|
for cpn in checkpoint_lines:
|
|
if cpn.strip():
|
|
if cpn.strip() in folder_paths.get_filename_list("checkpoints"):
|
|
checkpoints.append(cpn.strip())
|
|
else:
|
|
logger.warn(
|
|
"%s is in your DarkCheckpointRandomizer but it does not exist in your checkpoints directory"
|
|
% (cpn)
|
|
)
|
|
|
|
if not checkpoints:
|
|
raise Exception(
|
|
"You have no checkpoints that exist listed in your DarkCheckpointRandomizer"
|
|
)
|
|
|
|
with DarkData(filename="darkcheckpointrandomizer.json") as DFB:
|
|
random.seed(seed)
|
|
# make sure we have a checkpoint from the file or if it doesn't
|
|
# exist from the random function
|
|
checkpoint = DFB.get_key("checkpoint", random.choice(checkpoints))
|
|
|
|
if (
|
|
DFB.get_key("iteration", 0) >= use_for_iterations
|
|
or checkpoint not in checkpoints
|
|
):
|
|
# If we are over our iteration count or the checkpoint no
|
|
# longer exists in the list due to the user removing it, set
|
|
# the iteration back to 1, get a new checkpoint, and save it
|
|
DFB.set_key("iteration", 1)
|
|
checkpoint = random.choice(checkpoints)
|
|
else:
|
|
# If we aren't over our iteration count (or maybe this is our
|
|
# first iteration, increment the iteration counter
|
|
DFB.set_key("iteration", DFB.get_key("iteration", 0) + 1)
|
|
|
|
# Set our current checkpoint as it may have changed
|
|
DFB.set_key("checkpoint", checkpoint)
|
|
|
|
print("Checkpoint: %s" % (checkpoint))
|
|
|
|
return (checkpoint,)
|
|
|
|
|
|
class DarkCheckpointRotator(object):
|
|
"""
|
|
Steps through a checkpoint every iteration until the list is empty then starts again.
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xFFFFFFFFFFFFFFFF,
|
|
"forceInput": True,
|
|
},
|
|
),
|
|
"checkpoint_names": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (folder_paths.get_filename_list("checkpoints"),)
|
|
RETURN_NAMES = ("ckpt_name",)
|
|
FUNCTION = "action"
|
|
|
|
CATEGORY = "DarkPrompt"
|
|
|
|
def action(self, seed, checkpoint_names):
|
|
checkpoints = []
|
|
checkpoint = None
|
|
checkpoint_lines = strip_comments_from_lines(checkpoint_names.splitlines())
|
|
|
|
for cpn in checkpoint_lines:
|
|
if cpn.strip():
|
|
if cpn.strip() in folder_paths.get_filename_list("checkpoints"):
|
|
checkpoints.append(cpn.strip())
|
|
else:
|
|
logger.warn(
|
|
"%s is in your DarkCheckpointRotator but it does not exist in your checkpoints directory"
|
|
% (cpn)
|
|
)
|
|
|
|
if not checkpoints:
|
|
raise Exception(
|
|
"You have no checkpoints that exist listed in your DarkCheckpointRotator"
|
|
)
|
|
|
|
with DarkData(filename="darkcheckpointrotator.json") as DFB:
|
|
used_checkpoints = []
|
|
used_checkpoints_string = DFB.get_key("used_checkpoints", None)
|
|
if used_checkpoints_string:
|
|
used_checkpoints = used_checkpoints_string.split(",")
|
|
|
|
print("Used checkpoints: %s" % (used_checkpoints))
|
|
|
|
available_checkpoints = [
|
|
x for x in checkpoints if x not in used_checkpoints
|
|
]
|
|
print("Available checkpoints: %s" % (available_checkpoints))
|
|
|
|
if not available_checkpoints:
|
|
# We used 'em all up
|
|
available_checkpoints = checkpoints
|
|
used_checkpoints = []
|
|
|
|
checkpoint = available_checkpoints.pop()
|
|
used_checkpoints.append(checkpoint)
|
|
DFB.set_key("used_checkpoints", ",".join(used_checkpoints))
|
|
|
|
print("Checkpoint: %s" % (checkpoint))
|
|
|
|
return (checkpoint,)
|
|
|
|
|
|
class DarkCheckpointRotatorIterations(object):
|
|
"""
|
|
Returns a checkpoint for a number of iterations and then continues on to the next one until the entire list is empty. Then we start again.
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xFFFFFFFFFFFFFFFF,
|
|
"forceInput": True,
|
|
},
|
|
),
|
|
"use_for_iterations": ("INT", {"default": 10, "min": 1}),
|
|
"checkpoint_names": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (folder_paths.get_filename_list("checkpoints"),)
|
|
RETURN_NAMES = ("ckpt_name",)
|
|
FUNCTION = "action"
|
|
|
|
CATEGORY = "DarkPrompt"
|
|
|
|
def action(self, seed, use_for_iterations, checkpoint_names):
|
|
checkpoints = []
|
|
checkpoint = None
|
|
checkpoint_lines = strip_comments_from_lines(checkpoint_names.splitlines())
|
|
|
|
for cpn in checkpoint_lines:
|
|
if cpn.strip():
|
|
if cpn.strip() in folder_paths.get_filename_list("checkpoints"):
|
|
checkpoints.append(cpn.strip())
|
|
else:
|
|
logger.warn(
|
|
"%s is in your DarkCheckpointRotator but it does not exist in your checkpoints directory"
|
|
% (cpn)
|
|
)
|
|
|
|
if not checkpoints:
|
|
raise Exception(
|
|
"You have no checkpoints that exist listed in your DarkCheckpointRotator"
|
|
)
|
|
|
|
with DarkData(filename="darkcheckpointrotator.json") as DFB:
|
|
iterations = DFB.get_key("iterations", 0)
|
|
checkpoint = DFB.get_key("checkpoint", None)
|
|
|
|
used_checkpoints = []
|
|
used_checkpoints_string = DFB.get_key("used_checkpoints", None)
|
|
if used_checkpoints_string:
|
|
used_checkpoints = used_checkpoints_string.split(",")
|
|
|
|
available_checkpoints = [
|
|
x for x in checkpoints if x not in used_checkpoints
|
|
]
|
|
|
|
if checkpoint and checkpoint not in available_checkpoints:
|
|
if available_checkpoints:
|
|
checkpoint = available_checkpoints.pop()
|
|
iterations = 0
|
|
else:
|
|
available_checkpoints = checkpoints
|
|
checkpoint = available_checkpoints.pop()
|
|
iterations = 0
|
|
elif not checkpoint:
|
|
checkpoint = available_checkpoints.pop()
|
|
iterations = 0
|
|
|
|
if iterations >= use_for_iterations:
|
|
used_checkpoints.append(checkpoint)
|
|
available_checkpoints.pop()
|
|
if available_checkpoints:
|
|
checkpoint = available_checkpoints.pop()
|
|
iterations = 0
|
|
else:
|
|
available_checkpoints = checkpoints
|
|
used_checkpoints = []
|
|
checkpoint = available_checkpoints.pop()
|
|
iterations = 0
|
|
|
|
iterations += 1
|
|
|
|
DFB.set_key("used_checkpoints", ",".join(used_checkpoints))
|
|
DFB.set_key("checkpoint", checkpoint)
|
|
DFB.set_key("iterations", iterations)
|
|
|
|
return (checkpoint,)
|