Skip to main content

Adding a New Model Integration

This is the checklist we follow whenever a new model/provider (cloud API or open-weight/local) is added to OpenTryOn. Use it as a PR checklist -- most integrations touch every item below.

Before this checklist: read the agent-facing decision guide Model Integration Guidelines (Path A first-party API vs Path B local/HF/Ollama/LM Studio/Unsloth, dual-path rules, and adapter contracts).

1. Decide where the adapter lives​

KindLocationPattern
Cloud API, existing vendor with a dedicated packagetryon/api/<provider>/tryon/api/kimi/adapter.py, or a new class in an existing family file, e.g. NanoBanana2LiteAdapter added to tryon/api/nano_banana/adapter.py
Cloud API, brand-new vendor whose adapter's primary purpose is a single existing servicetryon/api/<service>/ (use-case directory, matches a SERVICES key in tryon/cli/registry.py)historically tryon/api/vton/p_image_tryon.py (Pruna try-on); once a vendor grows multiple adapters, promote to tryon/api/<provider>/ as we did for Pruna (tryon/api/pruna/)
Open-weight / local GPU modeltryon/models/<model>/tryon/models/kimi_vl/adapter.py

Why a use-case directory for new vendors: tryon/api/ already has a one-directory-per-vendor pattern for providers with multiple adapter classes or files (nano_banana/, kimi/, openAI/, lumaAI/, ben2/). As more VTON/generate/edit providers get added over time, minting a new top-level tryon/api/<vendor>/ directory for every single-file, single-purpose adapter would make that listing grow unbounded. Instead, park new single-file adapters whose primary purpose is one existing service inside the matching use-case directory (tryon/api/vton/, and analogously tryon/api/generate/, tryon/api/edit/, etc. if/when those are needed) -- the same small, fixed set of directories the CLI/MCP registry already organizes by (vton, generate, edit, understand, video-generate, bg-remove). A vendor whose model is already represented in one of these directories, or that ships multiple adapter classes/files, still gets its own tryon/api/<provider>/ package as before.

If a single adapter/model genuinely serves multiple services (e.g. an image model that can also do virtual try-on via multi-image composition), prefer registering the same class under multiple ModelSpec entries in tryon/cli/registry.py (different service, different method) over duplicating the adapter file into several use-case directories -- see how nano-banana/nano-banana-pro/nano-banana-2/nano-banana-2-lite are each registered under both generate and edit, and how nano-banana-2-lite is also registered under vton (calling generate_virtual_tryon(), a thin wrapper around generate_multi_image()) without adding any new files.

Single-file adapters (no extra helpers) can skip the subpackage and live directly as tryon/api/<provider>.py (see veo.py, segmind.py). Prefer a subpackage once you have more than one file (adapter + custom model code, multiple adapter classes, etc.).

2. Write the adapter class​

Follow the conventions already used across the repo so every adapter feels the same to callers and to the CLI:

  • Constructor: __init__(self, api_key: Optional[str] = None, ...). Fall back to an environment variable when api_key isn't passed, and raise ValueError if it's missing after that. Raise ImportError (with a pip install ... hint) if an optional SDK/library isn't installed -- see tryon/api/openAI/image_adapter.py for the pattern.
  • Flexible input types: methods that take an image/video should accept a file path, URL, PIL.Image.Image, raw bytes, or io.BytesIO (see _load_bytes/prepare_image helpers in existing adapters). Don't force callers to pre-load images themselves.
  • Method naming: match the existing vocabulary so the CLI/registry stay uniform -- generate_text_to_image / generate_image_edit / generate_multi_image for image generation, generate_and_decode for try-on, understand / understand_image / understand_video for understanding models, generate_text_to_video / generate_image_to_video for video.
  • Return types: prefer List[Image.Image] for image outputs (or raw bytes/List[bytes] if decoding needs to stay optional -- see output_kind below), raw bytes for video, and a dict (e.g. {"text": ..., "reasoning": ..., "model": ..., "usage": ...}) for understanding/captioning models.
  • Docstring header: include a one-paragraph summary, a link to the provider's API docs, the exact model ids supported, and 1-2 runnable examples (module-level docstring, mirrored by every existing adapter).

