文本到图像生成（Text-to-Image）是2024年最火的AI应用方向之一。但部署 Stable Diffusion (SD) 到昇腾NPU上，远比部署ResNet-50复杂得多。

它包含三个独立的子模型（CLIP文本编码器、UNet扩散模型、VAE解码器），总参数量超过1GB，推理链路长，且涉及20-50次迭代循环。显存管理、内存碎片化和动态Shape是核心难点。

这篇将手把手教你如何在昇腾NPU上高效部署Stable Diffusion，涵盖模型适配、显存优化、推理加速和工程化部署。

---

一、架构挑战：为什么SD在NPU上难跑？

阶段	模型	功能	耗时占比	NPU适配难点
Step 1	CLIP Text Encoder	将文本转为向量 [1, 77, 768]	~1%	静态图优化，需冻结参数
Step 2	UNet (Diffusion)	迭代去噪 (20-50步)	~98%	最大瓶颈：循环依赖、显存爆炸、动态Shape
Step 3	VAE Decoder	潜空间转像素 [1, 3, 512, 512]	~1%	卷积层多，需切片优化

核心痛点：UNet的20次迭代中，每一步都需要分配巨大的中间激活值（Activation）。在FP32下，单张512x512图的显存占用可能超过10GB，极易OOM。

---

二、核心代码实现：昇腾NPU上的SD部署

1. 配置与初始化

import torch
import torch.nn as nn
from dataclasses import dataclass
from typing import Optional, List
import numpy as np
import time
from diffusers import StableDiffusionPipeline, UNet2DConditionModel, AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import PNDMScheduler

@dataclass
class SDConfig:
    model_version: str = "runwayml/stable-diffusion-v1-5"
    num_inference_steps: int = 20
    guidance_scale: float = 7.5
    image_size: int = 512
    device: str = "npu:0"
    
    # 关键优化开关
    use_amp: bool = True           # FP16推理 (必开)
    use_memory_efficient_attention: bool = True # xformers/FlashAttention
    enable_vae_tiling: bool = True   # VAE分块解码 (省显存)
    enable_unet_slicing: bool = True # UNet切片 (省显存)
    compile_model: bool = True       # torch.compile (NPU专属优化)

