From 06256b62690b95778fb08c79159a73d952209e77 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Wed, 6 Sep 2023 19:15:29 +0200 Subject: [PATCH] Multi label + general jank This is probably not what the original paper had in mind for "interpolating between classes" but it mostly works. --- diffusion/gaussian_diffusion.py | 8 ++ models.py | 4 +- nodes.py | 125 +++++++++++++++++++++++++++----- 3 files changed, 116 insertions(+), 21 deletions(-) diff --git a/diffusion/gaussian_diffusion.py b/diffusion/gaussian_diffusion.py index 11edd7b..9cf4d90 100644 --- a/diffusion/gaussian_diffusion.py +++ b/diffusion/gaussian_diffusion.py @@ -492,6 +492,12 @@ class GaussianDiffusion: img = th.randn(*shape, device=device) indices = list(range(self.num_timesteps))[::-1] + y_inter = None + if "y_inter" in model_kwargs.keys(): + if len(model_kwargs["y_inter"]) > 1: + y_inter = model_kwargs["y_inter"] + del model_kwargs["y_inter"] + for i in indices: t = th.tensor([i] * shape[0], device=device) with th.no_grad(): @@ -506,6 +512,8 @@ class GaussianDiffusion: ) yield out img = out["sample"] + if y_inter: model_kwargs["y"] = y_inter[indices.index(i)%len(y_inter)] + if pbar: preview_bytes = None if previewer: diff --git a/models.py b/models.py index c90eeba..ab57d68 100644 --- a/models.py +++ b/models.py @@ -258,8 +258,8 @@ class DiT(nn.Module): # For exact reproducibility reasons, we apply classifier-free guidance on only # three channels by default. The standard approach to cfg applies it to all channels. # This can be done by uncommenting the following line and commenting-out the line following that. - # eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:] - eps, rest = model_out[:, :3], model_out[:, 3:] + eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:] + # eps, rest = model_out[:, :3], model_out[:, 3:] cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0) half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps) eps = torch.cat([half_eps, half_eps], dim=0) diff --git a/nodes.py b/nodes.py index df4ee3e..6c41200 100644 --- a/nodes.py +++ b/nodes.py @@ -17,7 +17,11 @@ folder_paths.folder_names_and_paths["dit"] = ( folder_paths.supported_pt_extensions ) + class DiTCheckpointLoader: + """ + Model loader with all possible options exposed. + """ @classmethod def INPUT_TYPES(s): return { @@ -35,20 +39,29 @@ class DiTCheckpointLoader: TITLE = "DiTCheckpointLoader" def load_checkpoint(self, ckpt_name, model, image_size, num_classes): - # note: switch to custom comfy.model_base eventually - model = DiT_models[model]( - input_size=image_size // 8, # latent size - num_classes=num_classes - ) - ckpt_path = folder_paths.get_full_path("dit", ckpt_name) state_dict = comfy.utils.load_torch_file(ckpt_path) + if "model" in state_dict.keys(): + state_dict = state_dict["model"] + dit = self.load_dit( + dit_model = DiT_models[model], + state_dict = state_dict, + latent_size = image_size // 8, + num_classes = num_classes, + ) + return (dit,) + + def load_dit(self, dit_model, state_dict, latent_size, num_classes): + model = dit_model( + input_size = latent_size, + num_classes = num_classes, + ) model.load_state_dict(state_dict) model.eval() # important, apparently # need these later anyway model.latent_format = comfy.latent_formats.SD15() - model.latent_size = image_size // 8 + model.latent_size = latent_size model.num_classes = num_classes # I didn't expect this to work but it looks like it does. @@ -58,12 +71,51 @@ class DiTCheckpointLoader: offload_device=comfy.model_management.unet_offload_device(), current_device="cpu" ) - - # return (model,) - return (model_patcher,) + return model_patcher -# todo: this needs fontend code to display properly +class DiTCheckpointLoaderSimple(DiTCheckpointLoader): + """ + Auto model loader. + To do: + - get image_size from pos_embed somehow + - guess model type from state_dict + """ + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "ckpt_name": (folder_paths.get_filename_list("dit"),), + "model": (list(DiT_models.keys()),), + "image_size": ([256, 512],), + } + } + TITLE = "DiTCheckpointLoaderSimple" + + def load_checkpoint(self, ckpt_name, model, image_size): + ckpt_path = folder_paths.get_full_path("dit", ckpt_name) + state_dict = comfy.utils.load_torch_file(ckpt_path) + if "model" in state_dict.keys(): + state_dict = state_dict["model"] + + num_classes, hidden_size = state_dict["y_embedder.embedding_table.weight"].shape + num_classes -= 1 # adj. for empty + + print("num_classes",num_classes) + print("hidden_size",hidden_size) + + latent_size = image_size // 8 + + dit = self.load_dit( + dit_model = DiT_models[model], + state_dict = state_dict, + latent_size = latent_size, + num_classes = num_classes, + ) + return (dit,) + + +# todo: this needs frontend code to display properly def get_label_data(label_file="labels/imagenet1000.json"): label_path = os.path.join( os.path.dirname(os.path.realpath(__file__)), @@ -93,9 +145,28 @@ class DiTLabelSelect: def label(self, label_name): global label_data - class_label = int([k for k,v in label_data.items() if v == label_name][0]) - return (class_label,) + class_labels = [int(k) for k,v in label_data.items() if v == label_name] + return (class_labels,) +class DiTLabelCombine: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "class_labels_a": ("DITLAB",), + "class_labels_b": ("DITLAB",), + } + } + + RETURN_TYPES = ("DITLAB",) + RETURN_NAMES = ("class_labels",) + FUNCTION = "label" + CATEGORY = "DiT" + TITLE = "DiTLabelCombine" + + def label(self, class_labels_a, class_labels_b): + class_labels = class_labels_a + class_labels_b + return (class_labels,) class DiTSampler: @classmethod @@ -117,7 +188,7 @@ class DiTSampler: TITLE = "DiTSampler" def sample(self, model, class_labels, latent_image, seed, steps, cfg, denoise): - device = comfy.model_management.get_torch_device() + device = model.load_device diffusion = create_diffusion(str(steps)) # pre @@ -132,17 +203,29 @@ class DiTSampler: zl = latent_image["samples"].to(device) zr = torch.randn(batch_size, 4, real_model.latent_size, real_model.latent_size, device=device) z = torch.lerp(zl,zr,denoise) # this is wrong - y = torch.tensor([class_labels] * batch_size, device=device) + + y_inter = [] + y_null = torch.tensor([real_model.num_classes] * batch_size, device=device) + for cl in class_labels: + cl = min(cl, real_model.num_classes) + y = torch.tensor([cl] * batch_size, device=device) + y = torch.cat([y, y_null], 0) + y_inter.append(y) # Setup classifier-free guidance: z = torch.cat([z, z], 0) - y_null = torch.tensor([real_model.num_classes] * batch_size, device=device) - y = torch.cat([y, y_null], 0) - model_kwargs = dict(y=y, cfg_scale=cfg) + model_kwargs = dict(y=y_inter[0], y_inter=y_inter, cfg_scale=cfg) # Sample images: samples = diffusion.p_sample_loop( - model.model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, pbar=pbar, previewer=previewer, device=device + model.model.forward_with_cfg, + z.shape, + z, + clip_denoised=False, + model_kwargs=model_kwargs, + pbar=pbar, + previewer=previewer, + device=device, ) samples, _ = samples.chunk(2, dim=0) # Remove null class samples samples = real_model.latent_format.process_out(samples.to(torch.float32)) @@ -151,13 +234,17 @@ class DiTSampler: return ({"samples": samples},) NODE_CLASS_MAPPINGS = { + "DiTCheckpointLoaderSimple": DiTCheckpointLoaderSimple, "DiTCheckpointLoader": DiTCheckpointLoader, + "DiTLabelCombine": DiTLabelCombine, "DiTLabelSelect": DiTLabelSelect, "DiTSampler": DiTSampler, } NODE_DISPLAY_NAME_MAPPINGS = { + "DiTCheckpointLoaderSimple": DiTCheckpointLoaderSimple.TITLE, "DiTCheckpointLoader": DiTCheckpointLoader.TITLE, + "DiTLabelCombine": DiTLabelCombine.TITLE, "DiTLabelSelect": DiTLabelSelect.TITLE, "DiTSampler": DiTSampler.TITLE, }