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_packages ではなく python_requirements を使用してください。
• cuda は torch のホイールに合わせてください（例：torch==2.7.1+cu128 には 12.8 を組み合わせます）。
• モデルをファインチューニングできる場合は、train: train.py:train を追加してください。
• cog push のみでプッシュできるようにするには、image: r8.im/owner/name を追加してください。

連続バッチ処理を行う非同期 predictor の場合：

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 のウェブ UI に表示されます。
• 数値の範囲には 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

連続バッチ処理を行う非同期 predictor（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

ディスク上のアセット（例：音声サンプルを格納した voices/ ディレクトリ）から動的に choices を設定する場合：

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: ...

重みを高速に読み込む

コールドブートは、ユーザーが体感するレイテンシの大部分を占めます。簡単な順に、次の 3 つの方法があります：

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 中に 1 回だけダウンロードします（例：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 や、ユーザーが predict の実行時に渡す重みの 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()

HF、CivitAI、Replicate、および任意の .safetensors URL に対応する、本番環境向けの実装については、replicate/cog-flux/weights.py を参照してください。

複数 LoRA の合成

URL が変更された場合にのみ再読み込みし、個別のスケールで 2 つの 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
  

• メタデバイスと 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)
  

• 3 つのコピーを読み込むのではなく、複数のパイプライン（例：ベース、img2img、インペイント）間で 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

ヒント：

• 重みが約 1 GB を超えるモデルには --separate-weights を使用します。コールドブートとレジストリへのプッシュが高速になります。
• run: ステップで --mount=type=cache,target=/root/.cache/pip を使用すると、ビルド間で pip のキャッシュを共有できます。
• トークンがイメージの履歴に残らないように、ARG ではなく --secret を使用します。
• デフォルトの Cog ベースイメージ（--use-cog-base-image=true）は、独自のイメージを作成するより高速です。

トレーニング

モデルがファインチューニングに対応している場合は、cog.yaml に train: train.py:train を追加し、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 を使用します。
• すべての入力について必ず説明を記述します。説明のないスキーマは 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 — メタデバイスによる読み込み、ウォームアップコンパイル
• https://github.com/replicate/cog-vllm — 継続的バッチ処理とトレーニングをパッケージ化する非同期 LLM サーバー
• https://github.com/replicate/cog-comfyui — Cog モデルとしての ComfyUI ワークフロー、カスタムノードのヘルパー
• https://github.com/replicate/flux-fine-tuner — 複数 LoRA の合成、パイプラインコンポーネントの共有
• https://github.com/replicate/vibevoice — 動的な choices を使った TTS、最小限の cog.yaml
• https://github.com/replicate/pget — 重みを並列に取得するツール