Stable Diffusion 架构与文生图流程

# -*- coding: utf-8 -*-
"""最小文生图：13 行跑通 Stable Diffusion。

环境：
    pip install diffusers transformers accelerate torch safetensors
显存：
    fp16 下约 4~6 GB；没有 NVIDIA 显卡时把 device 改成 "cpu"，一张图要几分钟。
"""
import torch
from diffusers import StableDiffusionPipeline

# ① 载入整条流水线：翻译官（text_encoder）、雕刻师（unet + scheduler）、
#    放大师（vae）三个子模型一次性被打包进 pipe 里。
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",   # 本地已下好就写目录路径
    torch_dtype=torch.float16,          # fp16 省一半显存
)
pipe = pipe.to("cuda")                  # 没显卡就改 "cpu"，并把 dtype 改回 float32

# ② 固定随机种子。同一个 seed + 同一套参数 = 同一张图，这是复现的唯一开关。
generator = torch.Generator(device="cuda").manual_seed(1024)

# ③ 出图。num_inference_steps 就是「雕刻师擦几轮」，guidance_scale 是「多听话」。
image = pipe(
    prompt="a photo of a corgi wearing sunglasses, studio lighting, 8k",
    negative_prompt="lowres, blurry, watermark, extra fingers",
    num_inference_steps=25,
    guidance_scale=7.5,
    height=512,
    width=512,
    generator=generator,
).images[0]

# ④ 这里拿到的已经是 PIL 图片，潜空间那一套在 pipe 内部走完了。
image.save("corgi.png")
print("saved corgi.png", image.size)   # (512, 512)

# -*- coding: utf-8 -*-
"""把 StableDiffusionPipeline 拆开手写一遍，并打印每一步的 tensor shape。

跑一次就能亲眼看到：
    文本  → (1, 77, 768)      翻译官的产物
    潜噪声 → (1, 4, 64, 64)    雕刻师的草稿纸
    成品  → (1, 3, 512, 512)  放大师的输出

这份代码不追求速度，只追求「每一步都看得见」。
"""
import torch
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler
from transformers import CLIPTextModel, CLIPTokenizer
from PIL import Image
import numpy as np

MODEL = "runwayml/stable-diffusion-v1-5"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

PROMPT = "a photo of a corgi wearing sunglasses, studio lighting, 8k"
NEGATIVE = "lowres, blurry, watermark"
STEPS = 25
GUIDANCE = 7.5
HEIGHT, WIDTH = 512, 512
SEED = 1024


