nodes/space-nuko-ComfyUI-Disco-Diffusion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Split into extra settings

Browse Source
This commit is contained in:
space-nuko
2023-05-15 14:51:38 -05:00
parent 92ab503b71
commit b10c317d42
2 changed files with 107 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
do_run.py
+17 -13
View File
@@ -560,19 +560,19 @@ def run_one_frame(diffusion, model, clip_model, clip_vision, args, batchNum, fra
# tqdm.write(f'Batch {i}, step {j}, output {k}:')
datetime.now().strftime('%y%m%d-%H%M%S_%f')
percent = math.ceil(j/total_steps*100)
if args.n_batches > 0:
# if intermediates are saved to the subfolder, don't append a step or percentage to the name
if cur_t == -1 and args.intermediates_in_subfolder is True:
save_num = f'{frame_num:04}' if args.animation_mode != "None" else i
filename = f'{args.batch_name}({batchNum})_{save_num}.png'
else:
# If we're working with percentages, append it
if args.steps_per_checkpoint is not None:
filename = f'{args.batch_name}({batchNum})_{i:04}-{percent:02}%.png'
# Or else, iIf we're working with specific steps, append those
else:
filename = f'{args.batch_name}({batchNum})_{i:04}-{j:03}.png'
save_image(image, j, cur_t, filename, frame_num, midas_model, midas_transform, args)
# if args.n_batches > 0:
# # if intermediates are saved to the subfolder, don't append a step or percentage to the name
# if cur_t == -1 and args.intermediates_in_subfolder is True:
# save_num = f'{frame_num:04}' if args.animation_mode != "None" else i
# filename = f'{args.batch_name}({batchNum})_{save_num}.png'
# else:
# # If we're working with percentages, append it
# if args.steps_per_checkpoint is not None:
# filename = f'{args.batch_name}({batchNum})_{i:04}-{percent:02}%.png'
# # Or else, iIf we're working with specific steps, append those
# else:
# filename = f'{args.batch_name}({batchNum})_{i:04}-{j:03}.png'
# save_image(image, j, cur_t, filename, frame_num, midas_model, midas_transform, args)
if cur_t == -1:
# We get back a tensor of size [C, H, W].
@@ -580,7 +580,11 @@ def run_one_frame(diffusion, model, clip_model, clip_vision, args, batchNum, fra
# So... Let's Transposing!
image = image.permute(1, 2, 0).add(1).div(2).clamp(0, 1)
# image = image.add(1).div(2).clamp(0, 1)
# [[H, W, C]] -> [B, H, W, C]
# All results will be wrapped in a Python list for use with OUTPUT_IS_LIST.
# B will always be 1 for each individual image tensor.
# Yes this is weird.
results.append(torch.stack([image]))
# plt.plot(np.array(loss_values), 'r')
nodes.py
+90 -52
View File
@@ -98,51 +98,86 @@ class GuidedDiffusionLoader:
return (model_settings,)
class DiscoDiffusionExtraSettings:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"eta": ("FLOAT", { "default": 0.8, "min": 0, "max": 100 }),
"cutn": ("INT", { "default": 16, "min": 1, "max": 32 }),
"cutn_batches": ("INT", { "default": 2, "min": 1, "max": 16 }),
"cut_overview": ("STRING", { "default": "[12]*400+[4]*600" }),
"cut_innercut": ("STRING", { "default": "[4]*400+[12]*600" }),
"cut_ic_pow": ("STRING", { "default": "[1]*1000" }),
"cut_icgray_p": ("STRING", { "default": "[0.2]*400+[0]*600" }),
}}
# These are technically different model formats so don't use them with vanilla nodes!
RETURN_TYPES = ("DISCO_DIFFUSION_EXTRA_SETTINGS",)
FUNCTION = "make_settings"
CATEGORY = "sampling"
def __init__(self):
pass
def make_settings(self, eta, cutn, cutn_batches, cut_overview, cut_innercut, cut_ic_pow, cut_icgray_p):
extra_settings = {
"eta": eta,
"cutn": cutn,
"cutn_batches": cutn_batches,
"cut_overview": cut_overview,
"cut_innercut": cut_innercut,
"cut_ic_pow": cut_ic_pow,
"cut_icgray_p": cut_icgray_p
}
return (extra_settings,)
DEFAULT_PROMPT = """\
# How to prompt:
# Each line is prefixed with the starting frame number of the prompt.
# More than one line with the same frame number concatenates the two prompts together.
# Each individual prompt can be no more than 77 characters long.
# Weights are parsed from the end of each prompt with "25:a fluffy fox:5" syntax
# Comments are written with the '#' character. Blank lines are ignored.
; How to prompt:
; Each line is prefixed with the starting frame number of the prompt.
; More than one line with the same frame number concatenates the two prompts together.
; Each individual prompt can be no more than 77 characters long.
; Weights are parsed from the end of each prompt with "25:a fluffy fox:5" syntax
; Comments are written with the ';' character. Blank lines are ignored.
0:A beautiful painting of a singular lighthouse, shining its light across a tumultuous sea of blood by greg rutkowski and thomas kinkade. Trending on artstation.
0:yellow color scheme
#100:This set of prompts start at frame 100.
#100:This prompt has weight five:5
;100:This set of prompts start at frame 100.
;100:This prompt has weight five:5
""".strip()
class DiscoDiffusion:
@classmethod
def INPUT_TYPES(s):
return {"required": {"text": ("STRING", {"default": DEFAULT_PROMPT, "multiline": True}),
"guided_diffusion": ("GUIDED_DIFFUSION_MODEL",),
"clip": ("CLIP",),
"clip_vision": ("CLIP_VISION",),
# Sane defaults:
# 1280x768 for 512x512 models
# 512x448 for 256x256 models
"width": ("INT", {"default": 1280, "min": 64, "max": 2048, "step": 64}),
"height": ("INT", {"default": 768, "min": 64, "max": 2048, "step": 64}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"steps": ("INT", {"default": 250, "min": 1, "max": 10000}),
"skip_steps": ("INT", {"default": 10, "min": 1, "max": 10000}),
"n_batches": ("INT", {"default": 1, "min": 1, "max": 16}),
# "max_frames": ("INT", {"default": 1, "min": 1, "max": 1000}),
"sampling_mode": (["plms", "ddim", "stsp", "ltsp"], {"default": "ddim"}),
"clip_guidance_scale": ("INT", { "default": 5000, "min": 1, "max": 10000000 }),
"tv_scale": ("INT", { "default": 0, "min": 0, "max": 100000 }),
"range_scale": ("INT", { "default": 150, "min": 0, "max": 100000 }),
"sat_scale": ("INT", { "default": 0, "min": 0, "max": 100000 }),
"eta": ("FLOAT", { "default": 0.8, "min": 0, "max": 100 }),
"cutn": ("INT", { "default": 16, "min": 1, "max": 32 }),
"cutn_batches": ("INT", { "default": 2, "min": 1, "max": 16 }),
"cut_overview": ("STRING", { "default": "[12]*400+[4]*600" }),
"cut_innercut": ("STRING", { "default": "[4]*400+[12]*600" }),
"cut_ic_pow": ("STRING", { "default": "[1]*1000" }),
"cut_icgray_p": ("STRING", { "default": "[0.2]*400+[0]*600" }),
}}
return {
"required": {
"text": ("STRING", {"default": DEFAULT_PROMPT, "multiline": True}),
"guided_diffusion": ("GUIDED_DIFFUSION_MODEL",),
"clip": ("CLIP",),
"clip_vision": ("CLIP_VISION",),
# Sane defaults:
# 1280x768 for 512x512 models
# 512x448 for 256x256 models
"width": ("INT", {"default": 1280, "min": 64, "max": 2048, "step": 64}),
"height": ("INT", {"default": 768, "min": 64, "max": 2048, "step": 64}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"steps": ("INT", {"default": 250, "min": 1, "max": 10000}),
"skip_steps": ("INT", {"default": 10, "min": 1, "max": 10000}),
"n_batches": ("INT", {"default": 1, "min": 1, "max": 16}),
# "max_frames": ("INT", {"default": 1, "min": 1, "max": 1000}),
"sampling_mode": (["plms", "ddim", "stsp", "ltsp"], {"default": "ddim"}),
"clip_guidance_scale": ("INT", { "default": 5000, "min": 1, "max": 10000000 }),
"tv_scale": ("INT", { "default": 0, "min": 0, "max": 100000 }),
"range_scale": ("INT", { "default": 150, "min": 0, "max": 100000 }),
"sat_scale": ("INT", { "default": 0, "min": 0, "max": 100000 }),
},
"optional": {
"extra_settings": ("DISCO_DIFFUSION_EXTRA_SETTINGS",),
}
}
RETURN_TYPES = ("IMAGE",)
OUTPUT_IS_LIST = (True,)
FUNCTION = "generate"
@@ -155,7 +190,7 @@ class DiscoDiffusion:
def parse_prompts(self, text):
result = {}
for line in text.split('\n'):
line = line.split('#')[0].strip()
line = line.split(';')[0].strip()
if line:
if ':' in line:
vals = line.split(':', 2)
@@ -199,7 +234,7 @@ class DiscoDiffusion:
model.load_state_dict(torch.load(model_settings.custom_path, map_location='cpu'))
else:
model.load_state_dict(torch.load(f'{model_settings.model_path}/{model_settings.get_model_filename(model_settings.diffusion_model)}', map_location='cpu'))
model.requires_grad_(False).eval().to(device)
model.requires_grad_(False).eval().to(device)
for name, param in model.named_parameters():
if 'qkv' in name or 'norm' in name or 'proj' in name:
@@ -211,8 +246,7 @@ class DiscoDiffusion:
return model, diffusion
def generate(self, text, guided_diffusion, clip, clip_vision, width, height, seed, steps, skip_steps, n_batches, sampling_mode,
clip_guidance_scale, tv_scale, range_scale, sat_scale, eta,
cutn, cutn_batches, cut_overview, cut_innercut, cut_ic_pow, cut_icgray_p):
clip_guidance_scale, tv_scale, range_scale, sat_scale, extra_settings=None):
settings = DiscoDiffusionSettings()
settings.seed = seed
settings.steps = steps
@@ -226,15 +260,17 @@ class DiscoDiffusion:
settings.tv_scale = tv_scale
settings.range_scale = range_scale
settings.sat_scale = sat_scale
settings.eta = eta
settings.cutn = cutn
settings.cutn_batches = cutn_batches
settings.cut_overview = cut_overview
settings.cut_innercut = cut_innercut
settings.cut_ic_pow = cut_ic_pow
settings.cut_icgray_p = cut_icgray_p
guided_diffusion.diffusion_sampling_mode = sampling_mode
if extra_settings is not None:
settings.eta = extra_settings["eta"]
settings.cutn = extra_settings["cutn"]
settings.cutn_batches = extra_settings["cutn_batches"]
settings.cut_overview = extra_settings["cut_overview"]
settings.cut_innercut = extra_settings["cut_innercut"]
settings.cut_ic_pow = extra_settings["cut_ic_pow"]
settings.cut_icgray_p = extra_settings["cut_icgray_p"]
print("[Disco Diffusion] Parsed Prompts:")
pp(settings.text_prompts)
@@ -250,13 +286,15 @@ class DiscoDiffusion:
NODE_CLASS_MAPPINGS = {
"ComfyUI_OpenAICLIPLoader": OpenAICLIPLoader,
"ComfyUI_GuidedDiffusionLoader": GuidedDiffusionLoader,
"ComfyUI_DiscoDiffusion": DiscoDiffusion,
"DiscoDiffusion_OpenAICLIPLoader": OpenAICLIPLoader,
"DiscoDiffusion_GuidedDiffusionLoader": GuidedDiffusionLoader,
"DiscoDiffusion_DiscoDiffusion": DiscoDiffusion,
"DiscoDiffusion_DiscoDiffusionExtraSettings": DiscoDiffusionExtraSettings,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ComfyUI_OpenAICLIPLoader": "OpenAI CLIP Loader",
"ComfyUI_GuidedDiffusionLoader": "Guided Diffusion Loader",
"ComfyUI_DiscoDiffusion": "Disco Diffusion",
"DiscoDiffusion_OpenAICLIPLoader": "OpenAI CLIP Loader",
"DiscoDiffusion_GuidedDiffusionLoader": "Guided Diffusion Loader",
"DiscoDiffusion_DiscoDiffusion": "Disco Diffusion",
"DiscoDiffusion_DiscoDiffusionExtraSettings": "Disco Diffusion Extra Settings",
}