生成式 AI 与 Diffusion 模型完全指南：从 Stable Diffusion、ControlNet 到视频生成

# 基本 GAN 结构
import torch
import torch.nn as nn

class Generator(nn.Module):
    def __init__(self, latent_dim=100, img_size=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, img_size * img_size * 3),
            nn.Tanh()
        )

    def forward(self, z):
        return self.net(z).view(-1, 3, 64, 64)

import torch
import torch.nn.functional as F

class DDPMScheduler:
    def __init__(self, num_timesteps=1000, beta_start=1e-4, beta_end=0.02):
        self.T = num_timesteps
        # 线性噪声调度
        self.betas = torch.linspace(beta_start, beta_end, num_timesteps)
        self.alphas = 1.0 - self.betas
        self.alpha_bar = torch.cumprod(self.alphas, dim=0)

    def add_noise(self, x0, noise, t):
        """对 x0 添加 t 步噪声（reparameterization trick）"""
        sqrt_alpha_bar = self.alpha_bar[t] ** 0.5
        sqrt_one_minus = (1 - self.alpha_bar[t]) ** 0.5
        # 调整 shape 以便广播
        sqrt_alpha_bar = sqrt_alpha_bar.view(-1, 1, 1, 1)
        sqrt_one_minus = sqrt_one_minus.view(-1, 1, 1, 1)
        return sqrt_alpha_bar * x0 + sqrt_one_minus * noise

def ddpm_training_step(model, scheduler, x0, optimizer):
    batch_size = x0.shape[0]
    # 随机采样时间步
    t = torch.randint(0, scheduler.T, (batch_size,))
    # 采样高斯噪声
    noise = torch.randn_like(x0)
    # 添加噪声（forward process）
    xt = scheduler.add_noise(x0, noise, t)
    # 预测噪声
    predicted_noise = model(xt, t)
    # MSE 损失
    loss = F.mse_loss(predicted_noise, noise)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

from diffusers import StableDiffusionPipeline
import torch

# 使用基础 pipeline
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16
).to("cuda")

# 生成图像
image = pipe(
    prompt="a serene mountain landscape at sunset, photorealistic",
    negative_prompt="blurry, low quality, distorted",
    num_inference_steps=30,
    guidance_scale=7.5,
    width=512,
    height=512
).images[0]

image.save("output.png")

# 可视化 VAE 潜在空间
from diffusers import AutoencoderKL
from PIL import Image
import torchvision.transforms as T

vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
vae = vae.to("cuda").eval()

transform = T.Compose([T.Resize((512, 512)), T.ToTensor(),
                        T.Normalize([0.5], [0.5])])

img = transform(Image.open("input.png")).unsqueeze(0).to("cuda")
with torch.no_grad():
    # 像素 → 潜在（编码）
    latent = vae.encode(img).latent_dist.sample()
    latent = latent * vae.config.scaling_factor
    print(f"潜在空间尺寸: {latent.shape}")  # [1, 4, 64, 64]

from transformers import CLIPTextModel, CLIPTokenizer

tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")

prompt = "a fantasy castle in the clouds"
tokens = tokenizer(prompt, padding="max_length", max_length=77,
                   return_tensors="pt")
with torch.no_grad():
    text_emb = text_encoder(tokens.input_ids)[0]
print(f"文本嵌入尺寸: {text_emb.shape}")  # [1, 77, 768]

from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from peft import LoraConfig, get_peft_model
import torch

# LoRA 配置
lora_config = LoraConfig(
    r=16,                          # 秩
    lora_alpha=32,                 # 缩放参数
    target_modules=["to_q", "to_v", "to_k", "to_out.0"],
    lora_dropout=0.05,
    bias="none",
)

# 把 LoRA 应用到模型上
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5"
)
unet_lora = get_peft_model(pipe.unet, lora_config)
unet_lora.print_trainable_parameters()
# 可训练参数: 约 3M / 总参数: 约 860M（约占 0.3%）

from diffusers import DiffusionPipeline
import torch

# 加载 DreamBooth 训练好的模型
pipe = DiffusionPipeline.from_pretrained(
    "./dreambooth-sks-dog",  # 训练好的 checkpoint
    torch_dtype=torch.float16
).to("cuda")

