Fix loader depth

2024-06-18 23:56:10 +01:00
parent 5b082f424d
commit 91340dedc7
1 changed files with 2 additions and 0 deletions
@@ -76,6 +76,8 @@ def load_pixart(model_path, model_conf):
 		device=model_management.get_torch_device()
 	)

+	model_conf.unet_config['depth'] = sum(key.endswith('.scale_shift_table') for key in state_dict.keys())
+
 	if model_conf.model_target == "PixArtMS":
 		from .models.PixArtMS import PixArtMS
 		model.diffusion_model = PixArtMS(**model_conf.unet_config)