3. Export it from the subpackage __init__.py​

# tryon/api/<provider>/__init__.py
from .adapter import YourAdapter

__all__ = ["YourAdapter"]

4. Register the import so tryon.api / tryon.models stay lazy-safe​

  • Cloud adapters: add an entry to _LAZY_ATTRS in tryon/api/__init__.py (PEP 562 lazy loading). Never add a plain top-level from .provider import Adapter there -- that would force every adapter's transitive dependencies (including heavy ones like torch/timm for BEN2) to import just because someone touched tryon.api.

    _LAZY_ATTRS = {
    ...
    "YourAdapter": ".your_provider",
    }
  • Local models: add a normal eager import to tryon/models/__init__.py (this package already assumes torch/transformers are installed via the local extra, so there's no lazy-loading concern here).

5. Add dependencies to setup.py​

  • Lightweight SDK/client library (e.g. a REST client, the OpenAI SDK) → install_requires (core install).
  • Heavy ML dependency (torch, transformers, diffusers, a model-specific package) → LOCAL_INFERENCE_DEPS, which feeds the local / training / all extras. Keep pip install opentryon free of GPU dependencies.

6. Add environment variables​

  • Add the new variable(s) to env.template with a comment explaining where to get the key and which CLI models need it.
  • Document the same variable in the README's environment/setup section and in the adapter's own docs page (see step 9).

7. Wire it into the opentryon CLI​

Add a ModelSpec entry to the right service dict in tryon/cli/registry.py (vton, generate, edit, understand, video-generate, bg-remove, or a new service if it's a genuinely new category):

"your-model": ModelSpec(
id="your-model", label="Human-readable name",
import_path="tryon.api.your_provider", class_name="YourAdapter",
method="generate_text_to_image", # or whatever public method to call
output_kind="images", # "images" | "image_bytes" | "video_bytes" | "text"
env_hint="YOUR_API_KEY", # omit for local models
extra="local", # only for GPU/local models
args=[
Arg(("--prompt", "-p"), "prompt", required=True, help="Text prompt"),
...
],
),

Watch out for the one recurring gotcha: dest collisions with the service-level --model selector. If the adapter itself takes a model= kwarg (e.g. a provider-specific model version string), you cannot name the CLI arg's dest="model" -- that's already used by the top-level --model <provider> flag and argparse will silently let one clobber the other. Use call_name to decouple them:

Arg(("--model-version",), "model_version", call_name="model", default="...", choices=[...]),

8. Validate and test​

  • Run python3 -c "from tryon.cli.registry import validate_registry; validate_registry()" (or tests/test_cli.py) to catch reserved/duplicate flag or dest collisions before they ship.
  • Add a dry-run check to tests/test_cli.py (--dry-run prints the resolved Adapter(**init_kwargs).method(**call_kwargs) call without hitting the network/GPU).
  • If a cheap/free-tier API call is feasible, gate a real end-to-end test behind the relevant env var being set (see check_flux_vto_real_call / check_kimi_k26_real_call in tests/test_cli.py) so CI/local runs skip it gracefully when the key isn't configured.

9. Update documentation​

  • README.md: add the adapter to the relevant section (Virtual Try-On / Image Generation / Video Generation / Understanding / Background Removal), the CLI service/model table, and the environment variables list.
  • Docusaurus site (docs/docs/):
    • New page under api-reference/<name>.md (cloud) or local-models/<name>.md (local), following the structure of api-reference/flux2.md / local-models/flux2-turbo.md.
    • Add the new page to docs/sidebars.ts.
    • Add a short mention to api-reference/overview.md or local-models/overview.md.
    • Update docs/docs/intro.md's feature bullets/keywords if the model adds a new capability or domain.
  • CHANGELOG.md: add an entry under "Added" for the release.

10. Sanity-check the full picture​

python3 -c "from tryon.api import YourAdapter"     # lazy import still works
python3 -m tryon.cli.main <service> --model your-model --help
python3 -m tryon.cli.main <service> --model your-model ... --dry-run
python3 tests/test_cli.py

If all of the above pass and the docs are updated, the integration is done.