build-models

文档

• Cog 参考（单文件）：https://cog.run/llms.txt
• cog.yaml 参考：https://cog.run/yaml
• Python predictor 参考：https://cog.run/python
• 示例：https://github.com/replicate/cog-examples
• 模板：https://github.com/replicate/cog-template

何时使用此技能

• 你有想托管在 Replicate 上的模型代码、权重或 HuggingFace/GitHub 项目。
• 你正在编写或编辑 cog.yaml、predict.py 或 train.py。
• 如需将构建好的模型推送到 Replicate，请参阅 publish-models。
• 如需运行现有的 Replicate 模型，请参阅 run-models。

前置条件

• 本地已运行 Docker。
• 已安装 Cog：brew install replicate/tap/cog 或 sh <(curl -fsSL https://cog.run/install.sh)。
• 可选：运行 cog init 以生成 cog.yaml 和 predict.py。

项目布局

规范的 Replicate 模型布局：

cog.yaml
predict.py
weights.py                 # optional download helpers
requirements.txt
cog-safe-push-configs/
  default.yaml             # see publish-models skill
.github/workflows/
  ci.yaml
script/                    # github.com/github/scripts-to-rule-them-all
  lint
  test
  push

cog.yaml 要点

适用于 GPU 模型的现代配置：

build:
  gpu: true
  cuda: "12.8"
  python_version: "3.12"
  python_requirements: requirements.txt
  system_packages:
    - libgl1
    - libglib2.0-0
predict: predict.py:Predictor

说明：

• 将 Python 固定为特定的次版本，并固定 requirements.txt 中的每个版本。浮动版本会导致冷启动失败。
• 列表变长后，使用 python_requirements，而不是内联的 python_packages。
• cuda 版本应与 PyTorch wheel 匹配（例如，12.8 对应 torch==2.7.1+cu128）。
• 如果模型支持微调，请添加 train: train.py:train。
• 添加 image: r8.im/owner/name 以启用直接运行 cog push。

支持连续批处理的异步预测器：

concurrency:
  max: 32

predict.py 要点

from cog import BasePredictor, Input, Path

class Predictor(BasePredictor):
    def setup(self) -> None:
        """One-time loads. Heavy work goes here, not in predict()."""
        self.model = load_model("weights/")

    def predict(
        self,
        prompt: str = Input(description="Text prompt for generation"),
        seed: int = Input(description="Random seed; leave blank for random", default=None),
        num_steps: int = Input(description="Number of denoising steps", ge=1, le=50, default=20),
        output_format: str = Input(description="Output image format", choices=["webp", "jpg", "png"], default="webp"),
    ) -> Path:
        """Run a single prediction."""
        if not prompt.strip():
            raise ValueError("prompt cannot be empty")
        out = self.model.generate(prompt, seed=seed, steps=num_steps)
        return Path(out)

输入规则：

• 每项输入都需要有 description。该描述会显示在模型架构和 Replicate 的网页界面中。
• 数值范围使用 ge/le，枚举使用 choices=[...]，字符串使用 regex=。
• 文件输入和输出使用 cog.Path，不要使用原始字节。
• 类似令牌的输入（HF 令牌、API 密钥）使用 cog.Secret，不要使用普通的 str。
• 分类输入的默认值必须包含在 choices 中。
• 在 predict() 中尽早验证输入，并引发 ValueError。

文本流式输出（适用于 LLM）：

from cog import BasePredictor, Input, ConcatenateIterator

class Predictor(BasePredictor):
    def predict(self, prompt: str = Input(description="Prompt")) -> ConcatenateIterator[str]:
        for token in self.model.stream(prompt):
            yield token

支持连续批处理的异步预测器（与 cog.yaml 中的 concurrency.max 配合使用）：

from cog import BasePredictor, Input, AsyncConcatenateIterator

class Predictor(BasePredictor):
    async def setup(self) -> None:
        self.engine = await load_async_engine()

    async def predict(
        self,
        prompt: str = Input(description="Prompt"),
    ) -> AsyncConcatenateIterator[str]:
        async for token in self.engine.generate(prompt):
            yield token

从磁盘上的资源动态生成 choices（例如，存放音频样本的 voices/ 目录）：

from pathlib import Path as _P
AVAILABLE_VOICES = sorted(p.stem for p in _P("voices").glob("*.wav"))

class Predictor(BasePredictor):
    def predict(
        self,
        speaker: str = Input(description="Voice", choices=AVAILABLE_VOICES, default=AVAILABLE_VOICES[0]),
    ) -> Path: ...

快速加载权重

冷启动是影响用户感知延迟的主要因素。以下是三种模式，按简单程度排序：

1. 在构建时将权重打包进镜像

适用于你希望实现零冷启动的小型或中型权重（< 5GB）。

适用于 torchvision：

import os
os.environ["TORCH_HOME"] = "."  # set before importing torch
import torch
from torchvision import models

适用于 HuggingFace：

import os
os.environ["HF_HUB_CACHE"] = "./.cache"
os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"

然后在 cog build 期间只下载一次（例如，在 run: 步骤中下载，或在构建过程中运行一个小型下载脚本）。权重会成为镜像层的一部分。

2. 使用 pget 从 weights.replicate.delivery 拉取

适用于大型权重，或需要在多个模型间共享权重的情况。pget 是 Replicate 的并行 HTTP 下载工具。

在 cog.yaml 中：

build:
  run:
    - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.8.2/pget_linux_x86_64"
    - chmod +x /usr/local/bin/pget

在 setup() 中：

import subprocess
from pathlib import Path

WEIGHTS_URL = "https://weights.replicate.delivery/default/my-model/weights.tar"
WEIGHTS_DIR = Path("weights")

class Predictor(BasePredictor):
    def setup(self) -> None:
        if not WEIGHTS_DIR.exists():
            # -x extracts tar in-memory; default concurrency is 4 * NumCPU
            subprocess.check_call(["pget", "-x", WEIGHTS_URL, str(WEIGHTS_DIR)])
        self.model = load_from(WEIGHTS_DIR)

一次下载多个文件：

manifest = "\n".join([
    f"{base}/unet.safetensors weights/unet.safetensors",
    f"{base}/vae.safetensors  weights/vae.safetensors",
    f"{base}/text_encoder.safetensors weights/text_encoder.safetensors",
])
subprocess.run(["pget", "multifile", "-"], input=manifest, text=True, check=True)

3. 使用 hf_transfer 从 HuggingFace Hub 下载

设置 HF_HUB_ENABLE_HF_TRANSFER=1，并使用 huggingface_hub.snapshot_download 或 from_pretrained。这比原版 HF 下载更快。对于受限模型，请使用 cog.Secret 输入。

用户提供权重的缓存

对于 LoRA 或用户在预测时传入的任何权重 URL，使用以 sha256 为键、带 LRU 淘汰策略的磁盘缓存：

import hashlib, shutil, subprocess
from pathlib import Path

class WeightsDownloadCache:
    def __init__(self, cache_dir: str = "/tmp/weights-cache", min_disk_free_gb: int = 10):
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.min_disk_free = min_disk_free_gb * 1024**3

    def ensure(self, url: str) -> Path:
        key = hashlib.sha256(url.encode()).hexdigest()
        target = self.cache_dir / key
        if target.exists():
            target.touch()  # bump LRU mtime
            return target
        self._evict_until_room()
        subprocess.check_call(["pget", url, str(target)])
        return target

    def _evict_until_room(self) -> None:
        while shutil.disk_usage(self.cache_dir).free < self.min_disk_free:
            entries = sorted(self.cache_dir.iterdir(), key=lambda p: p.stat().st_mtime)
            if not entries:
                return
            entries[0].unlink()

请参阅 replicate/cog-flux/weights.py，其中提供了适用于生产环境的实现，可处理 HF、CivitAI、Replicate 和任意 .safetensors URL。

多个 LoRA 组合

仅在 URL 发生变化时重新加载；使用单独的缩放比例组合两个 LoRA：

class Predictor(BasePredictor):
    def setup(self) -> None:
        self.pipe = load_base_pipeline()
        self.loaded = {"main": None, "extra": None}

    def _ensure_lora(self, slot: str, url: str | None) -> None:
        if url == self.loaded[slot]:
            return
        if self.loaded[slot] is not None:
            self.pipe.unload_lora_weights(adapter_name=slot)
        if url:
            path = self.cache.ensure(url)
            self.pipe.load_lora_weights(str(path), adapter_name=slot)
        self.loaded[slot] = url

    def predict(
        self,
        prompt: str = Input(description="Prompt"),
        lora_url: str = Input(description="Primary LoRA URL", default=None),
        lora_scale: float = Input(description="Primary LoRA scale", ge=0.0, le=2.0, default=1.0),
        extra_lora_url: str = Input(description="Optional second LoRA URL", default=None),
        extra_lora_scale: float = Input(description="Second LoRA scale", ge=0.0, le=2.0, default=1.0),
    ) -> Path:
        self._ensure_lora("main", lora_url)
        self._ensure_lora("extra", extra_lora_url)
        adapters = [s for s, u in self.loaded.items() if u]
        scales = [lora_scale if s == "main" else extra_lora_scale for s in adapters]
        if adapters:
            self.pipe.set_adapters(adapters, adapter_weights=scales)
        return Path(self.pipe(prompt).images[0].save("/tmp/out.png"))

冷启动技巧

来自生产环境中的扩散模型，例如 replicate/cog-flux 和 replicate/cog-flux-kontext：

• 在 setup() 中设置一次性能标志：

  import torch
  torch.set_float32_matmul_precision("high")
  torch.backends.cuda.matmul.allow_tf32 = True
  torch.backends.cudnn.benchmark = True
  

• 编译并预热：

  self.model = torch.compile(self.model, dynamic=True)
  _ = self.predict(prompt="warmup", num_steps=1)  # absorbs compile cost in setup
  

• 使用 meta device + assign=True 加载大型权重，避免重复分配：

  with torch.device("meta"):
      model = build_model_skeleton()
  state = torch.load("weights.pt", map_location="cpu")
  model.load_state_dict(state, assign=True)
  

• 在多个流水线（例如 base + img2img + inpaint）之间共享 VAE / 文本编码器，而不是加载三份副本。
• 对于 fp8/int8，提前保存量化后的权重并直接加载；不要在启动时量化。

本地开发

cog init                                    # scaffold cog.yaml + predict.py
cog predict -i prompt="hello"               # build + run a single prediction
cog predict -i [email protected] -o out.png  # file inputs and outputs
cog serve -p 8393                           # HTTP server matching production
cog exec python                             # interactive shell inside the build env

构建

cog build -t my-model
cog build --separate-weights -t my-model    # weights in their own image layer
cog build --secret id=hf,src=$HOME/.hf_token -t my-model

提示：

• 对权重超过约 1GB 的模型，使用 --separate-weights。这能加快冷启动和注册表推送。
• 在 run: 步骤中使用 --mount=type=cache,target=/root/.cache/pip，以便在构建之间缓存 pip。
• 使用 --secret 而不是 ARG，避免令牌保留在镜像历史记录中。
• 默认 Cog 基础镜像（--use-cog-base-image=true）比自行构建基础镜像更快。

训练

如果模型支持微调，请将 train: train.py:train 添加到 cog.yaml，并编写一个返回 TrainingOutput(weights=Path("model.tar")) 的 train() 函数。随后，预测器会通过 setup(self, weights) 或 COG_WEIGHTS 环境变量接收 URL。完整示例请参见 https://cog.run/training 和 replicate/flux-fine-tuner。

指南

• 将一次性加载放在 setup() 中；让 predict() 保持快速，且输出形状确定。
• 固定 Python 和所有依赖项的版本。如果 torch 版本较旧，请使用 numpy<2。
• 始终为每个输入添加说明。没有说明的 schema 无法在 Web UI 中使用。
• 文件使用 cog.Path，令牌使用 cog.Secret。
• 将 pget 固定到特定版本（v0.8.2），以确保可复现性。
• 每当调用 HuggingFace Hub 时，设置 HF_HUB_ENABLE_HF_TRANSFER=1。
• 权重加载后，设置 TRANSFORMERS_OFFLINE=1，以防止运行时查询 HF。
• 推送前使用 cog predict 进行测试。如果它无法在本地运行，就无法在生产环境中运行。

生产环境参考

• https://github.com/replicate/cog-examples — 最简模式（resnet、hello-world、流式传输、训练）
• https://github.com/replicate/cog-template — 用于创建新模型仓库的脚手架工具
• https://github.com/replicate/cog-flux — 多变体 FLUX 模型、权重缓存、fp8 + torch.compile
• https://github.com/replicate/cog-flux-kontext — meta-device 加载、预热编译
• https://github.com/replicate/cog-vllm — 支持连续批处理的异步 LLM 服务、将训练作为打包步骤
• https://github.com/replicate/cog-comfyui — 将 ComfyUI 工作流作为 Cog 模型、custom-node 辅助工具
• https://github.com/replicate/flux-fine-tuner — 多 LoRA 组合、共享流水线组件
• https://github.com/replicate/vibevoice — 使用动态 choices 的 TTS、精简的 cog.yaml
• https://github.com/replicate/pget — 并行权重获取器