diffusers库学习--pipeline，模型，调度器的基础使用

pipeline，模型，调度器的基础使用

⚙️ 1. 基础设置

首先，我们从 diffusers、torch 和其他工具库中导入所有必要的组件，并设置好在 GPU 上运行的设备。

from diffusers import DiffusionPipeline
from diffusers import EulerDiscreteScheduler
from diffusers import AutoModel
from diffusers import UNet2DModel
from diffusers import DDPMScheduler
import torch
from tqdm.auto import tqdm
import PIL.Image
import numpy as np

device="cuda" if torch.cuda.is_available() else "cpu"

🚀 Demo 1: 高度封装的 Pipeline 用法

这是生成图像最简单、最直接的方式。DiffusionPipeline (扩散管道) 为我们封装了所有复杂的步骤。

核心概念： 使用 DiffusionPipeline.from_pretrained() 来加载一个像 Stable Diffusion 一样配置齐全的“开箱即用”模型。之后，只需要提供一个提示词 (Prompt) 来调用这个 pipeline 对象就可以生成图像。

def demo1(): # 测试SD1.5默认管道设置
    pipeline=DiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        torch_dtype=torch.float16,
        use_safetensors=True
        ) # 模型id ，精度，是否用safetensors
    
    pipeline.to(device)
    image1=pipeline(
        prompt="A cute cat",
        num_inference_steps=25,
        guidance_scale=7.5,
        negative_prompt=""
        ).images[0]
    image1.save("./test_cat.png")

🔄 Demo 2: 更换调度器 (Scheduler)

Pipeline 是模块化的。你可以轻松地替换其中的组件，例如“调度器”，来改变去噪过程的行为和效果。

核心概念： 加载 pipeline 后，可以直接替换它的 .scheduler 属性。EulerDiscreteScheduler.from_config(pipeline.scheduler.config) 这个用法非常巧妙，它能方便地创建一个新的调度器，并让它继承旧调度器的所有配置。

def demo2():# 更换调度器
    pipeline=DiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        torch_dtype=torch.float16,
        use_safetensors=True
        ) # 模型id ，精度，是否用safetensors
    pipeline.scheduler=EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
    pipeline.to(device)
    image1=pipeline(
        prompt="A cute cat",
        num_inference_step=25,
        guidance_scale=7.5,
        negative_prompt=""
        ).images[0]
    image1.save("./test_cat2.png")

🧩 Demo 3: 使用独立的 U-Net 模型

除了使用完整的 pipeline，我们也可以只和独立的组件打交道。扩散模型最核心的组件就是 U-Net，它负责预测噪声。

核心概念： 直接从仓库加载一个模型，比如 UNet2DModel。要使用它，你必须手动创建一个输入（一个随机噪声张量），并提供一个特定的 timestep (时间步)。模型的输出就是在该时间步预测出的噪声。

def demo3(): # AutoModel的使用和 UNet2DModel 模型
    repo_id="google/ddpm-cat-256"
    model1=AutoModel.from_pretrained(
        repo_id,
        use_safetensors=True
        ).to(device)
    
    noise_sample=torch.randn(1,model1.config.in_channels,model1.config.sample_size,model1.config.sample_size).to(device)
    
    # 有了噪声图像然后就是推理了，推理还需要一个时间步
    with torch.no_grad():
        # 直接把时间步和噪声图像传给模型降噪
        noisy_residual = model1(sample=noise_sample, timestep=2).sample 
    print(noisy_residual)

🧠 Demo 4: 理解调度器的作用

调度器是指导整个去噪过程的“大脑”。它利用模型的输出来计算出图像在上一个时间步的状态。

核心概念： scheduler.step() 方法是调度器的核心。它接收 model_output (模型预测的噪声)、当前的 timestep (时间步) 和当前的 sample (样本，即 x_t)，然后计算出 prev_sample (上一步的样本，即 x_{t-1}). 这个过程为我们展示了去噪过程中的一个独立的、精确控制的步骤。

def demo4(): # 调度器
    repo_id="google/ddpm-cat-256"
    model1=AutoModel.from_pretrained(
        repo_id,
        use_safetensors=True
        ).to(device)
    noise_sample=torch.randn(1,model1.config.in_channels,model1.config.sample_size,model1.config.sample_size).to(device)
    with torch.no_grad():
        noisy_residual = model1(sample=noise_sample, timestep=2).sample 

    scheduler=DDPMScheduler.from_config(repo_id)
    less_noise_sample=scheduler.step(model_output=noisy_residual,timestep=2,sample=noise_sample).prev_sample

📜 Demo 5: 完整的“从零开始”推理循环

通过在一个循环中结合 U-Net 模型和调度器，我们可以从零开始复现 pipeline 内部的完整逻辑。这给了我们最大程度的控制权。

核心概念： 一个完整的推理过程，就是一个不断重复“预测噪声”和“时间步回退”的循环。

具体流程如下：

1. 加载 U-Net 模型和一个兼容的调度器。
2. 使用 scheduler.set_timesteps() 告诉调度器你想要进行多少步推理。
3. 创建一个正确尺寸的初始随机噪声张量。
4. 使用 for 循环遍历 scheduler.timesteps 数组。在每次循环中:
   a. 用 U-Net 模型根据当前样本预测噪声。
   b. 用 scheduler.step() 计算出上一个时间步的样本。
   c. 用上一步的结果更新当前样本，为下一次循环做准备。
5. 循环结束后，将最终的张量结果转换为可以查看和保存的图像。

def demo5():# 综合练习，进行一个完整推理
    repo_id="google/ddpm-cat-256"
    model=UNet2DModel.from_pretrained(repo_id,use_safetensors=True).to(device)
    scheduler=DDPMScheduler.from_config(repo_id)
    
    # 关键：设置推理步数
    num_inference_steps = 40
    scheduler.set_timesteps(num_inference_steps)

    # 创建初始的随机噪声
    sample_size = model.config.sample_size
    sample = torch.randn(1,model.config.in_channels,model.config.sample_size,model.config.sample_size).to(device)

    # 使用tqdm来显示进度条
    for t in tqdm(scheduler.timesteps):
        with torch.no_grad():
            # 1. 预测噪声
            noisy_residual = model(sample, t).sample
        # 2. 计算上一步的样本
        sample = scheduler.step(noisy_residual, t, sample).prev_sample

    # 循环结束后，将最终的样本转换为图像
    image_processed = sample.cpu().permute(0, 2, 3, 1)
    image_processed = (image_processed + 1.0) * 127.5
    image_processed = image_processed.numpy().astype(np.uint8)
    image_pil = PIL.Image.fromarray(image_processed[0])
    
    # 保存最终的图像
    image_pil.save("cat_after_40_steps.png")