Add support for PixArt diffusers weights

PixArt/loader.py

@@ -5,8 +5,7 @@ import comfy.model_base
 import comfy.utils
 import torch
 from comfy import model_management
-
-from .models import PixArtMS
+from .diffusers_convert import convert_pixart_state_dict
 
 class EXM_PixArt(comfy.supported_models_base.BASE):
 	unet_config = {}
@@ -27,18 +26,18 @@ class EXM_PixArt(comfy.supported_models_base.BASE):
 def load_pixart(model_path, model_conf):
 	state_dict = comfy.utils.load_torch_file(model_path)
 	state_dict = state_dict.get("model", state_dict)
+	if "caption_projection.y_embedding" in state_dict:
+		state_dict = convert_pixart_state_dict(state_dict) # Diffusers
 	parameters = comfy.utils.calculate_parameters(state_dict)
 	unet_dtype = model_management.unet_dtype(model_params=parameters)
 
 	model_conf = EXM_PixArt(model_conf) # convert to object
-
 	model = comfy.model_base.BaseModel(
 		model_conf,
 		model_type=comfy.model_base.ModelType.EPS,
 		device=model_management.get_torch_device()
 	)
 
-	model.pixart_config = model_conf
 	if model_conf.model_target == "PixArtMS":
 		from .models.PixArtMS import PixArtMS
 		model.diffusion_model = PixArtMS(**model_conf.unet_config)
@@ -48,7 +47,9 @@ def load_pixart(model_path, model_conf):
 	else:
 		raise NotImplementedError(f"Unknown model target '{model_conf.model_target}'")
 
-	model.diffusion_model.load_state_dict(state_dict)
+	m, u = model.diffusion_model.load_state_dict(state_dict, strict=False)
+	if len(m) > 0: print("Missing UNET keys", m)
+	if len(u) > 0: print("Leftover UNET keys", u)
 	model.diffusion_model.dtype = unet_dtype
 	model.diffusion_model.eval()
 	model.diffusion_model.to(unet_dtype)