From 91340dedc7d431608f1b6741cbc8e9b632f2e640 Mon Sep 17 00:00:00 2001 From: gchapman Date: Tue, 18 Jun 2024 23:56:10 +0100 Subject: [PATCH] Fix loader depth --- PixArt/loader.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/PixArt/loader.py b/PixArt/loader.py index 3e26711..e085d63 100644 --- a/PixArt/loader.py +++ b/PixArt/loader.py @@ -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)