Fix DiT sampling

Apply the PixArt fix here as well.

Move models to checkpoints folder as that makes more sense in this case.

DiT/loader.py

@@ -1,5 +1,4 @@
 import comfy.supported_models_base
-import comfy.supported_models
 import comfy.latent_formats
 import comfy.model_patcher
 import comfy.model_base
@@ -7,11 +6,17 @@ import comfy.utils
 import torch
 from comfy import model_management
 
-from .model import DiT
-
-class EXMDiT(comfy.supported_models.SD15):
+class EXM_DiT(comfy.supported_models_base.BASE):
 	unet_config = {}
 	unet_extra_config = {}
+	latent_format = comfy.latent_formats.SD15
+
+	def __init__(self, model_conf):
+		self.unet_config = model_conf.get("unet_config", {})
+		self.sampling_settings = model_conf.get("sampling_settings", {})
+		self.latent_format = self.latent_format()
+		# UNET is handled by extension
+		self.unet_config["disable_unet_model_creation"] = True
 
 	def model_type(self, state_dict, prefix=""):
 		return comfy.model_base.ModelType.EPS
@@ -22,18 +27,21 @@ def load_dit(model_path, model_conf):
 	parameters = comfy.utils.calculate_parameters(state_dict)
 	unet_dtype = model_management.unet_dtype(model_params=parameters)
 
-	offload_device = model_management.unet_offload_device()
+	model_conf["unet_config"]["num_classes"] = state_dict["y_embedder.embedding_table.weight"].shape[0] - 1 # adj. for empty
+
+	model_conf = EXM_DiT(model_conf)
 	model = comfy.model_base.BaseModel(
-		EXMDiT({"disable_unet_model_creation" : True }),
+		model_conf,
 		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()
+
+	from .model import DiT
+	model.diffusion_model = DiT(**model_conf.unet_config)
+
 	model.diffusion_model.load_state_dict(state_dict)
-	model.diffusion_model.eval()
 	model.diffusion_model.dtype = unet_dtype
+	model.diffusion_model.eval()
 	model.diffusion_model.to(unet_dtype)
 
 	model_patcher = comfy.model_patcher.ModelPatcher(