build-models

Docs

• Cog reference (fichier unique): https://cog.run/llms.txt
• Référence cog.yaml: https://cog.run/yaml
• Référence Python predictor: https://cog.run/python
• Exemples: https://github.com/replicate/cog-examples
• Template: https://github.com/replicate/cog-template

Quand utiliser cette skill

• Vous avez du code de modèle, des poids, ou un projet HuggingFace/GitHub que vous voulez héberger sur Replicate.
• Vous écrivez ou modifiez un cog.yaml, predict.py, ou train.py.
• Pour pousser un modèle construit vers Replicate, voir publish-models.
• Pour exécuter des modèles Replicate existants, voir run-models.

Prérequis

• Docker exécuté localement.
• Cog installé: brew install replicate/tap/cog ou sh <(curl -fsSL https://cog.run/install.sh).
• Optionnel: cog init pour générer cog.yaml et predict.py.

Disposition du projet

La disposition canonique d'un modèle Replicate:

cog.yaml
predict.py
weights.py                 # assistants de téléchargement optionnels
requirements.txt
cog-safe-push-configs/
  default.yaml             # voir skill publish-models
.github/workflows/
  ci.yaml
script/                    # github.com/github/scripts-to-rule-them-all
  lint
  test
  push

Essentiels de cog.yaml

Une config moderne pour un modèle 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

Notes:

• Épinglez Python à une version mineure spécifique, et épinglez chaque ligne dans requirements.txt. Les versions flottantes cassent les démarrages à froid.
• Utilisez python_requirements plutôt que python_packages en ligne une fois la liste plus longue.
• cuda suit votre wheel torch (par ex. 12.8 associé à torch==2.7.1+cu128).
• Ajoutez train: train.py:train si votre modèle est fine-tunable.
• Ajoutez image: r8.im/owner/name pour activer le simple cog push.

Pour les predictors asynchrones avec batching continu:

concurrency:
  max: 32

Essentiels de predict.py

from cog import BasePredictor, Input, Path

class Predictor(BasePredictor):
    def setup(self) -> None:
        """Chargements ponctuels. Le travail lourd va ici, pas dans 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:
        """Exécuter une seule prédiction."""
        if not prompt.strip():
            raise ValueError("prompt cannot be empty")
        out = self.model.generate(prompt, seed=seed, steps=num_steps)
        return Path(out)

Règles des inputs:

• Chaque input a besoin d'une description. La description s'affiche dans le schéma du modèle et sur l'UI web de Replicate.
• Utilisez ge/le pour les limites numériques, choices=[...] pour les énums, regex= pour les strings.
• Utilisez cog.Path pour les inputs et outputs de fichiers, jamais de bytes bruts.
• Utilisez cog.Secret pour tout input de type token (tokens HF, clés API), jamais plain str.
• Fournissez une valeur par défaut qui est dans choices pour les inputs catégoriques.
• Validez les inputs tôt dans predict() et levez ValueError.

Sortie texte en streaming (pour les LLMs):

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 asynchrone avec batching continu (associé à concurrency.max dans cog.yaml):

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 dynamique à partir d'assets sur disque (par ex. un répertoire voices/ d'échantillons audio):

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

Charger les poids rapidement

Le démarrage à froid domine la latence perçue par l'utilisateur. Trois motifs, classés par simplicité:

1. Cuire les poids dans l'image au moment de la construction

Meilleur pour les poids petits ou moyens (< 5 Go) pour lesquels vous voulez zéro démarrage à froid.

Pour torchvision:

import os
os.environ["TORCH_HOME"] = "."  # définir avant d'importer torch
import torch
from torchvision import models

Pour HuggingFace:

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

Puis téléchargez une fois pendant cog build (par ex. dans une étape run: ou en exécutant un petit script de récupération dans la construction). Les poids deviennent partie d'une couche d'image.

2. Tirer de weights.replicate.delivery avec pget

Meilleur pour les poids volumineux, ou quand vous voulez partager des poids sur plusieurs modèles. pget est le fetcher HTTP parallèle de Replicate.

Dans 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

Dans 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 extrait tar en mémoire; concurrence par défaut est 4 * NumCPU
            subprocess.check_call(["pget", "-x", WEIGHTS_URL, str(WEIGHTS_DIR)])
        self.model = load_from(WEIGHTS_DIR)

Pour plusieurs fichiers en une seule fois:

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. HuggingFace Hub avec hf_transfer

Définissez HF_HUB_ENABLE_HF_TRANSFER=1 et utilisez huggingface_hub.snapshot_download ou from_pretrained. Plus rapide que les téléchargements HF vanille. Utilisez un input cog.Secret pour les modèles gated.

Cache de poids pour poids fournis par l'utilisateur

Pour les LoRAs ou toute URL de poids que l'utilisateur passe au moment de la prédiction, utilisez un cache disque keyed sha256 avec éviction 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()

Voir replicate/cog-flux/weights.py pour une version production qui gère HF, CivitAI, Replicate, et des URLs .safetensors arbitraires.

Composition multi-LoRA

Rechargez seulement quand l'URL change; composez deux LoRAs avec des scales séparées:

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"))

Astuces de démarrage à froid

À partir de modèles de diffusion en production comme replicate/cog-flux et replicate/cog-flux-kontext:

• Définissez les flags perf une fois dans setup():

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

• Compilez et échauffez:

  self.model = torch.compile(self.model, dynamic=True)
  _ = self.predict(prompt="warmup", num_steps=1)  # absorbe le coût de compilation dans setup

• Chargez les gros poids avec meta device + assign=True pour éviter la double allocation:

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

• Partagez VAE / text encoder sur plusieurs pipelines (par ex. base + img2img + inpaint) au lieu de charger trois copies.
• Pour fp8/int8, sauvegardez les poids quantizés à l'avance et chargez directement; ne quantizez pas au démarrage.

Développement local

cog init                                    # générer cog.yaml + predict.py
cog predict -i prompt="hello"               # construire + exécuter une seule prédiction
cog predict -i image=@input.jpg -o out.png  # inputs et outputs de fichiers
cog serve -p 8393                           # serveur HTTP correspondant à la production
cog exec python                             # shell interactif dans l'environnement de construction

Construction

cog build -t my-model
cog build --separate-weights -t my-model    # poids dans leur propre couche d'image
cog build --secret id=hf,src=$HOME/.hf_token -t my-model

Astuces:

• Utilisez --separate-weights pour tout modèle avec des poids > ~1 Go. Cela accélère les démarrages à froid et les pushes de registre.
• Utilisez --mount=type=cache,target=/root/.cache/pip dans les étapes run: pour cacher pip entre les constructions.
• Utilisez --secret au lieu de ARG pour garder les tokens hors de l'historique d'image.
• L'image de base Cog par défaut (--use-cog-base-image=true) est plus rapide que d'en créer une vôtre.

Entraînement

Si votre modèle supporte le fine-tuning, ajoutez train: train.py:train à cog.yaml et écrivez une fonction train() qui retourne TrainingOutput(weights=Path("model.tar")). Le predictor accepte alors l'URL via setup(self, weights) ou la variable d'env COG_WEIGHTS. Voir https://cog.run/training et replicate/flux-fine-tuner pour un exemple complet.

Directives

• Conservez setup() pour les chargements ponctuels; gardez predict() rapide et déterministe en forme.
• Épinglez Python et chaque dépendance. Utilisez numpy<2 si votre torch est plus ancien.
• Décrivez toujours chaque input. Les schémas sans descriptions sont inutilisables sur l'UI web.
• Utilisez cog.Path pour les fichiers et cog.Secret pour les tokens.
• Épinglez pget à une version spécifique (v0.8.2) pour la reproductibilité.
• Définissez HF_HUB_ENABLE_HF_TRANSFER=1 chaque fois que vous appelez HuggingFace Hub.
• Définissez TRANSFORMERS_OFFLINE=1 après le chargement des poids pour prévenir les lookups HF au runtime.
• Testez avec cog predict avant de pousser. Si ça ne marche pas localement, ça ne marchera pas en production.

Références production

• https://github.com/replicate/cog-examples — motifs minimaux (resnet, hello-world, streaming, training)
• https://github.com/replicate/cog-template — scaffolder pour les nouveaux dépôts de modèles
• https://github.com/replicate/cog-flux — modèles FLUX multi-variante, cache de poids, fp8 + torch.compile
• https://github.com/replicate/cog-flux-kontext — chargement avec meta device, compilation d'échauffement
• https://github.com/replicate/cog-vllm — serveur LLM asynchrone avec batching continu, training-as-packaging
• https://github.com/replicate/cog-comfyui — workflows ComfyUI en tant que modèle Cog, assistants custom-node
• https://github.com/replicate/flux-fine-tuner — composition multi-LoRA, composants pipeline partagés
• https://github.com/replicate/vibevoice — TTS avec choices dynamique, cog.yaml minimal
• https://github.com/replicate/pget — fetcher de poids parallèle