class AscendStableDiffusion:
    def __init__(self, config: SDConfig):
        self.config = config
        self.device = config.device
        
        print(f"🚀 初始化 Stable Diffusion on Ascend NPU...")
        print(f"  设备: {self.device}")
        print(f"  精度: {'FP16' if config.use_amp else 'FP32'}")
        
        # 设置NPU环境
        torch.npu.set_device(0)
        torch.npu.set_benchmark_mode(True)
        
        self.tokenizer = None
        self.text_encoder = None
        self.unet = None
        self.vae = None
        self.scheduler = None
    
    def load_models(self):
        """加载并优化模型"""
        print("\n=== 加载模型 ===")
        
        # 1. 加载Tokenizer & CLIP Text Encoder
        print("Loading CLIP Text Encoder...")
        self.tokenizer = CLIPTokenizer.from_pretrained(
            self.config.model_version, subfolder="tokenizer"
        )
        self.text_encoder = CLIPTextModel.from_pretrained(
            self.config.model_version, subfolder="text_encoder"
        ).to(self.device).eval()
        
        # 冻结CLIP参数
        for param in self.text_encoder.parameters():
            param.requires_grad = False
            
        if self.config.use_amp:
            self.text_encoder = self.text_encoder.half()
        
        # 2. 加载UNet (核心瓶颈)
        print("Loading UNet...")
        self.unet = UNet2DConditionModel.from_pretrained(
            self.config.model_version, subfolder="unet"
        ).to(self.device).eval()
        
        if self.config.use_amp:
            self.unet = self.unet.half()
            
        # 启用UNet优化 (如果支持)
        if self.config.enable_unet_slicing:
            self.unet.enable_attention_slicing()
            
        # 3. 加载VAE Decoder
        print("Loading VAE Decoder...")
        self.vae = AutoencoderKL.from_pretrained(
            self.config.model_version, subfolder="vae"
        ).to(self.device).eval()
        
        if self.config.use_amp:
            self.vae = self.vae.half()
            
        if self.config.enable_vae_tiling:
            self.vae.enable_tiling()
            
        # 4. 加载调度器
        print("Loading Scheduler...")
        self.scheduler = PNDMScheduler.from_pretrained(
            self.config.model_version, subfolder="scheduler"
        )
        self.scheduler.set_timesteps(self.config.num_inference_steps)
        
        total_params = sum(p.numel() for p in self.text_encoder.parameters()) + \
                       sum(p.numel() for p in self.unet.parameters()) + \
                       sum(p.numel() for p in self.vae.parameters())
        print(f"\n✅ 模型加载完成，总参数量: {total_params / 1e6:.2f}M")
        
        # 尝试编译UNet (NPU性能提升关键)
        if self.config.compile_model:
            print("⚡ 正在编译 UNet (torch.compile)...")
            try:
                self.unet = torch.compile(self.unet, mode="max-autotune", backend="ascend")
                print("   UNet编译成功！")
            except Exception as e:
                print(f"   编译失败 (可能是版本问题), 使用原生模式: {e}")

    @torch.no_grad()
    def generate(self, prompt: str, seed: Optional[int] = None) -> np.ndarray:
        """
        生成图像主流程
        
        流程：
          1. 编码Prompt (CLIP)
          2. 迭代去噪 (UNet)
          3. 解码潜变量 (VAE)
        """
        if seed is not None:
            torch.manual_seed(seed)
            torch.npu.manual_seed(seed)
            
        t_start = time.time()
        
        # --- Step 1: 文本编码 ---
        t1 = time.time()
        text_input = self.tokenizer(
            prompt, padding="max_length", max_length=77, 
            truncation=True, return_tensors="pt"
        ).input_ids.to(self.device)
        
        with torch.npu.amp.autocast(enabled=self.config.use_amp):
            text_embeddings = self.text_encoder(text_input)[0]
        
        # CFG (Classifier-Free Guidance)
        uncond_input = self.tokenizer(
            "", padding="max_length", max_length=77, 
            truncation=True, return_tensors="pt"
        ).input_ids.to(self.device)
        
        with torch.npu.amp.autocast(enabled=self.config.use_amp):
            uncond_embeddings = self.text_encoder(uncond_input)[0]
            
        text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
        t_clip = time.time() - t1
        
        # --- Step 2: UNet 迭代去噪 (核心瓶颈) ---
        t2 = time.time()
        latents = self._denoise_loop(text_embeddings)
        t_unet = time.time() - t2
        
        # --- Step 3: VAE 解码 ---
        t3 = time.time()
        with torch.npu.amp.autocast(enabled=self.config.use_amp):
            image_latents = latents / self.vae.config.scaling_factor
            decoded_image = self.vae.decode(image_latents).sample
        t_vae = time.time() - t3
        
        # Post-processing
        image = (decoded_image / 2 + 0.5).clamp(0, 1)
        image = image.cpu().permute(0, 2, 3, 1).numpy()
        
        total_time = time.time() - t_start
        print(f"\n🎨 生成完成:")
        print(f"  总耗时: {total_time*1000:.1f}ms")
        print(f"  CLIP: {t_clip*1000:.1f}ms | UNet: {t_unet*1000:.1f}ms | VAE: {t_vae*1000:.1f}ms")
        
        return image[0]

    def _denoise_loop(self, text_embeddings: torch.Tensor) -> torch.Tensor:
        """
        UNet 迭代去噪
        
        优化策略：
          1. 复用噪声缓冲区
          2. 混合精度推理
          3. 避免不必要的梯度计算
        """
        batch_size = text_embeddings.shape[0] // 2
        latent_channels = self.unet.config.in_channels
        height = width = self.config.image_size // 8
        
        # 初始化噪声
        noise = torch.randn((batch_size, latent_channels, height, width), device=self.device)
        
        self.scheduler.set_timesteps(self.config.num_inference_steps)
        
        for i, t in enumerate(self.scheduler.timesteps):
            # 扩展噪声以匹配CFG
            latent_model_input = torch.cat([noise] * 2)
            
            # 预测噪声残差
            with torch.npu.amp.autocast(enabled=self.config.use_amp):
                noise_pred = self.unet(
                    latent_model_input, t, encoder_hidden_states=text_embeddings
                ).sample
            
            # 执行CFG
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + self.config.guidance_scale * (noise_pred_text - noise_pred_uncond)
            
            # 调度器更新
            noise = self.scheduler.step(noise_pred, t, noise).prev_sample
            
            # 可选：每10步打印进度
            if (i + 1) % 10 == 0:
                print(f"  去噪进度: {i+1}/{self.config.num_inference_steps}")
                
        return noise

---

三、昇腾NPU专用优化技巧

1. 显存优化三板斧

