From ccf8d62c2298ff41d02ec8dd67739ac81e48cfb5 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Fri, 29 Dec 2023 14:24:27 +0000 Subject: [PATCH] update --- src/diffusers/loaders/single_file.py | 9 ++++- src/diffusers/loaders/single_file_utils.py | 45 +++++++++++++--------- 2 files changed, 34 insertions(+), 20 deletions(-) diff --git a/src/diffusers/loaders/single_file.py b/src/diffusers/loaders/single_file.py index cc211f0a4b..00fa73a852 100644 --- a/src/diffusers/loaders/single_file.py +++ b/src/diffusers/loaders/single_file.py @@ -339,7 +339,13 @@ class FromSingleFileMixin: pipeline_components = {} for component in component_names: components = build_component( - pipeline_components, pipeline_name, component, checkpoint, original_config, **kwargs + pipeline_components, + pipeline_name, + component, + original_config, + checkpoint, + pretrained_model_link_or_path, + **kwargs, ) pipeline_components.update(components) @@ -354,4 +360,3 @@ class FromSingleFileMixin: pipe.to(dtype=torch_dtype) return pipe - diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index 242bbc2d9d..e6ab616087 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -222,6 +222,23 @@ def get_default_scheduler_config(): return SCHEDULER_DEFAULT_CONFIG +def determine_image_size(pipeline_class_name, original_config, checkpoint, **kwargs): + image_size = kwargs.get("image_size", 512) + global_step = checkpoint["global_step"] if "global_step" in checkpoint else None + + if pipeline_class_name == "StableDiffusionUpscalePipeline": + image_size = image_size or original_config.model.params.unet_config.params.image_size + + elif ( + "parameterization" in original_config["model"]["params"] + and original_config["model"]["params"]["parameterization"] == "v" + ): + # NOTE: For stable diffusion 2 base one has to pass `image_size==512` + # as it relies on a brittle global step parameter here + image_size = 512 if global_step == 875000 else 768 + + return image_size + def shave_segments(path, n_shave_prefix_segments=1): """ Removes segments. Positive values shave the first segments, negative shave the last segments. @@ -1268,19 +1285,7 @@ def create_unet_model(pipeline_class_name, original_config, checkpoint, checkpoi else: num_in_channels = 4 - image_size = kwargs.get("image_size", 512) - global_step = checkpoint["global_step"] if "global_step" in checkpoint else None - - if pipeline_class_name == "StableDiffusionUpscalePipeline": - image_size = image_size or original_config.model.params.unet_config.params.image_size - - elif ( - "parameterization" in original_config["model"]["params"] - and original_config["model"]["params"]["parameterization"] == "v" - ): - # NOTE: For stable diffusion 2 base one has to pass `image_size==512` - # as it relies on a brittle global step parameter here - image_size = 512 if global_step == 875000 else 768 + image_size = determine_image_size(pipeline_class_name, original_config, checkpoint, **kwargs) upcast_attention = kwargs.get("upcast_attention", False) extract_ema = kwargs.get("extract_ema", False) @@ -1303,7 +1308,7 @@ def create_unet_model(pipeline_class_name, original_config, checkpoint, checkpoi else: unet.load_state_dict(diffusers_format_unet_checkpoint) - return unet + return {"unet": unet} def create_controlnet_model(pipeline_class_name, original_config, checkpoint, checkpoint_path_or_dict, image_size, **kwargs): @@ -1321,10 +1326,13 @@ def create_controlnet_model(pipeline_class_name, original_config, checkpoint, ch return {"controlnet": controlnet} -def create_vae_model(original_config, checkpoint, checkpoint_path_or_dict, **kwargs): - vae_config = create_vae_diffusers_config(original_config) +def create_vae_model(pipeline_class_name, original_config, checkpoint, checkpoint_path_or_dict, **kwargs): + image_size = determine_image_size(pipeline_class_name, original_config, checkpoint, **kwargs) + + vae_config = create_vae_diffusers_config(original_config, image_size=image_size) diffusers_format_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) ctx = init_empty_weights if is_accelerate_available() else nullcontext + with ctx(): vae = AutoencoderKL(**vae_config) @@ -1334,7 +1342,7 @@ def create_vae_model(original_config, checkpoint, checkpoint_path_or_dict, **kwa else: vae.load_state_dict(diffusers_format_vae_checkpoint) - return vae + return {"vae": vae} def create_text_encoders_and_tokenizers(pipeline_class_name, original_config, checkpoint, checkpoint_path_or_dict, **kwargs): @@ -1425,7 +1433,8 @@ def create_text_encoders_and_tokenizers(pipeline_class_name, original_config, ch local_files_only=local_files_only, **config_kwargs, ) - except Exception: + except Exception as e: + print(e) raise ValueError( f"With local_files_only set to {local_files_only}, you must first locally save the text_encoder_2 and tokenizer_2 in the following path: {config_name} with `pad_token` set to '!'." )