# 用 "sks dog" 生成特定的狗
images = pipe(
    "a photo of sks dog in the Eiffel Tower",
    num_inference_steps=50,
    guidance_scale=7.5
).images

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import torch
import cv2
import numpy as np

# 加载 ControlNet 模型（Canny edge）
controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/sd-controlnet-canny",
    torch_dtype=torch.float16
)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    controlnet=controlnet,
    torch_dtype=torch.float16
).to("cuda")

# 提取 Canny edge
image = load_image("input.png")
image_np = np.array(image)
low_threshold, high_threshold = 100, 200
canny_image = cv2.Canny(image_np, low_threshold, high_threshold)
canny_image = canny_image[:, :, None]  # 增加通道
canny_image = np.concatenate([canny_image] * 3, axis=2)

# ControlNet 推理
result = pipe(
    prompt="a beautiful landscape, detailed, 8k",
    image=canny_image,
    num_inference_steps=30,
    controlnet_conditioning_scale=1.0,
).images[0]

from diffusers import StableDiffusionInstructPix2PixPipeline
import torch
from diffusers.utils import load_image

pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    "timbrooks/instruct-pix2pix",
    torch_dtype=torch.float16,
    safety_checker=None
).to("cuda")

image = load_image("horse.png")
result = pipe(
    "turn the horse into a zebra",
    image=image,
    num_inference_steps=30,
    image_guidance_scale=1.5,  # 对原图的保真度
    guidance_scale=7.5          # 文本指令强度
).images[0]

from diffusers import CogVideoXPipeline
import torch

pipe = CogVideoXPipeline.from_pretrained(
    "THUDM/CogVideoX-5b",
    torch_dtype=torch.bfloat16
).to("cuda")

video = pipe(
    prompt="A serene lake with rippling water, birds flying overhead",
    num_videos_per_prompt=1,
    num_inference_steps=50,
    num_frames=49,
    guidance_scale=6,
).frames[0]

from audiocraft.models import MusicGen
import torchaudio

model = MusicGen.get_pretrained("facebook/musicgen-large")
model.set_generation_params(duration=30)  # 生成 30 秒

descriptions = ["happy jazz piano with upbeat rhythm"]
wav = model.generate(descriptions)
torchaudio.save("music.wav", wav[0].cpu(), sample_rate=32000)

from diffusers import StableDiffusionPipeline
import torch

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16
)

# 内存优化
pipe.enable_attention_slicing()           # 注意力切片
pipe.enable_vae_slicing()                 # VAE 切片
pipe.enable_model_cpu_offload()           # CPU 卸载

# xformers 加速（如已安装）
try:
    pipe.enable_xformers_memory_efficient_attention()
    print("xformers 已启用")
except:
    print("未找到 xformers，使用默认注意力机制")

import json
import urllib.request
import urllib.parse

def queue_prompt(prompt_workflow, server_address="127.0.0.1:8188"):
    """通过 ComfyUI API 执行工作流"""
    p = {"prompt": prompt_workflow}
    data = json.dumps(p).encode("utf-8")
    req = urllib.request.Request(
        f"http://{server_address}/prompt",
        data=data,
        headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req) as response:
        return json.loads(response.read())

# ComfyUI 工作流（JSON 格式）
workflow = {
    "1": {
        "class_type": "CheckpointLoaderSimple",
        "inputs": {"ckpt_name": "v1-5-pruned-emaonly.ckpt"}
    },
    "2": {
        "class_type": "CLIPTextEncode",
        "inputs": {
            "text": "a beautiful sunset over mountains",
            "clip": ["1", 1]
        }
    },
    "3": {
        "class_type": "KSampler",
        "inputs": {
            "model": ["1", 0],
            "positive": ["2", 0],
            "negative": ["4", 0],
            "latent_image": ["5", 0],
            "seed": 42,
            "steps": 30,
            "cfg": 7.5,
            "sampler_name": "euler",
            "scheduler": "karras",
            "denoise": 1.0
        }
    }
}

result = queue_prompt(workflow)
print(f"Prompt ID: {result['prompt_id']}")

from diffusers import OnnxStableDiffusionPipeline
import numpy as np

# 用 ONNX 运行时推理（CPU/GPU 均可）
pipe = OnnxStableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    revision="onnx",
    provider="CUDAExecutionProvider",
)

image = pipe("a mountain lake at dawn").images[0]