DiT Support

DiT/loader.py

@@ -0,0 +1,45 @@
+import comfy.supported_models_base
+import comfy.supported_models
+import comfy.latent_formats
+import comfy.model_patcher
+import comfy.model_base
+import comfy.utils
+import torch
+from comfy import model_management
+
+from .model import DiT
+
+class EXMDiT(comfy.supported_models.SD15):
+	unet_config = {}
+	unet_extra_config = {}
+
+	def model_type(self, state_dict, prefix=""):
+		return comfy.model_base.ModelType.EPS
+
+def load_dit(model_path, model_conf):
+	state_dict = comfy.utils.load_torch_file(model_path)
+	state_dict = state_dict.get("model", state_dict)
+	parameters = comfy.utils.calculate_parameters(state_dict)
+	unet_dtype = model_management.unet_dtype(model_params=parameters)
+
+	offload_device = model_management.unet_offload_device()
+	model = comfy.model_base.BaseModel(
+		EXMDiT({"disable_unet_model_creation" : True }),
+		model_type=comfy.model_base.ModelType.EPS,
+		device=model_management.get_torch_device()
+	)
+	model_conf["num_classes"] = state_dict["y_embedder.embedding_table.weight"].shape[0] - 1 # adj. for empty
+	model.dit_config = model_conf
+	model.diffusion_model = DiT(**model_conf).eval()
+	model.diffusion_model.load_state_dict(state_dict)
+	model.diffusion_model.eval()
+	model.diffusion_model.dtype = unet_dtype
+	model.diffusion_model.to(unet_dtype)
+
+	model_patcher = comfy.model_patcher.ModelPatcher(
+		model,
+		load_device=comfy.model_management.get_torch_device(),
+		offload_device=comfy.model_management.unet_offload_device(),
+		current_device="cpu"
+	)
+	return model_patcher