diff --git a/examples/textual_inversion/.gitignore b/examples/textual_inversion/.gitignore new file mode 100644 index 000000000000..8de1bf1e9f21 --- /dev/null +++ b/examples/textual_inversion/.gitignore @@ -0,0 +1,3 @@ +frida +wandb +textual_inversion_frida \ No newline at end of file diff --git a/examples/textual_inversion/textual_inversion.py b/examples/textual_inversion/textual_inversion.py index 0b3515ff8efa..27670f8e69e7 100644 --- a/examples/textual_inversion/textual_inversion.py +++ b/examples/textual_inversion/textual_inversion.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and import argparse +import importlib import logging import math import os @@ -34,7 +35,14 @@ from accelerate import Accelerator from accelerate.logging import get_logger from accelerate.utils import set_seed -from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel +from diffusers import ( + AutoencoderKL, + DDPMScheduler, + DiffusionPipeline, + DPMSolverMultistepScheduler, + StableDiffusionPipeline, + UNet2DConditionModel, +) from diffusers.optimization import get_scheduler from diffusers.utils import check_min_version from diffusers.utils.import_utils import is_xformers_available @@ -48,6 +56,28 @@ from transformers import CLIPTextModel, CLIPTokenizer +_wandb_available = importlib.util.find_spec("wandb") is not None + + +def is_wandb_available(): + return _wandb_available + + +wandb = None +if is_wandb_available(): + import wandb + + +def wandb_setup( + args: dict, + project_name: str = "glide-text2im-finetune", +): + return wandb.init( + project=project_name, + config=args, + ) + + if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"): PIL_INTERPOLATION = { "linear": PIL.Image.Resampling.BILINEAR, @@ -83,6 +113,42 @@ def save_progress(text_encoder, placeholder_token_id, accelerator, args, save_pa def parse_args(): parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser.add_argument( + "--num_inference_steps", + type=int, + default=25, + help="Number of steps for inference.", + ) + parser.add_argument( + "--validation_prompt", + type=str, + default="A photo of ", + help="Prompt to use for validation.", + ) + parser.add_argument( + "--num_validation_images", + type=int, + default=1, + help="Number of images to generate for validation.", + ) + parser.add_argument( + "--guidance_scale", + type=float, + default=7.5, + help="Guidance scale for inference.", + ) + parser.add_argument( + "--project_name", + type=str, + default="huggingface_textual_inv", + help="Name of wandb run", + ) + parser.add_argument( + "--log_frequency", + type=int, + default=100, + help="Frequency to log/save the model.", + ) parser.add_argument( "--save_steps", type=int, @@ -404,10 +470,7 @@ def __getitem__(self, i): if self.center_crop: crop = min(img.shape[0], img.shape[1]) - ( - h, - w, - ) = ( + (h, w,) = ( img.shape[0], img.shape[1], ) @@ -437,14 +500,12 @@ def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: def main(): args = parse_args() logging_dir = os.path.join(args.output_dir, args.logging_dir) - accelerator = Accelerator( gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision, log_with=args.report_to, logging_dir=logging_dir, ) - # Make one log on every process with the configuration for debugging. logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", @@ -724,6 +785,45 @@ def main(): if accelerator.sync_gradients: progress_bar.update(1) global_step += 1 + if args.validation_prompt is not None and global_step % args.log_frequency == 0: + logger.info( + f"Running validation... \n Generating {args.num_validation_images} images with prompt:" + f" {args.validation_prompt}." + ) + # create pipeline + pipeline = DiffusionPipeline.from_pretrained( + args.pretrained_model_name_or_path, + unet=unet, + vae=vae, + text_encoder=accelerator.unwrap_model(text_encoder), + tokenizer=tokenizer, + revision=args.revision, + torch_dtype=weight_dtype, + ) + pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config) + pipeline = pipeline.to(accelerator.device) + pipeline.set_progress_bar_config(disable=True) + # run inference + generator = torch.Generator(device=accelerator.device).manual_seed(args.seed or 0) + prompt = args.num_validation_images * [args.validation_prompt] + images = pipeline( + prompt, + num_inference_steps=args.num_inference_steps, + generator=generator, + guidance_scale=args.guidance_scale, + ).images + for tracker in accelerator.trackers: + if tracker.name == "wandb": + tracker.log( + { + "validation": [ + wandb.Image(image, caption=f"{i}: {args.validation_prompt}") + for i, image in enumerate(images) + ] + } + ) + del pipeline + torch.cuda.empty_cache() if global_step % args.save_steps == 0: save_path = os.path.join(args.output_dir, f"learned_embeds-steps-{global_step}.bin") save_progress(text_encoder, placeholder_token_id, accelerator, args, save_path) @@ -752,9 +852,9 @@ def main(): if save_full_model: pipeline = StableDiffusionPipeline.from_pretrained( args.pretrained_model_name_or_path, - text_encoder=accelerator.unwrap_model(text_encoder), - vae=vae, unet=unet, + vae=vae, + text_encoder=accelerator.unwrap_model(text_encoder), tokenizer=tokenizer, ) pipeline.save_pretrained(args.output_dir) diff --git a/examples/textual_inversion/textual_inversion_flax.py b/examples/textual_inversion/textual_inversion_flax.py index 948346c4f74e..d84dd4f11692 100644 --- a/examples/textual_inversion/textual_inversion_flax.py +++ b/examples/textual_inversion/textual_inversion_flax.py @@ -306,10 +306,7 @@ def __getitem__(self, i): if self.center_crop: crop = min(img.shape[0], img.shape[1]) - ( - h, - w, - ) = ( + (h, w,) = ( img.shape[0], img.shape[1], ) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index b38ca866d58d..bd378784079c 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -482,7 +482,7 @@ def __call__( # 3. Encode input prompt text_embeddings = self._encode_prompt( prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt - ) + ).to(dtype=self.vae.dtype) # 4. Prepare timesteps self.scheduler.set_timesteps(num_inference_steps, device=device)