# ---------------------------------------------------------------- 子模型
# 四个部件分开加载，正好对应架构图里的三个子模型 + 一个采样器。
tokenizer = CLIPTokenizer.from_pretrained(MODEL, subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(
    MODEL, subfolder="text_encoder", torch_dtype=DTYPE).to(DEVICE)
unet = UNet2DConditionModel.from_pretrained(
    MODEL, subfolder="unet", torch_dtype=DTYPE).to(DEVICE)
vae = AutoencoderKL.from_pretrained(
    MODEL, subfolder="vae", torch_dtype=DTYPE).to(DEVICE)
scheduler = PNDMScheduler.from_pretrained(MODEL, subfolder="scheduler")

# CLIP 的参数在出图过程中是冻结的，这三个模型都只做前向推理。
text_encoder.requires_grad_(False)
unet.requires_grad_(False)
vae.requires_grad_(False)


# ---------------------------------------------- 第一步：翻译官把文字变成向量
def encode(text):
    """任何长度的 prompt 都会被 padding / 截断到固定 77 个 token。"""
    ids = tokenizer(
        text,
        padding="max_length",
        max_length=tokenizer.model_max_length,   # 77
        truncation=True,
        return_tensors="pt",
    ).input_ids.to(DEVICE)
    with torch.no_grad():
        return text_encoder(ids)[0]              # (1, 77, 768)


cond = encode(PROMPT)
uncond = encode(NEGATIVE)
print("[1] 文本嵌入          cond:", tuple(cond.shape), " uncond:", tuple(uncond.shape))

# classifier-free guidance：把「要的」和「不要的」拼成一个 batch 一起算，
# 这样一次 U-Net 前向就能同时拿到两个方向的噪声预测。
text_embeddings = torch.cat([uncond, cond])
print("[2] 拼接后送入 U-Net      :", tuple(text_embeddings.shape))


# ---------------------------------------------- 第二步：在潜空间里搓一张噪声
generator = torch.Generator(device="cpu").manual_seed(SEED)
latents = torch.randn(
    (1, unet.config.in_channels, HEIGHT // 8, WIDTH // 8),   # 8 倍下采样
    generator=generator,
).to(DEVICE, dtype=DTYPE)
print("[3] 初始潜噪声            :", tuple(latents.shape),
      " 元素数 =", latents.numel())

scheduler.set_timesteps(STEPS)
latents = latents * scheduler.init_noise_sigma


# ---------------------------------------------- 第三步：擦 STEPS 轮
for i, t in enumerate(scheduler.timesteps):
    # 同一份潜变量复制两份，分别配 uncond 与 cond 的文本嵌入
    model_input = torch.cat([latents] * 2)
    model_input = scheduler.scale_model_input(model_input, t)

    with torch.no_grad():
        noise_pred = unet(model_input, t,
                          encoder_hidden_states=text_embeddings).sample

    # 拆回两半，按 guidance_scale 在「没有提示词的方向」与
    # 「有提示词的方向」之间外推
    noise_uncond, noise_cond = noise_pred.chunk(2)
    noise_pred = noise_uncond + GUIDANCE * (noise_cond - noise_uncond)

    # 由 scheduler 决定这一步到底减掉多少——这就是「采样算法」的位置
    latents = scheduler.step(noise_pred, t, latents).prev_sample

    if i in (0, STEPS // 2, STEPS - 1):
        print("    step %2d/%d  t=%-4s  latents:" % (i + 1, STEPS, int(t)),
              tuple(latents.shape))

print("[4] 降噪完成的潜变量      :", tuple(latents.shape))


# ---------------------------------------------- 第四步：放大师解码一次
# 0.18215 是 SD v1 系列训练时给潜变量用的缩放因子，解码前必须除回去
latents = latents / vae.config.scaling_factor
with torch.no_grad():
    image = vae.decode(latents).sample
print("[5] VAE 解码后            :", tuple(image.shape))

# [-1, 1] → [0, 1] → uint8 → PIL
image = (image / 2 + 0.5).clamp(0, 1)
image = image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
image = (image * 255).round().astype(np.uint8)
Image.fromarray(image[0]).save("corgi_stepwise.png")
print("[6] 已保存 corgi_stepwise.png  像素尺寸:", (WIDTH, HEIGHT))

# -*- coding: utf-8 -*-
"""图生图：不从纯噪声起步，而是从一张已有图片的潜表示起步。

和文生图唯一的结构差别在 strength：
    strength=1.0  → 原图的潜表示被噪声完全淹没，等价于文生图
    strength=0.6  → 只走 60% 的降噪步数，原图的构图与色块还看得出来
    strength=0.2  → 几乎只做轻微润色
"""
import torch
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=DTYPE).to(DEVICE)

# ① 读入起点图。SD v1 的潜空间是 8 倍下采样，边长建议是 8 的倍数，
#    否则内部会自动裁剪，出图和你给的图对不齐。
init = Image.open("input.png").convert("RGB").resize((512, 512))

generator = torch.Generator(device="cpu").manual_seed(1024)

# ② strength 决定「往原图上撒多厚的一层噪声」。
#    实际执行的步数 = int(num_inference_steps * strength)，
#    所以 strength 调小时，出图也会更快。
out = pipe(
    prompt="watercolor painting, soft pastel colors, delicate brush strokes",
    negative_prompt="photo, 3d render, lowres",
    image=init,
    strength=0.6,
    num_inference_steps=30,     # 实际只跑 30 * 0.6 ≈ 18 步
    guidance_scale=7.5,
    generator=generator,
).images[0]

out.save("img2img_out.png")
print("saved img2img_out.png", out.size)

# -*- coding: utf-8 -*-
"""图像修复（inpaint）：只让潜空间里被 mask 圈中的那一块重新生成。

mask 的约定：
    白色（255）= 要重画的区域
    黑色（0）  = 要原样保留的区域
mask 必须和原图同尺寸，否则内部 resize 会让边界错位。
"""
import torch
from diffusers import StableDiffusionInpaintPipeline
from PIL import Image

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

# 注意模型不是 v1-5 那个：inpaint 用的是专门多训了一路 mask 输入的权重，
# U-Net 的输入通道数从 4 变成 9（4 潜变量 + 4 掩码图潜变量 + 1 掩码）。
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-inpainting", torch_dtype=DTYPE).to(DEVICE)

init = Image.open("room.png").convert("RGB").resize((512, 512))
mask = Image.open("room_mask.png").convert("L").resize((512, 512))

generator = torch.Generator(device="cpu").manual_seed(2048)

out = pipe(
    prompt="a green potted plant on the table, natural light",
    negative_prompt="blurry, distorted, artifacts",
    image=init,
    mask_image=mask,
    num_inference_steps=30,
    guidance_scale=7.5,
    generator=generator,
).images[0]

out.save("inpaint_out.png")
print("saved inpaint_out.png", out.size)

# -*- coding: utf-8 -*-
"""把降噪循环中途的潜变量解码出来，看雕刻师是怎么一轮轮把雪花擦掉的。

这份脚本的唯一目的是「眼见为实」：
它在循环里每隔几步就临时调一次 VAE 解码，把当前潜变量还原成图片存盘。
跑完会得到一组 step_00.png … step_24.png，连起来就是一张图的诞生过程。

注意：中途解码只是为了观察。正式出图时 VAE 只在最后跑一次——
每解码一次就多一次完整的放大开销，中途解码会让出图慢好几倍。
"""
import os

import numpy as np
import torch
from diffusers import StableDiffusionPipeline
from PIL import Image

MODEL = "runwayml/stable-diffusion-v1-5"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

PROMPT = "a photo of a corgi wearing sunglasses, studio lighting, 8k"
NEGATIVE = "lowres, blurry, watermark"
STEPS = 25
GUIDANCE = 7.5
SEED = 1024
PEEK_EVERY = 3            # 每隔几步存一张
OUT_DIR = "peek"

pipe = StableDiffusionPipeline.from_pretrained(MODEL, torch_dtype=DTYPE).to(DEVICE)
os.makedirs(OUT_DIR, exist_ok=True)


def latents_to_image(latents):
    """把 (1, 4, 64, 64) 的潜变量解码成 PIL 图片。

    0.18215 这个缩放因子必须先除回去，否则颜色整体偏掉。
    """
    with torch.no_grad():
        img = pipe.vae.decode(latents / pipe.vae.config.scaling_factor).sample
    img = (img / 2 + 0.5).clamp(0, 1)
    img = img.detach().cpu().permute(0, 2, 3, 1).float().numpy()
    return Image.fromarray((img * 255).round().astype(np.uint8)[0])


def on_step(pipe_ref, step_index, timestep, callback_kwargs):
    """diffusers 的回调钩子：每一步结束后被调用一次。

    callback_kwargs 里带着当前的 latents，可以看、也可以改；
    原样返回就等于不干预循环。
    """
    latents = callback_kwargs["latents"]
    if step_index % PEEK_EVERY == 0 or step_index == STEPS - 1:
        path = os.path.join(OUT_DIR, "step_%02d.png" % step_index)
        latents_to_image(latents).save(path)
        # 顺便打印潜变量的统计量：标准差会随着噪声被擦掉而单调下降
        print("step %2d  t=%-4d  std=%.3f  -> %s"
              % (step_index, int(timestep), float(latents.std()), path))
    return callback_kwargs


generator = torch.Generator(device="cpu").manual_seed(SEED)
image = pipe(
    prompt=PROMPT,
    negative_prompt=NEGATIVE,
    num_inference_steps=STEPS,
    guidance_scale=GUIDANCE,
    generator=generator,
    callback_on_step_end=on_step,
    callback_on_step_end_tensor_inputs=["latents"],
).images[0]

image.save(os.path.join(OUT_DIR, "final.png"))
print("完成，逐步图在 %s/ 下" % OUT_DIR)

# -*- coding: utf-8 -*-
"""参数对照实验：固定 seed，一次只改一个参数，出一张网格图。

这是调参唯一靠谱的做法。seed 一旦随机，你就永远分不清
「这次好看」是参数调对了，还是单纯抽到了好牌。

跑完得到三张网格：
    grid_steps.png      只改 steps
    grid_guidance.png   只改 guidance_scale
    grid_sampler.png    只改采样器
"""
import torch
from diffusers import (StableDiffusionPipeline, DDIMScheduler,
                       EulerDiscreteScheduler, EulerAncestralDiscreteScheduler,
                       DPMSolverMultistepScheduler)
from PIL import Image

MODEL = "runwayml/stable-diffusion-v1-5"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

# 全程不变的基准设置
PROMPT = "a photo of a corgi wearing sunglasses, studio lighting, 8k"
NEGATIVE = "lowres, bad anatomy, watermark"
SEED = 1024
BASE_STEPS = 25
BASE_GUIDANCE = 7.5

pipe = StableDiffusionPipeline.from_pretrained(MODEL, torch_dtype=DTYPE).to(DEVICE)
default_scheduler_config = pipe.scheduler.config


def draw(**over):
    """每次都用同一个 seed 新建 generator——这是「唯一变量」的保证。

    generator 用一次就被消耗掉一部分随机状态，
    复用同一个对象会导致第二张图的起点噪声和第一张不同。
    """
    generator = torch.Generator(device="cpu").manual_seed(SEED)
    kwargs = dict(prompt=PROMPT, negative_prompt=NEGATIVE,
                  num_inference_steps=BASE_STEPS,
                  guidance_scale=BASE_GUIDANCE, generator=generator)
    kwargs.update(over)
    return pipe(**kwargs).images[0]


def grid(images, labels, path, cell=512):
    """把若干张图横向拼成一张网格，便于一眼横向对比。"""
    canvas = Image.new("RGB", (cell * len(images), cell), "white")
    for i, img in enumerate(images):
        canvas.paste(img.resize((cell, cell)), (i * cell, 0))
    canvas.save(path)
    print("%s  <-  %s" % (path, " | ".join(str(x) for x in labels)))


def sweep_steps():
    """步数：20~30 之后收益急剧递减，耗时却是线性增长的。"""
    values = [5, 10, 20, 30, 50]
    imgs = [draw(num_inference_steps=v) for v in values]
    grid(imgs, values, "grid_steps.png")


def sweep_guidance():
    """引导强度：1 几乎不听提示词，7~8 通用，15 以上过饱和、细节崩坏。"""
    values = [1.0, 4.0, 7.5, 12.0, 20.0]
    imgs = [draw(guidance_scale=v) for v in values]
    grid(imgs, values, "grid_guidance.png")


def sweep_sampler():
    """采样器：换它不需要重新加载模型，因为它没有权重。

    Euler a 带随机性（ancestral），步数改一格图就大变；
    其余几个是确定性采样器，适合先把 prompt 和 seed 调稳。
    """
    schedulers = {
        "DDIM": DDIMScheduler,
        "Euler": EulerDiscreteScheduler,
        "Euler_a": EulerAncestralDiscreteScheduler,
        "DPMpp2M": DPMSolverMultistepScheduler,
    }
    imgs, labels = [], []
    for name, cls in schedulers.items():
        pipe.scheduler = cls.from_config(default_scheduler_config)
        imgs.append(draw())
        labels.append(name)
    # 用完把采样器换回默认，免得影响后续调用
    pipe.scheduler = type(pipe.scheduler).from_config(default_scheduler_config)
    grid(imgs, labels, "grid_sampler.png")


if __name__ == "__main__":
    sweep_steps()
    sweep_guidance()
    sweep_sampler()

# -*- coding: utf-8 -*-
"""可复制的出图脚本骨架：所有可调参数集中在顶部，下面的逻辑一个字都不用改。

用法：
    python3 sd_generate_skeleton.py
    改 TODO 处即可；批量出图把 SEEDS 写成多个值。
"""
import os
import json
import torch
from diffusers import StableDiffusionPipeline

# ==========================================================================
# 参数区：只改这一块
# ==========================================================================
# TODO 1：模型路径。可以是 HuggingFace 仓库名，也可以是本地目录。
MODEL_PATH = os.environ.get("SD_MODEL", "runwayml/stable-diffusion-v1-5")

# TODO 2：正向提示词。主体 → 细节 → 风格 → 画质，逗号分隔，越靠前权重越高。
PROMPT = "a photo of a corgi wearing sunglasses, studio lighting, 8k, highly detailed"

# TODO 3：负向提示词。写你「不想要」的东西，这是最省事的质量开关。
NEGATIVE_PROMPT = "lowres, bad anatomy, extra fingers, watermark, text, blurry"

# TODO 4：采样步数。20~30 足够，再高收益递减、耗时线性增长。
STEPS = 25

# TODO 5：提示词引导强度。7~8 是通用值；调太高画面会过饱和、细节崩坏。
GUIDANCE_SCALE = 7.5

# TODO 6：随机种子列表。同一个 seed + 同一套参数 = 同一张图。
#         写 None 表示每次都随机（不可复现，调参阶段不要用）。
SEEDS = [1024, 2048, 4096]

# TODO 7：出图尺寸。必须是 8 的倍数；SD v1 系列在 512 附近训练，
#         直接拉到 1024 容易出现「双头/重复肢体」，要放大请用 upscale 流程。
WIDTH, HEIGHT = 512, 512

# TODO 8：输出目录
OUT_DIR = "outputs"
# ==========================================================================


def build_pipe():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32
    pipe = StableDiffusionPipeline.from_pretrained(MODEL_PATH, torch_dtype=dtype)
    pipe = pipe.to(device)
    # 显存紧张时打开下面两行，用时间换空间
    # pipe.enable_attention_slicing()
    # pipe.enable_vae_slicing()
    return pipe, device


def main():
    os.makedirs(OUT_DIR, exist_ok=True)
    pipe, device = build_pipe()

    seeds = SEEDS if SEEDS else [torch.seed() % (2 ** 32)]
    for seed in seeds:
        generator = torch.Generator(device="cpu").manual_seed(int(seed))
        image = pipe(
            prompt=PROMPT,
            negative_prompt=NEGATIVE_PROMPT,
            num_inference_steps=STEPS,
            guidance_scale=GUIDANCE_SCALE,
            width=WIDTH,
            height=HEIGHT,
            generator=generator,
        ).images[0]

        stem = os.path.join(OUT_DIR, "seed_%d" % seed)
        image.save(stem + ".png")

        # 出图参数与图片一起落盘。没有这一步，三天后你就再也复现不出这张图。
        meta = {
            "model": MODEL_PATH, "prompt": PROMPT,
            "negative_prompt": NEGATIVE_PROMPT, "steps": STEPS,
            "guidance_scale": GUIDANCE_SCALE, "seed": int(seed),
            "size": [WIDTH, HEIGHT], "device": device,
        }
        with open(stem + ".json", "w", encoding="utf-8") as f:
            json.dump(meta, f, ensure_ascii=False, indent=2)
        print("saved", stem + ".png")


if __name__ == "__main__":
    main()