From b10c317d4262c850fdac52aeae8e28de4162ab9e Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Mon, 15 May 2023 14:51:38 -0500 Subject: [PATCH] Split into extra settings --- do_run.py | 30 +++++++----- nodes.py | 142 ++++++++++++++++++++++++++++++++++-------------------- 2 files changed, 107 insertions(+), 65 deletions(-) diff --git a/do_run.py b/do_run.py index bce9a2e..41ff156 100644 --- a/do_run.py +++ b/do_run.py @@ -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') diff --git a/nodes.py b/nodes.py index 2c363a4..8103460 100644 --- a/nodes.py +++ b/nodes.py @@ -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", }