AI绘画Stable Diffusion代码重构：从烂代码到优雅代码

标题提到的Stable Diffusion在本文写作时（2022年6月）尚未正式发布（预计2022年8月开源）。本文基于AI绘画项目的一般架构和代码重构的经验，分享如何把一个"能跑但很烂"的AI绘画项目，重构成优雅、可维护、可扩展的代码。

最近我参与了一个AI绘画项目的重构。项目最开始是几个算法工程师快速搭起来的，能跑通，但代码质量很差：一个文件几千行、函数几百行、变量名看不懂、没有注释、没有测试、改一个地方要翻半天。

我花了两周时间，把这个项目从头到尾重构了一遍。重构之后，代码结构清晰了，功能扩展容易了，新人也能快速上手了。

本文分享这次重构的经验，包括代码坏味道识别、重构步骤、设计模式应用、测试保障等。不管你做的是不是AI绘画，这些重构的思路和方法都是通用的。

一、为什么要重构

先说说，为什么要重构。

1. 烂代码的代价

烂代码，短期看能跑，长期看代价很大：

改一个功能要花很长时间，因为看不懂代码
容易引入bug，因为不知道改了会影响什么
新人上手慢，要花很多时间理解代码
无法扩展，加新功能要改很多地方
技术债越积越多，最后可能要重写

很多项目，最开始快速迭代，代码写得很烂。后来功能越来越多，代码越来越乱，最后改不动了，只能推倒重来。

2. 重构的时机

什么时候该重构？

加新功能很困难的时候
bug频繁出现，修了一个又来一个的时候
新人上手很慢的时候
代码 review 很痛苦的时候
你自己都不想看自己写的代码的时候

重构不是"等有时间了再做"，而是"现在不做，以后更难做"。

3. 重构的目标

重构的目标，不是"代码更漂亮"，而是：

更容易理解
更容易修改
更容易测试
更容易扩展

代码漂亮是手段，不是目的。

二、识别代码坏味道

重构之前，先识别代码中的"坏味道"。

1. 过长的函数

一个函数超过50行，就要考虑拆分了。超过100行，基本肯定有问题。

过长的函数，往往做了不止一件事，难以理解和测试。

2. 过大的类/文件

一个类或文件超过500行，就要考虑拆分了。超过1000行，基本肯定有问题。

过大的类，往往承担了太多职责，违反单一职责原则。

3. 命名不清晰

变量名、函数名、类名，不能清晰表达含义。

比如：

data、info、temp：不知道是什么数据
func1、doSomething：不知道做什么
a、b、c：完全看不懂

好的命名，是代码可读性的基础。

4. 重复代码

同样或类似的代码，出现在多个地方。

重复代码的问题是：改一个地方，要记得改所有重复的地方。漏改一个，就是bug。

5. 过长的参数列表

函数参数超过5个，就要考虑封装成对象或配置类。

参数太多，调用的时候容易传错，也难以理解每个参数的含义。

6. 魔法数字和魔法字符串

代码里出现莫名其妙的数字或字符串，不知道是什么意思。

比如：

if status == 3:  # 3是什么意思？
    do_something()

应该用常量或枚举代替。

7. 嵌套过深

if里面套if，套了三四层，很难看懂。

嵌套过深，应该用提前返回（early return）或卫语句（guard clause）来简化。

8. 注释过多或过少

注释过多：代码本身不清楚，靠注释解释，应该优化代码让它自解释
注释过少：复杂的逻辑没有注释，别人看不懂

最好的代码是自解释的，复杂的地方加注释说明"为什么"，而不是"做什么"。

三、重构的步骤

说说我的重构步骤。

1. 第一步：理解代码

重构之前，先理解代码在做什么。

通读代码，画出流程图
理解每个函数的输入输出
理解模块之间的调用关系
找到核心逻辑和辅助逻辑

不要一上来就改，先理解清楚。理解错了，重构就是在制造bug。

2. 第二步：建立测试

重构之前，一定要有测试。

没有测试的重构，就是赌博。你不知道改完之后，功能是不是还正常。

如果项目没有测试，先补测试：

核心功能写单元测试
关键流程写集成测试
至少要有"冒烟测试"，确保主流程能跑通

