稳固扩散模型因其从文本形容生成高质量、多样化图像的能力而取得了极大的关注。然而这些预训练模型在生成高度定制或个性化主题的图像时可能会有所有余。

这时就须要咱们进行手动的微调。微调能够依据绝对较小的图像集向模型传授新的、独特的主题。咱们明天应用DreamBooth在不影响模型原始性能的状况下实现微调过程。

根底概念

1、生成模型和文本到图像的合成

生成模型是一类机器学习模型，旨在生成与给定数据集类似的新数据实例。他们捕获潜在的数据分布，产生新的样本。

文本到图像模型是生成模型的一个子集，因为它们以极高的准确性和保真度将文本形容转换为相应的视觉示意而特地乏味。

Stable Diffusion是一种文本到图像的模型，它利用Transformer架构的一种变体来依据文本输出生成图像。

这些模型的生成过程能够形容如下:

给定一个文本形容T，模型的指标是生成一个图像I，使联结概率P(I,T)最大化。这通常是通过训练模型来最大化条件概率P(I∣T)来实现的，以确保生成的图像与文本形容统一。

2、生成模型的微调

微调是指在一个新的、通常更小的数据集上调整预训练的模型，以使模型适应特定的需要，而不会失去从原始数据集学习到的泛化性。这种办法在数据稀缺或须要定制的应用程序中至关重要。

在数学上，微调调整模型的参数以优化新数据集Dnew上的损失函数L，同时避免与原始参数 origin的显著偏差。这能够被表述为一个正则化问题:

3、DreamBooth

DreamBooth提出了一种新的微调办法，容许生成具备特定主题或对象的图像，同时放弃模型生成不同图像的能力。

传统的微调可能会导致适度拟合或灾难性忘记(遗记原始数据散布)，DreamBooth则能确保模型保留其个别性能。

该过程包含训练特定于主题的标记以及原始模型参数。这在概念上相似于向模型的词汇表中增加一个代表新主题的新“单词”。训练的指标能够形容为:

其中示意主题特定参数，1,21，2为正则化参数。

更具体内容请看原论文 《DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation》

技术筹备

应用DreamBooth微调像Stable Diffusion这样的生成模型须要大量的计算，并且须要大量的内存。为了确保训练过程的效率并防止潜在的瓶颈，强烈建议应用高性能GPU。

如果没有GPU，咱们能够间接应用Google Colab，他就能够满足本文的需要。

而后须要装置以下库

Diffusers:扩散模型库，专门用于微调和利用预训练模型。

Accelerate:一个用于分布式训练和混合精度的库。

TensorBoard:用于可视化训练进度和指标。

transformer、FTFY和Gradio:用于模型组件、文本处理和创立用于模型交互的web ui。

Bitsandbytes:用于内存高效和疾速训练，特地是用于优化特定GPU架构上的模型训练。

 !pip install -U -qq git+https://github.com/huggingface/diffusers.git !pip install -qq accelerate tensorboard transformers ftfy gradio !pip install -qq "ipywidgets>=7,<8" !pip install -qq bitsandbytes !pip install huggingface_hub

为了更快、更节俭内存的训练，特地是如果应用特定类型的gpu (T4、P100、V100、A100)，还能够抉择装置以下组件:

Xformers:提供高效transformers 组件的库。

Triton:用于gpu编程。