技术	原理	效果	昇腾适配
FP16 混合精度	权重和激活值用FP16	显存↓50%, 速度↑2x	model.half() + torch.npu.amp
VAE Tiling	将大图切分成小块解码	显存↓70% (解决OOM)	vae.enable_tiling()
UNet Attention Slicing	分块计算Attention矩阵	显存↓40%	unet.enable_attention_slicing()

注意：在昇腾上，务必开启 torch.npu.set_optimize_mode(True) 以启用内存碎片整理。

2. torch.compile 加速

昇腾NPU对静态图优化能力极强。使用 torch.compile 可以将Python层面的控制流转化为高效的NPU指令。

# 在加载UNet后调用
if hasattr(torch, 'compile'):
    try:
        # 指定Ascend后端 (需安装 ascend-pytorch 或特定版本)
        self.unet = torch.compile(self.unet, mode="max-autotune", backend="ascend")
        print("✅ UNet 已编译为NPU高效算子")
    except:
        # 降级方案：使用inductor或native
        self.unet = torch.compile(self.unet, mode="reduce-overhead")

3. 自定义ATC编译 (进阶)

对于生产环境，建议将模型导出为ONNX，然后使用ATC工具编译为.om文件。

# 1. 导出ONNX (仅UNet部分)
python export_onnx.py --model_path ./unet.pt --output ./unet.onnx

# 2. ATC 编译 (开启INT8量化或FP16融合)
atc \
  --model=./unet.onnx \
  --output=./unet_ascend \
  --framework=5 \
  --input_shape="latent:1,4,64,64;time:1;text_emb:1,77,768" \
  --precision_mode=mixed_precision \
  --op_select_implmode=high_precision \
  --soc_version=Ascend910B

---

四、常见陷阱与解决方案

问题现象	原因分析	解决方案
显存瞬间爆满 (OOM)	UNet中间激活值过大	1. 强制开启 enable_tiling 2. 降低 image_size (如512→384) 3. 减小 num_inference_steps
推理速度慢 (<5 FPS)	Python循环开销大	1. 使用 torch.compile 2. 检查是否频繁CPU↔NPU拷贝 3. 开启 benchmark_mode
生成图片模糊/崩坏	量化误差或精度不足	1. 关闭INT8，使用FP16 2. 增加 guidance_scale 3. 校准CLIP输入范围
动态Shape报错	不同Prompt长度不一致	1. 统一Pad到77 tokens 2. 使用 torch.jit.script 处理变长序列
多用户并发冲突	单卡资源争抢	1. 使用 模型实例池 (每个用户独立进程) 2. 限制 max_batch_size

---

五、工程化部署：高并发服务架构

为了支撑生产流量，不能直接运行Python脚本，需要构建微服务。

1. 异步推理服务 (FastAPI)

from fastapi import FastAPI, HTTPException
import asyncio

app = FastAPI()
sd_service = AscendStableDiffusion(SDConfig())

@app.post("/generate")
async def generate_image(prompt: str, seed: int = 42):
    # 异步执行推理，不阻塞IO
    loop = asyncio.get_event_loop()
    image = await loop.run_in_executor(None, sd_service.generate, prompt, seed)
    
    # 返回Base64图片
    import base64
    from PIL import Image
    pil_img = Image.fromarray((image * 255).astype(np.uint8))
    buffer = io.BytesIO()
    pil_img.save(buffer, format="PNG")
    img_str = base64.b64encode(buffer.getvalue()).decode()
    
    return {"image": f"data:image/png;base64,{img_str}"}

2. 请求队列与动态Batching

当多个用户同时请求时，利用 Dynamic Batching 合并UNet的多次调用，大幅提升吞吐量。

# 伪代码：在UNet层前加入Batcher
class BatchedInferencer:
    async def infer(self, requests):
        # 收集一批请求
        # 合并它们的prompt (如果需要) 或并行处理
        # 调用一次UNet批量去噪
        pass

---

六、总结：昇腾NPU部署SD最佳实践

1. 精度优先: 必须使用 FP16 (half())，这是提速和减显存的基础。
2. 显存急救: 遇到OOM先开 VAE Tiling 和 Attention Slicing，不要盲目降分辨率。
3. 编译加速: 务必尝试 torch.compile 或 ATC编译，NPU的静态图优化能带来30%-50%的性能提升。
4. 迭代次数: 默认20步，生产环境可尝试 DDIM 采样器将步数降至10-15步，速度翻倍。
5. 监控显存: 实时监控 npu-smi info，确保显存碎片率低于20%。

一句话建议：在昇腾上做SD，“先FP16，再Tiling，最后Compile”。先用FP16跑通，再用Tiling解决OOM，最后用Compile压榨性能。