DiT Support

DiT/nodes.py

@@ -0,0 +1,109 @@
+import os
+import json
+import torch
+import folder_paths
+
+from .conf import dit_conf
+from .loader import load_dit
+
+# initialize custom folder path
+# TODO: integrate with `extra_model_paths.yaml`
+os.makedirs(
+	os.path.join(folder_paths.models_dir,"dit"),
+	exist_ok = True,
+)
+folder_paths.folder_names_and_paths["dit"] = (
+	[os.path.join(folder_paths.models_dir,"dit")],
+	folder_paths.supported_pt_extensions
+)
+
+class DitCheckpointLoader:
+	@classmethod
+	def INPUT_TYPES(s):
+		return {
+			"required": {
+				"ckpt_name": (folder_paths.get_filename_list("dit"),),
+				"model": (list(dit_conf.keys()),),
+				"image_size": ([256, 512],),
+				# "num_classes": ("INT", {"default": 1000, "min": 0,}),
+			}
+		}
+	RETURN_TYPES = ("MODEL",)
+	RETURN_NAMES = ("model",)
+	FUNCTION = "load_checkpoint"
+	CATEGORY = "ExtraModels/DiT"
+	TITLE = "DitCheckpointLoader"
+
+	def load_checkpoint(self, ckpt_name, model, image_size):
+		ckpt_path = folder_paths.get_full_path("dit", ckpt_name)
+		model_conf = dit_conf[model]
+		model_conf["input_size"]  = image_size // 8
+		# model_conf["num_classes"] = num_classes
+		dit = load_dit(
+			model_path = ckpt_path,
+			model_conf = model_conf,
+		)
+		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__)),
+		label_file,
+	)
+	label_data = {0: "None"}
+	with open(label_path, "r") as f:
+		label_data = json.loads(f.read())
+	return label_data
+label_data = get_label_data()
+
+class DiTCondLabelSelect:
+	@classmethod
+	def INPUT_TYPES(s):
+		global label_data
+		return {
+			"required": {
+				"model" : ("MODEL",),
+				"label_name": (list(label_data.values()),),
+			}
+		}
+
+	RETURN_TYPES = ("CONDITIONING",)
+	RETURN_NAMES = ("class",)
+	FUNCTION = "cond_label"
+	CATEGORY = "ExtraModels/DiT"
+	TITLE = "DiTCondLabelSelect"
+
+	def cond_label(self, model, label_name):
+		global label_data
+		class_labels = [int(k) for k,v in label_data.items() if v == label_name]
+		y = torch.tensor([[class_labels[0]]]).to(torch.int)
+		return ([[y, {"pooled_output": []}]], )
+
+class DiTCondLabelEmpty:
+	@classmethod
+	def INPUT_TYPES(s):
+		global label_data
+		return {
+			"required": {
+				"model" : ("MODEL",),
+			}
+		}
+
+	RETURN_TYPES = ("CONDITIONING",)
+	RETURN_NAMES = ("empty",)
+	FUNCTION = "cond_empty"
+	CATEGORY = "ExtraModels/DiT"
+	TITLE = "DiTCondLabelEmpty"
+
+	def cond_empty(self, model):
+		# [ID of last class + 1] == [num_classes]
+		y_null = model.model.dit_config["num_classes"]
+		y = torch.tensor([[y_null]]).to(torch.int)
+		return ([[y, {"pooled_output": []}]], )
+
+NODE_CLASS_MAPPINGS = {
+	"DitCheckpointLoader" : DitCheckpointLoader,
+	"DiTCondLabelSelect"  : DiTCondLabelSelect,
+	"DiTCondLabelEmpty"   : DiTCondLabelEmpty,
+}