Update docstrings

Author: corystephenson-db

diffusion/models/models.py

@@ -366,6 +366,9 @@ def stable_diffusion_xl(
         use_xformers (bool): Whether to use xformers for attention. Defaults to True.
         lora_rank (int, optional): If not None, the rank to use for LoRA finetuning. Defaults to None.
         lora_alpha (int, optional): If not None, the alpha to use for LoRA finetuning. Defaults to None.
+        cache_dir (str): Directory to cache local files in. Default: `'/tmp/hf_files'`.
+        local_files_only (bool): Whether to only use local files. Default: `False`.
+
     """
     latent_mean, latent_std = _parse_latent_statistics(latent_mean), _parse_latent_statistics(latent_std)
 
@@ -379,7 +382,9 @@ def stable_diffusion_xl(
         val_metrics = [MeanSquaredError()]
 
     # Make the tokenizer and text encoder
-    tokenizer = MultiTokenizer(tokenizer_names_or_paths=tokenizer_names, cache_dir=cache_dir, local_files_only=local_files_only)
+    tokenizer = MultiTokenizer(tokenizer_names_or_paths=tokenizer_names,
+                               cache_dir=cache_dir,
+                               local_files_only=local_files_only)
     text_encoder = MultiTextEncoder(model_names=text_encoder_names,
                                     encode_latents_in_fp16=encode_latents_in_fp16,
                                     pretrained_sdxl=pretrained,
@@ -412,9 +417,15 @@ def stable_diffusion_xl(
         downsample_factor = 2**(len(vae.config['channel_multipliers']) - 1)
 
     # Make the unet
-    unet_config = PretrainedConfig.get_config_dict(unet_model_name, subfolder='unet', cache_dir=cache_dir,  local_files_only=local_files_only)[0]
+    unet_config = PretrainedConfig.get_config_dict(unet_model_name,
+                                                   subfolder='unet',
+                                                   cache_dir=cache_dir,
+                                                   local_files_only=local_files_only)[0]
     if pretrained:
-        unet = UNet2DConditionModel.from_pretrained(unet_model_name, subfolder='unet', cache_dir=cache_dir, local_files_only=local_files_only)
+        unet = UNet2DConditionModel.from_pretrained(unet_model_name,
+                                                    subfolder='unet',
+                                                    cache_dir=cache_dir,
+                                                    local_files_only=local_files_only)
         if isinstance(vae, AutoEncoder) and vae.config['latent_channels'] != 4:
             raise ValueError(f'Pretrained unet has 4 latent channels but the vae has {vae.latent_channels}.')
     else:
@@ -667,6 +678,7 @@ def precomputed_text_latent_diffusion(
         use_xformers (bool): Whether to use xformers for attention. Defaults to True.
         lora_rank (int, optional): If not None, the rank to use for LoRA finetuning. Defaults to None.
         lora_alpha (int, optional): If not None, the alpha to use for LoRA finetuning. Defaults to None.
+        local_files_only (bool): Whether to only use local files. Default: `False`.
     """
     latent_mean, latent_std = _parse_latent_statistics(latent_mean), _parse_latent_statistics(latent_std)
 
@@ -700,7 +712,10 @@ def precomputed_text_latent_diffusion(
         downsample_factor = 2**(len(vae.config['channel_multipliers']) - 1)
 
     # Make the unet
-    unet_config = PretrainedConfig.get_config_dict(unet_model_name, subfolder='unet', cache_dir=cache_dir, local_files_only=local_files_only)[0]
+    unet_config = PretrainedConfig.get_config_dict(unet_model_name,
+                                                   subfolder='unet',
+                                                   cache_dir=cache_dir,
+                                                   local_files_only=local_files_only)[0]
 
     if isinstance(vae, AutoEncoder):
         # Adapt the unet config to account for differing number of latent channels if necessary
@@ -797,7 +812,9 @@ def precomputed_text_latent_diffusion(
     if include_text_encoders:
         dtype_map = {'float32': torch.float32, 'float16': torch.float16, 'bfloat16': torch.bfloat16}
         dtype = dtype_map[text_encoder_dtype]
-        t5_tokenizer = AutoTokenizer.from_pretrained('google/t5-v1_1-xxl', cache_dir=cache_dir, local_files_only=local_files_only)
+        t5_tokenizer = AutoTokenizer.from_pretrained('google/t5-v1_1-xxl',
+                                                     cache_dir=cache_dir,
+                                                     local_files_only=local_files_only)
         clip_tokenizer = AutoTokenizer.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0',
                                                        subfolder='tokenizer',
                                                        cache_dir=cache_dir,