当前位置：首页 > article >正文

Diffusers使用笔记

article 2024/12/26 14:51:03

Diffusers 是用于生成图像、音频等最先进预训练扩散模型的库。它既支持推理解决方案，也支持训练自己的扩散模型，Diffusers 是一个支持这两者的模块化工具箱。区别与ComfyUI与webUI这类UI类的应用，Diffusers实际上是更底层的库，可以支持更好的建立自己的工作流而不仅仅是应用。这一点使它更为开放人员所接受。

1. 安装

pip install --upgrade diffusers

2. 模型加载

自动下载模型

import torch
import requests
from PIL import Image
from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler

# Load the pipeline
pipeline = DiffusionPipeline.from_pretrained(
    "sudo-ai/zero123plus-v1.1", custom_pipeline="sudo-ai/zero123plus-pipeline",
    torch_dtype=torch.float16
)

从本地加载模型

import torch
import requests
from PIL import Image
from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler

# Load the pipeline
pipeline = DiffusionPipeline.from_pretrained(
    "/xxxx/models--sudo-ai--zero123plus-v1.1/snapshots/36df7de980afd15f80b2e1a4e9a920d7020e2654", 
    custom_pipeline="/xxxx/models--sudo-ai--zero123plus-pipeline/snapshots/983e66d28a3637ddd8e3e2fd8165cdff32230872",
    local_files_only=True,
    torch_dtype=torch.float16
)

3. 推理

def infer(self):
  # 0. 定义unet中的高宽，这里要考虑经过vae后的缩小系数
  height = height or self.unet.config.sample_size * self.vae_scale_factor
  width = width or self.unet.config.sample_size * self.vae_scale_factor


  # 1. 检查输入是否合规
  self.check_inputs(
      prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
  )

  # 2. 定义调用的batch，device以及classifier-free的监督系数
  if prompt is not None and isinstance(prompt, str):
      batch_size = 1
  elif prompt is not None and isinstance(prompt, list):
      batch_size = len(prompt)
  else:
      batch_size = prompt_embeds.shape[0]

  device = self._execution_device
  do_classifier_free_guidance = guidance_scale > 1.0

  # 3. 将输入的文字prompt进行encoding
  text_encoder_lora_scale = (
      cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
  )
  prompt_embeds = self._encode_prompt(
      prompt,
      device,
      num_images_per_prompt,
      do_classifier_free_guidance,
      negative_prompt,
      prompt_embeds=prompt_embeds,
      negative_prompt_embeds=negative_prompt_embeds,
      lora_scale=text_encoder_lora_scale,
  )

  # 4. 准备scheduler中的时间步数
  self.scheduler.set_timesteps(num_inference_steps, device=device)
  timesteps = self.scheduler.timesteps

  # 5. 生成对应尺寸的初始噪声图latent
  num_channels_latents = self.unet.config.in_channels
  latents = self.prepare_latents(
      batch_size * num_images_per_prompt,
      num_channels_latents,
      height,
      width,
      prompt_embeds.dtype,
      device,
      generator,
      latents,
  )

  # 6. 额外的去噪参数
  extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

  # 7. 循环多个步长进行去噪
  num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
  with self.progress_bar(total=num_inference_steps) as progress_bar:
      for i, t in enumerate(timesteps):
          latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
          latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
          noise_pred = self.unet(
              latent_model_input,
              t,
              encoder_hidden_states=prompt_embeds,
              cross_attention_kwargs=cross_attention_kwargs,
              return_dict=False,
          )[0]

          if do_classifier_free_guidance:
              noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
              noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

          if do_classifier_free_guidance and guidance_rescale > 0.0:
              noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)

          latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]

          if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
              progress_bar.update()
              if callback is not None and i % callback_steps == 0:
                  callback(i, t, latents)

  if not output_type == "latent":
      image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
      image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
  else:
      image = latents
      has_nsfw_concept = None

  if has_nsfw_concept is None:
      do_denormalize = [True] * image.shape[0]
  else:
      do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]

  image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)

  if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
      self.final_offload_hook.offload()

  if not return_dict:
      return (image, has_nsfw_concept)

  return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)