而后咱们开始进行代码的编写，首先导入库:

 importargparse importitertools importmath importos fromcontextlibimportnullcontext importrandom  importnumpyasnp importtorch importtorch.nn.functionalasF importtorch.utils.checkpoint fromtorch.utils.dataimportDataset  importPIL fromaccelerateimportAccelerator fromaccelerate.loggingimportget_logger fromaccelerate.utilsimportset_seed fromdiffusersimportAutoencoderKL, DDPMScheduler, PNDMScheduler, StableDiffusionPipeline, UNet2DConditionModel fromdiffusers.optimizationimportget_scheduler fromdiffusers.pipelines.stable_diffusionimportStableDiffusionSafetyChecker fromPILimportImage fromtorchvisionimporttransforms fromtqdm.autoimporttqdm fromtransformersimportCLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer  importbitsandbytesasbnb  defimage_grid(imgs, rows, cols):     assertlen(imgs) ==rows*cols      w, h=imgs[0].size     grid=Image.new('RGB', size=(cols*w, rows*h))     grid_w, grid_h=grid.size      fori, imginenumerate(imgs):         grid.paste(img, box=(i%cols*w, i//cols*h))     returngrid

数据集的筹备

这个过程包含抉择有代表性的图像，对它们进行预处理，并将它们组织成一个结构化的格局以供训练。在本文中，咱们将应用以下4个训练图像作为示例。

1、下载和可视化训练图像

上面的download_image函数用于从指定的url列表中检索图像，而后下载这些图像并将其转换为RGB格局以放弃一致性。

咱们将应用他来下载下面的4个图片

 urls= [         "https://huggingface.co/datasets/Entreprenerdly/finetunestablediffusion/resolve/main/2.jpeg",         "https://huggingface.co/datasets/Entreprenerdly/finetunestablediffusion/resolve/main/3.jpeg",         "https://huggingface.co/datasets/Entreprenerdly/finetunestablediffusion/resolve/main/5.jpeg",         "https://huggingface.co/datasets/Entreprenerdly/finetunestablediffusion/resolve/main/6.jpeg",         ## Add additional images here         ]  importrequests importglob fromioimportBytesIO  defdownload_image(url):   try:     response=requests.get(url)   except:     returnNone   returnImage.open(BytesIO(response.content)).convert("RGB")  images=list(filter(None,[download_image(url) forurlinurls])) save_path="./my_concept" ifnotos.path.exists(save_path):   os.mkdir(save_path) [image.save(f"{save_path}/{i}.jpeg") fori, imageinenumerate(images)] image_grid(images, 1, len(images))

2、创立图像文字对

为了应用DreamBooth对生成模型进行微调，须要配置特定的设置以无效地定义新概念。instance_prompt十分重要，因为它蕴含一个描述性标识符，模型应用它来辨认和生成新概念——在本例中是cat_toy。

prior_preservation标记批示模型在训练期间是否应该保留更宽泛的类属性。这有助于提高质量和泛化，但可能会缩短训练工夫。

 instance_prompt="<cat-toy> toy"  # Descriptive prompt with unique identifier prior_preservation=False  # Flag for enabling class characteristics preservation prior_preservation_class_prompt="a photo of a cat clay toy"  # Prompt for the class of the concept  # Parameters for class image generation and loss weighting num_class_images=12 sample_batch_size=2 prior_loss_weight=0.5 prior_preservation_class_folder="./class_images"  # Directories for storing class images class_data_root=prior_preservation_class_folder class_prompt=prior_preservation_class_prompt

3、自定义DataSet类

DreamBoothDataset类扩大了PyTorch的Dataset，用于治理用于训练模型的图像数据。它负责从指定的目录加载图像，利用所需的转换，并应用提供的标记器对提醒进行编码。

另一个类PromptDataset被设置为解决类图像提醒的生成。这个简略的数据集构造存储了提醒和要生成的样本数量。

这些类有助于为DreamBooth微调过程构建训练数据，确保模型以预期的格局接收数据并进行必要的扩大。

 # Initialization of DreamBoothDataset with directory paths and settings classDreamBoothDataset(Dataset):     ...     def__init__(self, instance_data_root, instance_prompt, tokenizer, class_data_root=None, class_prompt=None, size=512, center_crop=False):         ...         self.image_transforms=transforms.Compose([...])      def__getitem__(self, index):         ...         returnexample  # Class for prompt dataset classPromptDataset(Dataset):     ...

优于篇幅太长，咱们就不贴残缺代码了，请在最初的残缺代码中查看

模型加载与配置

咱们将应用stable-diffusion-2

 pretrained_model_name_or_path = "stabilityai/stable-diffusion-2"

而后就是与DreamBooth微调相干的配置参数包含加载模型体系结构的各种组件:文本编码器、变分主动编码器(VAE)和U-Net。

每个组件都是从预训练的模型中加载的，以确保兼容性并保留学习到的特色。

 # Load models and create wrapper for stable diffusion text_encoder=CLIPTextModel.from_pretrained(     pretrained_model_name_or_path, subfolder="text_encoder" ) vae=AutoencoderKL.from_pretrained(     pretrained_model_name_or_path, subfolder="vae" ) unet=UNet2DConditionModel.from_pretrained(     pretrained_model_name_or_path, subfolder="unet" ) tokenizer=CLIPTokenizer.from_pretrained(     pretrained_model_name_or_path,     subfolder="tokenizer", )

这里的CLIPTextModel负责将文本形容编码为嵌入；AutoencoderKL解决图像的潜在空间示意；UNet2DConditionModel是模型的次要生成网络；tokenizer则是解决词元令牌的标记器。

微调

1、设置训练参数

训练参数封装在argparse模块的Namespace类中。这些参数包含预训练模型的门路、图像的分辨率、是否训练文本编码器、学习率以及与启用之前保留相干的细节。

 fromargparseimportNamespace args=Namespace(     pretrained_model_name_or_path=pretrained_model_name_or_path,     resolution=vae.sample_size,     center_crop=True,     train_text_encoder=False,     instance_data_dir=save_path,     instance_prompt=instance_prompt,     learning_rate=5e-06,     max_train_steps=300,     save_steps=50,     train_batch_size=2, # set to 1 if using prior preservation     gradient_accumulation_steps=2,     max_grad_norm=1.0,     mixed_precision="fp16", # set to "fp16" for mixed-precision training.     gradient_checkpointing=True, # set this to True to lower the memory usage.     use_8bit_adam=True, # use 8bit optimizer from bitsandbytes     seed=3434554,     with_prior_preservation=prior_preservation,     prior_loss_weight=prior_loss_weight,     sample_batch_size=2,     class_data_dir=prior_preservation_class_folder,     class_prompt=prior_preservation_class_prompt,     num_class_images=num_class_images,     lr_scheduler="constant",     lr_warmup_steps=100,     output_dir="dreambooth-concept", )

2、定义训练函数

训练函数负责微调过程。它初始化Accelerator以解决分布式训练，为调试复现设置随机种子，配置优化器(应用8位精度以进步内存效率)。

该函数还筹备数据集和数据加载器，设置学习率调度器，并定义训练循环，其中包含损失计算和模型更新。

上面是训练函数设置的代码片段:

 fromaccelerate.utilsimportset_seed deftraining_function(text_encoder, vae, unet):     logger=get_logger(__name__)     set_seed(args.seed)     ...     ifargs.use_8bit_adam:         optimizer_class=bnb.optim.AdamW8bit     else:         optimizer_class=torch.optim.AdamW     ...

3、训练

理论的训练是应用accelerate库中的notebook_launcher启动的，它承受训练函数和先前定义的参数。

 importaccelerate accelerate.notebook_launcher(training_function, args=(text_encoder, vae, unet)) forparaminitertools.chain(unet.parameters(), text_encoder.parameters()):   ifparam.gradisnotNone:     delparam.grad  # free some memory   torch.cuda.empty_cache()

推理

应用StableDiffusionPipeline设置推理管道:

 fromdiffusersimportStableDiffusionPipeline, DPMSolverMultistepScheduler  try:     pipe exceptNameError:     pipe=StableDiffusionPipeline.from_pretrained(         args.output_dir,         scheduler=DPMSolverMultistepScheduler.from_pretrained(args.output_dir, subfolder="scheduler"),         torch_dtype=torch.float16,     ).to("cuda")

在Gradio和/或Google Colab上运行Pipeline

 importgradioasgr  definference(prompt, num_samples):     images=pipe(prompt, num_images_per_prompt=num_samples, num_inference_steps=25).images     returnimages  withgr.Blocks() asdemo:     prompt=gr.Textbox(label="prompt")     samples=gr.Slider(label="Samples", value=1)     run=gr.Button(value="Run")     gallery=gr.Gallery(show_label=False)     run.click(inference, inputs=[prompt, samples], outputs=gallery) demo.launch()

如果是Colab，则应用上面代码

 prompt = "a <cat-toy> in mad max fury road" num_samples = 2 all_images = [] images = pipe(prompt, num_images_per_prompt=num_samples, num_inference_steps=25, guidance_scale=9).images all_images.extend(images) grid = image_grid(all_images, num_rows, num_samples)

总结

dreambooth通过向模型注入自定义的主题来fine-tune diffusion model，它简化了咱们微调自定义模型的老本，而Google Colab的收费GPU能够让咱们进行更多的测试，以下是本文的残缺代码，能够间接在线测试：

https://avoid.overfit.cn/post/1b02c6c3d93749558941777786b64513

作者：Cris Velasquez