fix example style

Author: jiqing-feng

examples/textual_inversion/textual_inversion_sdxl.py

@@ -111,7 +111,9 @@ def save_model_card(repo_id: str, images=None, base_model=str, repo_folder=None)
         f.write(yaml + model_card)
 
 
-def log_validation(text_encoder_1, text_encoder_2, tokenizer_1, tokenizer_2, unet, vae, args, accelerator, weight_dtype, epoch):
+def log_validation(
+    text_encoder_1, text_encoder_2, tokenizer_1, tokenizer_2, unet, vae, args, accelerator, weight_dtype, epoch
+):
     logger.info(
         f"Running validation... \n Generating {args.num_validation_images} images with prompt:"
         f" {args.validation_prompt}."
@@ -644,7 +646,6 @@ def main():
         args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, variant=args.variant
     )
 
-
     # Add the placeholder token in tokenizer_1
     placeholder_tokens = [args.placeholder_token]
 
@@ -875,17 +876,27 @@ def main():
                 noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
                 # Get the text embedding for conditioning
-                encoder_hidden_states_1 = text_encoder_1(batch["input_ids_1"], output_hidden_states=True).hidden_states[-2].to(dtype=weight_dtype)
-                encoder_output_2 = text_encoder_2(batch["input_ids_2"].reshape(batch["input_ids_1"].shape[0], -1), output_hidden_states=True)
+                encoder_hidden_states_1 = (
+                    text_encoder_1(batch["input_ids_1"], output_hidden_states=True)
+                    .hidden_states[-2]
+                    .to(dtype=weight_dtype)
+                )
+                encoder_output_2 = text_encoder_2(
+                    batch["input_ids_2"].reshape(batch["input_ids_1"].shape[0], -1), output_hidden_states=True
+                )
                 encoder_hidden_states_2 = encoder_output_2.hidden_states[-2].to(dtype=weight_dtype)
                 sample_size = unet.config.sample_size * (2 ** (len(vae.config.block_out_channels) - 1))
                 original_size = (sample_size, sample_size)
-                add_time_ids = torch.tensor([list(original_size + (0, 0) + original_size)], dtype=weight_dtype, device=accelerator.device)
+                add_time_ids = torch.tensor(
+                    [list(original_size + (0, 0) + original_size)], dtype=weight_dtype, device=accelerator.device
+                )
                 added_cond_kwargs = {"text_embeds": encoder_output_2[0], "time_ids": add_time_ids}
                 encoder_hidden_states = torch.cat([encoder_hidden_states_1, encoder_hidden_states_2], dim=-1)
 
                 # Predict the noise residual
-                model_pred = unet(noisy_latents, timesteps, encoder_hidden_states, added_cond_kwargs=added_cond_kwargs).sample
+                model_pred = unet(
+                    noisy_latents, timesteps, encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
+                ).sample
 
                 # Get the target for loss depending on the prediction type
                 if noise_scheduler.config.prediction_type == "epsilon":
@@ -961,7 +972,16 @@ def main():
 
                     if args.validation_prompt is not None and global_step % args.validation_steps == 0:
                         images = log_validation(
-                            text_encoder_1, text_encoder_2, tokenizer_1, tokenizer_2, unet, vae, args, accelerator, weight_dtype, epoch
+                            text_encoder_1,
+                            text_encoder_2,
+                            tokenizer_1,
+                            tokenizer_2,
+                            unet,
+                            vae,
+                            args,
+                            accelerator,
+                            weight_dtype,
+                            epoch,
                         )
 
             logs = {"loss": loss.detach().item(), "lr": lr_scheduler_1.get_last_lr()[0]}
@@ -1020,4 +1040,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-