有了测试，重构才有安全网。改完之后跑一遍测试，通过了就说明没改坏。

3. 第三步：小步重构

重构要小步进行，不要一下子大改。

每次只改一个地方
改完跑测试，确保没问题
提交一次代码
再改下一个地方

小步重构的好处是：如果出了问题，很容易定位是哪一步改坏的。大改的话，出了问题都不知道从哪查。

4. 第四步：先易后难

从简单的地方开始重构，建立信心，再处理复杂的部分。

先改命名：把看不懂的变量名、函数名改清楚
再拆函数：把过长的函数拆成小函数
再拆模块：把过大的文件拆成多个模块
最后优化架构：调整模块之间的关系

5. 第五步：持续验证

重构过程中，持续验证：

跑单元测试
跑集成测试
手动测试核心功能
对比重构前后的输出是否一致

AI绘画项目，我会用同一张图、同一个seed，对比重构前后的生成结果，确保完全一致。

四、AI绘画项目的重构实践

说说我在AI绘画项目中的具体重构实践。

1. 原始代码的问题

原始代码的问题：

一个main.py文件，3000多行
一个generate()函数，500多行，做了加载模型、预处理、推理、后处理、保存图片所有事情
配置写死在代码里，改参数要改代码
没有模块化，模型、数据、推理混在一起
没有错误处理，出错直接崩溃
没有日志，出问题不知道哪一步错了

2. 重构后的结构

重构后的目录结构：

ai_painting/
├── config/
│   ├── __init__.py
│   └── settings.py       # 配置管理
├── models/
│   ├── __init__.py
│   ├── base.py           # 模型基类
│   ├── diffusion.py      # 扩散模型
│   └── vae.py            # VAE模型
├── data/
│   ├── __init__.py
│   ├── loader.py         # 数据加载
│   └── processor.py      # 数据预处理
├── inference/
│   ├── __init__.py
│   ├── pipeline.py       # 推理流水线
│   └── sampler.py        # 采样器
├── utils/
│   ├── __init__.py
│   ├── image.py          # 图像处理工具
│   └── logger.py         # 日志工具
├── tests/
│   ├── test_models.py
│   ├── test_inference.py
│   └── test_utils.py
└── main.py               # 入口

每个模块职责单一，结构清晰。

3. 具体的重构动作

动作一：拆分大函数

把500行的generate()函数拆成：

load_model()：加载模型
preprocess()：预处理输入
diffuse()：扩散推理
decode()：VAE解码
postprocess()：后处理
save_image()：保存图片

每个函数只做一件事，不超过50行。

动作二：提取配置

把写死在代码里的配置，提取到配置文件：

# config/settings.py
from dataclasses import dataclass

@dataclass
class ModelConfig:
    model_path: str = "models/stable-diffusion"
    device: str = "cuda"
    dtype: str = "float16"

@dataclass
class GenerateConfig:
    steps: int = 50
    guidance_scale: float = 7.5
    seed: int = 42
    width: int = 512
    height: int = 512

用dataclass管理配置，类型安全，有默认值，容易扩展。

动作三：引入基类和接口

给模型定义基类，方便以后扩展新模型：

# models/base.py
from abc import ABC, abstractmethod

class BaseModel(ABC):
    @abstractmethod
    def load(self, config):
        pass

    @abstractmethod
    def predict(self, input_data):
        pass

以后加新模型，只要继承BaseModel，实现load和predict方法就行。

动作四：错误处理

给关键步骤加错误处理：

try:
    model = load_model(config)
except Exception as e:
    logger.error(f"加载模型失败: {e}")
    raise ModelLoadError("模型加载失败，请检查模型路径") from e

自定义异常类型，出错的时候有明确的错误信息，而不是一堆看不懂的堆栈。

动作五：加日志

关键步骤加日志，方便排查问题：

logger.info("开始加载模型...")
model = load_model(config)
logger.info("模型加载完成")

logger.info(f"开始生成，steps={config.steps}, seed={config.seed}")
image = pipeline.generate(prompt, config)
logger.info("生成完成")

出问题的时候，看日志就知道卡在哪一步了。

动作六：写测试

给核心模块写单元测试：

# tests/test_inference.py
def test_generate():
    config = GenerateConfig(steps=5, seed=42)
    pipeline = InferencePipeline()
    image = pipeline.generate("a cat", config)
    assert image is not None
    assert image.size == (512, 512)

用小steps快速测试，确保核心逻辑没问题。

五、重构中用到的设计原则

说说重构中用到的设计原则。

1. 单一职责原则（SRP）

每个类、每个函数，只做一件事。

模型类只负责模型的加载和推理
数据类只负责数据的加载和预处理
推理类只负责推理流程的编排

不要让一个类做太多事情。

2. 开闭原则（OCP）

对扩展开放，对修改关闭。

加新模型，不需要改推理流程，只要加一个新的模型类
加新的采样器，不需要改模型，只要加一个新的采样器类

通过接口和基类，实现扩展时不需要修改已有代码。

3. 依赖倒置原则（DIP）

高层模块不依赖低层模块，都依赖抽象。

推理流程依赖模型基类，不依赖具体的模型实现
这样换模型的时候，推理流程不需要改

4. 不要重复自己（DRY）

不要有重复代码。

相同的逻辑，提取成函数
相同的配置，提取成常量
相同的模式，提取成基类或工具函数

5. 保持简单（KISS）

不要过度设计。

重构是为了让代码更简单，不是更复杂。不要为了用设计模式而用设计模式。

如果一个简单的函数就能解决问题，就不要搞一堆类和接口。

六、重构的注意事项

说说重构中需要注意的事项。

1. 不要在重构的同时加新功能

重构就是重构，不要同时加新功能。

同时加新功能，出了问题不知道是重构改坏的，还是新功能有bug。

先重构，重构完了，再加新功能。

2. 不要追求完美

重构不是一次就能做到完美的。

先把最痛的地方改了，其他地方以后慢慢改。追求完美，可能永远改不完。

3. 要有耐心

重构是慢功夫，不要急。

一个3000行的文件，不是一天就能拆完的。慢慢来，小步迭代，确保每一步都正确。

4. 和团队沟通

如果是团队项目，重构之前要和团队沟通。

告诉大家为什么要重构
告诉大家重构的计划
让大家知道重构期间怎么协作
重构完了，告诉大家新的代码结构

不要一个人偷偷重构，否则别人还在改旧代码，冲突会很多。

七、重构后的效果

说说重构后的效果。

1. 代码可读性提高了

以前看代码要翻半天，现在结构清晰，命名清楚，看文件名就知道大概是做什么的。

新人上手，以前要一周，现在两三天就能理解核心代码。

2. 加新功能更容易了

以前加一个新的采样器，要改generate()函数里的很多地方。现在只要加一个新的采样器类，实现接口就行。

加新模型也是一样，不需要改推理流程。

3. bug减少了

重构过程中，发现并修复了几个隐藏的bug。重构后，代码结构清晰，bug也更容易发现和定位。

4. 测试覆盖了

重构后，核心功能都有了单元测试。以后改代码，跑一遍测试就知道有没有改坏。

5. 性能也提升了

重构过程中，顺便优化了一些性能问题：

避免了重复加载模型
减少了不必要的数据拷贝
优化了GPU内存使用

重构后，生成速度提升了约20%。

八、写在最后

代码重构，是每个程序员的必修课。

烂代码不可怕，可怕的是知道烂却不改，让技术债越积越多。

重构不是"等有时间了再做"的事情，而是"现在不做，以后更难做"的事情。小步重构，持续改进，代码质量会慢慢好起来。

AI绘画项目的重构，让我深刻体会到：好的代码，不是写出来的，是改出来的。第一版能跑通就行，然后在使用中不断重构、不断优化。

2022年了，AI绘画、大模型这些新技术发展很快。很多项目都是快速搭起来的，代码质量参差不齐。但越是快速发展，越需要好的代码质量，否则后面会越来越难维护。

最后，用一句话总结："重构的本质，不是让代码更漂亮，而是让代码更易变。易变，才能适应快速变化的需求。"

愿大家的代码，都能从烂代码，慢慢变成优雅代码。