46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
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
|