Fix setting scaling factor in VAE config (#6779)

fix
huggingface · Jan 31, 2024 · d7dc0ff · patrickvonplaten · Feb 12, 2024 · d7dc0ff
1 parent 97ee616
commit d7dc0ff
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/diffusers/loaders/single_file.py b/src/diffusers/loaders/single_file.py
@@ -61,8 +61,9 @@ def build_sub_model_components(
         return unet_components
 
     if component_name == "vae":
+        scaling_factor = kwargs.get("scaling_factor", None)
         vae_components = create_diffusers_vae_model_from_ldm(
-            pipeline_class_name, original_config, checkpoint, image_size
+            pipeline_class_name, original_config, checkpoint, image_size, scaling_factor
         )
         return vae_components
 

diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py
@@ -513,11 +513,12 @@ def create_controlnet_diffusers_config(original_config, image_size: int):
     return controlnet_config
 
 
-def create_vae_diffusers_config(original_config, image_size, scaling_factor=0.18125):
+def create_vae_diffusers_config(original_config, image_size, scaling_factor=None):
     """
     Creates a config for the diffusers based on the config of the LDM model.
     """
     vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"]
+    scaling_factor = scaling_factor or original_config["model"]["params"]["scale_factor"]
 
     block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]]
     down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels)