FastAPI-for-Machine-Learning-Live-Demo
FastAPI + Stable Diffusion 推理 API 演示项目,展示如何将 ML 模型
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
FastAPI + Stable Diffusion 推理 API 演示项目,展示如何将 ML 模型
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
图1:项目架构示意图(来源:FourthBrain官方仓库)
想象你熬了三天三夜训练好的 Stable Diffusion 模型,终于能生成「赛博朋克猫咪」这样炫酷的图,结果只能在 Jupyter Notebook 里跑 demo。每次换 prompt 都要重新运行整个 cell,队友想测试还得把笔记本接上显示器——这大概是 2022 年前后大多数 ML 工程师的日常。
FourthBrain(一家专注于 AI 实战培训的教育机构)的团队当时也遇到了同样的困境。他们需要一套能在云端运行的、RESTful 的 AI 推理接口,让学员不用装任何本地环境,只要发一个 HTTP 请求,就能得到一张 AI 生成的图片。解决方案就是本文要分析的这个小项目:FastAPI-for-Machine-Learning-Live-Demo。
FourthBrain 是一家 AI 教育科技公司,隶属于 deeplearning.ai 生态,专注于培养能将大模型(LLM)落地的工程人才。他们提供从模型微调到 LLMOps 全链路的实战课程,而这个项目正是课程配套的演示代码。
项目创建于 2022 年 12 月(当时 Stable Diffusion v1.4 刚发布不久),使用 FastAPI 作为 Web 框架,Stable Diffusion(CompVis/stable-diffusion-v1-4)作为推理后端,HuggingFace Diffusers 库作为推理引擎。整个仓库只有 8 个文件、11KB,却清晰展示了「如何将 ML 模型封装成 HTTP 服务」的完整工程路径。
项目提供两条生成图片的 HTTP 路径,均为 GET 接口:
GET /generate?prompt=赛博朋克猫咪&seed=42&num_inference_steps=50&guidance_scale=7.5
GET /generate-memory?prompt=赛博朋克猫咪&seed=42&num_inference_steps=50&guidance_scale=7.5
/generate 将图片写入本地文件 image.png 后通过 FileResponse 返回,适合调试和查看中间结果。/generate-memory 则直接将图片写入内存 BytesIO 流后通过 StreamingResponse 返回,不落盘,性能更优,是生产环境推荐用法。
两条路径的底层都调用了 ml.py 中的 obtain_image() 函数,签名完全一致:
def obtain_image(
prompt: str,
seed: int | None = None,
num_inference_steps: int = 50, # 推理步数,越高质量越好但越慢
guidance_scale: float = 7.5, # CFG引导强度,控制prompt遵循度
) -> Image
obtain_image 使用 torch.Generator("cuda") 设置随机种子以保证可复现性——给定相同 seed、相同步数、相同 prompt,每次生成的图片完全一致。这对调试和 A/B 测试非常重要。
项目虽然体量小,但技术栈覆盖了 ML 部署的典型组件:
| 组件 | 技术选型 | 说明 |
|---|---|---|
| Web 框架 | FastAPI | 异步、高性能、自动 OpenAPI 文档 |
| ML 推理 | diffusers 0.4.0 | HuggingFace 官方推理库 |
| 模型 | CompVis/stable-diffusion-v1-4 | HF Hub 模型,fp16 精度 |
| 硬件 | torch.float16 + CUDA | 减少显存占用,加速推理 |
| 图像处理 | PIL (Pillow) | 图片编解码 |
值得注意的是,项目使用 revision="fp16" 加载模型,这表示从 HuggingFace 下载的是 float16 精度的权重,比原生 float32 节省一半显存。对于 6-8GB 显存的消费级 GPU(如 RTX 3060)来说,这是让推理能够跑起来的关键。
尽管代码只有 80 行,但部署它需要跨越几道坎:
第一道:CUDA 环境。 代码硬编码 .to("cuda"),没有 GPU 则直接报错。这要求部署机器有 NVIDIA 显卡并安装 CUDA 驱动(>=11.8)。没有 GPU 的开发者根本无法运行。
第二道:HuggingFace 认证。 ml.py 从本地文件 token.txt 读取 HuggingFace access token,用于下载模型权重。用户必须:① 注册 HuggingFace 账号;② 申请 access token;③ 同意 CompVis/stable-diffusion-v1-4 的模型协议(该模型采用非商业 license)。这个流程对初学者并不友好。
第三道:模型权重。 fp16 精度的 Stable Diffusion v1.4 权重约 4-5GB,加上 CUDA 依赖,第一次部署可能需要 20GB+ 磁盘空间。
第四道:缺失容器化。 项目没有任何 Dockerfile 或 docker-compose.yml,无法「一键启动」。所有依赖和配置都需要手动处理。
不过,FourthBrain 的课程视频(YouTube)提供了完整的部署演示,学员可以在有指导的情况下完成整个流程。这也是这个项目定位的一部分——教育演示,而非生产级服务。
模型 license 限制。 CompVis/stable-diffusion-v1-4 采用非商业 license,禁止将其直接商用。任何想基于此做商业 AI 图片服务的开发者都需要换用明确支持商用的模型(如 Stable Diffusion XL 的 Apache 2.0 版本)。
v1.4 已过时。 项目使用 2022 年底的 v1.4 版本,当前(2026年)已迭代到 SDXL、SD 3、FLUX 等多代架构。v1.4 在图片质量、生成速度上已无优势。
无 Web 前端。 项目只有 API 端点,没有 HTML/JS 前端界面。需要自己写前端来调用 /generate 接口。对于想快速体验的用户来说,缺少「开箱即用」的 WebUI。
这个项目代表了 ML 工程化早期的一个缩影:把 Jupyter Notebook 中的推理代码,封装成 REST API,供其他服务调用。它所解决的问题——模型推理的可访问性和可复用性——至今仍是 LLMOps 的核心议题。
对比来看,2024-2026 年的主流方案已经演进到:Docker 容器化部署 -> Kubernetes 自动扩缩容 -> vLLM/TGI 等专用推理引擎优化吞吐量。而这个 2022 年的 demo 用最朴素的方式,展示了 ML 推理服务的本质:一个输入 prompt、输出图片的 HTTP 接口。
图2:项目实际运行效果(来源:FourthBrain官方仓库)
FourthBrain/FastAPI-for-Machine-Learning-Live-Demo 是一个典型的教学导向 ML 推理 API 项目,代码简洁、结构清晰,非常适合作为 ML 部署入门教材。但作为生产级服务,它在容器化、模型 license、版本迭代等方面存在明显不足。
如果你想学习如何用 FastAPI 封装 ML 模型,这个项目是很好的起点;如果你的目标是快速部署一个能商用的 AI 图片生成服务,这个项目还需要大量改造(换成 SDXL/LlamaGen 等商用兼容模型、增加 Docker 支持、加入错误处理和并发控制)。