如何在Google Colab保存微调后的StableDiffusion并部署到Amazon Sagemaker？

一、导出Colab中训练好的模型 #

你用Textual Inversion微调的是自定义嵌入文件（不是完整SD模型），导出步骤如下：

• 训练结束后，在Colab文件管理器中找到生成的learned_embeds.bin（嵌入权重）和token_identifier.txt（自定义触发词）两个文件，直接右键下载到本地；或者用代码压缩后存到Google Drive：

  import zipfile
  from google.colab import drive
  
  # 压缩嵌入文件
  with zipfile.ZipFile("textual_inversion_assets.zip", "w") as zip_ref:
      zip_ref.write("learned_embeds.bin")
      zip_ref.write("token_identifier.txt")
  
  # 挂载Drive并保存
  drive.mount('/content/drive')
  !cp textual_inversion_assets.zip /content/drive/MyDrive/
  

• 如果训练的是LoRA模块，直接导出包含adapter_config.json和adapter_model.bin的整个文件夹即可。

二、部署成公共端点的可行方案 #

方案1：Hugging Face推理端点（适配HF生态，步骤最简） #

1. 注册Hugging Face账号，创建一个新的模型仓库（类型选「文本到图像」）。
2. 上传你的嵌入/LoRA文件，同时上传推理脚本app.py，示例代码如下：

   from diffusers import StableDiffusionPipeline
   import torch
   from PIL import Image
   import base64
   from io import BytesIO
   
   # 加载基础SD模型
   pipe = StableDiffusionPipeline.from_pretrained(
       "runwayml/stable-diffusion-v1-5",
       torch_dtype=torch.float16
   )
   # 加载你的自定义嵌入
   pipe.load_textual_inversion("learned_embeds.bin", token="你的自定义触发词")
   pipe.to("cuda")
   
   def predict(prompt):
       # 生成图像
       image = pipe(prompt).images[0]
       # 转为Base64方便返回
       buffer = BytesIO()
       image.save(buffer, format="PNG")
       img_str = base64.b64encode(buffer.getvalue()).decode()
       return {"image_base64": img_str}
   

3. 在仓库页面开启「推理端点」，选择GPU实例（如T4）完成部署，即可获得一个公开API URL，用户通过POST请求传入提示词就能调用。

方案2：Amazon SageMaker（简化版流程） #

1. 上传模型到S3：登录AWS控制台，创建S3存储桶，将你的嵌入/LoRA文件上传至桶内。
2. 创建Sagemaker模型：
   • 选择「自定义模型」，使用Hugging Face官方提供的SD推理容器（可在AWS Marketplace找到）。
   • 配置模型数据为S3桶内的文件路径，同时上传适配Sagemaker格式的推理脚本（需处理input_fn和predict_fn函数）。
3. 部署实时端点：选择GPU实例类型（如ml.g4dn.xlarge）启动部署，完成后获得端点URL。
4. 调用端点：使用AWS SDK（boto3）或Postman发送POST请求，传入文本提示词即可获取生成结果。

三、核心注意事项 #

• 推理脚本需确保正确加载自定义嵌入/LoRA，并处理好输入输出格式（比如返回Base64编码的图像，方便前端解析）。
• 公共端点可设置权限控制：HF端点可添加API密钥，Sagemaker可配置IAM策略限制访问，避免无授权调用。
• 注意成本：GPU实例运行成本较高，闲置时可暂停端点，或选择按需实例降低开销。

内容的提问来源于stack exchange，提问作者Lodore66