# How to Integrate Custom Image or Video Generation Models into ViMax's RenderBackend

> Integrate custom models into ViMax RenderBackend easily. Implement a generator class, expose it via Python, and configure in YAML. No core pipeline changes needed.

- Repository: [✨Data Intelligence Lab@HKU✨/ViMax](https://github.com/HKUDS/ViMax)
- Tags: how-to-guide
- Published: 2026-05-20

---

**To integrate custom models into ViMax's RenderBackend, implement a generator class with the required async interface, expose it via a Python module path, and reference that path in the YAML configuration file—no core pipeline modifications needed.**

The ViMax framework (HKUDS/ViMax) uses a flexible, config-driven factory system that lets you swap image and video generation models without touching the core pipeline code. By following the interface contracts defined in [`tools/render_backend.py`](https://github.com/HKUDS/ViMax/blob/main/tools/render_backend.py), you can integrate custom generation models—from proprietary APIs to local diffusion models—using only a Python class and a YAML configuration update.

## Understanding the RenderBackend Architecture

ViMax's rendering pipeline centers on a **config-driven factory** that instantiates concrete generator objects at runtime. The `RenderBackend.from_config` method in [`tools/render_backend.py`](https://github.com/HKUDS/ViMax/blob/main/tools/render_backend.py) dynamically loads classes based on YAML specifications, handling module imports and optional rate limiting automatically.

### Core Components

| Component | Role | Source File |
|-----------|------|-------------|
| **RenderBackend** | Factory that bundles `image_generator` and `video_generator` instances | [`tools/render_backend.py`](https://github.com/HKUDS/ViMax/blob/main/tools/render_backend.py) |
| **RateLimiter** | Async throttling helper for API quotas | [`utils/rate_limiter.py`](https://github.com/HKUDS/ViMax/blob/main/utils/rate_limiter.py) |
| **YAML Config** | Declares `class_path` and `init_args` for each generator | [`configs/script2video.yaml`](https://github.com/HKUDS/ViMax/blob/main/configs/script2video.yaml) |
| **Output Interfaces** | Data containers (`ImageOutput`, `VideoOutput`) returned by generators | [`interfaces/image_output.py`](https://github.com/HKUDS/ViMax/blob/main/interfaces/image_output.py), [`interfaces/video_output.py`](https://github.com/HKUDS/ViMax/blob/main/interfaces/video_output.py) |

When `RenderBackend.from_config` executes, it parses the configuration dictionary, constructs a `RateLimiter` if limits are specified, and calls `_instantiate` to dynamically import the class from `class_path` and inject dependencies.

## Required Interface for Custom Generators

To integrate custom image or video generation models into ViMax's RenderBackend, your class must expose specific async methods and accept a standard constructor signature.

### Image Generator Contract

Your custom image generator must implement:

```python
def __init__(self, *, api_key: str, rate_limiter: Optional[RateLimiter] = None, **kwargs)

```

And provide this async method:

```python
async def generate_single_image(
    self,
    prompt: str,
    reference_image_paths: List[str] = [],
    aspect_ratio: Optional[str] = "16:9",
    **kwargs,
) -> ImageOutput

```

### Video Generator Contract

Similarly, video generators require:

```python
async def generate_single_video(
    self,
    prompt: str,
    reference_video_paths: List[str] = [],
    aspect_ratio: Optional[str] = "16:9",
    **kwargs,
) -> VideoOutput

```

Both must return instances of the respective output data classes defined in [`interfaces/image_output.py`](https://github.com/HKUDS/ViMax/blob/main/interfaces/image_output.py) and [`interfaces/video_output.py`](https://github.com/HKUDS/ViMax/blob/main/interfaces/video_output.py).

## Step-by-Step Integration Guide

Follow these steps to wire your custom model into the ViMax pipeline.

### 1. Create the Generator Class

Place your implementation in a module accessible from the repository root (e.g., [`tools/my_custom_generator.py`](https://github.com/HKUDS/ViMax/blob/main/tools/my_custom_generator.py)). The class must be importable via a dotted path.

### 2. Implement the Constructor and Rate Limiter Check

Accept `api_key` and an optional `rate_limiter` parameter. Store the rate limiter for use in generation methods:

```python
from typing import Optional
from utils.rate_limiter import RateLimiter

class ImageGeneratorCustom:
    def __init__(self, *, api_key: str, model_endpoint: str, rate_limiter: Optional[RateLimiter] = None):
        self.api_key = api_key
        self.model_endpoint = model_endpoint
        self.rate_limiter = rate_limiter

```

### 3. Add the Async Generation Method

Before calling external APIs, check and acquire the rate limiter if present:

```python
async def generate_single_image(self, prompt: str, **kwargs):
    if self.rate_limiter:
        await self.rate_limiter.acquire()
    # Your generation logic here

    return ImageOutput(fmt="pil", ext="png", data=image)

```

### 4. Update the YAML Configuration

In your configuration file (e.g., [`configs/script2video.yaml`](https://github.com/HKUDS/ViMax/blob/main/configs/script2video.yaml)), add a `class_path` pointing to your module and class, plus `init_args` for instantiation:

```yaml
image_generator:
  class_path: tools.my_custom_generator.ImageGeneratorCustom
  init_args:
    api_key: ${CUSTOM_API_KEY}
    model_endpoint: "https://api.custom-model.com/v1/generate"
  max_requests_per_minute: 10
  max_requests_per_day: 1000

```

The `max_requests_per_minute` and `max_requests_per_day` fields are optional. When present, the factory automatically constructs a `RateLimiter` and injects it into your class constructor.

### 5. Run the Pipeline

Existing pipeline scripts (such as [`main_script2video.py`](https://github.com/HKUDS/ViMax/blob/main/main_script2video.py)) already invoke `RenderBackend.from_config`. They will now instantiate your custom generator without further code changes.

## Complete Implementation Example

This example demonstrates a complete custom image generator integration.

### Custom Generator Implementation

```python

# file: tools/image_generator_custom_api.py

import logging
from typing import List, Optional
from PIL import Image
from io import BytesIO
import aiohttp

from interfaces.image_output import ImageOutput
from utils.rate_limiter import RateLimiter

class ImageGeneratorCustomAPI:
    def __init__(
        self,
        *,
        api_key: str,
        model_version: str = "v1",
        rate_limiter: Optional[RateLimiter] = None,
    ):
        self.api_key = api_key
        self.model_version = model_version
        self.rate_limiter = rate_limiter
        self.endpoint = f"https://api.custom-gen.com/{model_version}/image"

    async def generate_single_image(
        self,
        prompt: str,
        reference_image_paths: List[str] = [],
        aspect_ratio: Optional[str] = "16:9",
        **kwargs,
    ) -> ImageOutput:
        # Apply rate limiting if configured

        if self.rate_limiter:
            await self.rate_limiter.acquire()
        
        # Call external API (simplified)

        async with aiohttp.ClientSession() as session:
            payload = {
                "prompt": prompt,
                "aspect_ratio": aspect_ratio,
                "reference_images": reference_image_paths,
            }
            headers = {"Authorization": f"Bearer {self.api_key}"}
            
            async with session.post(self.endpoint, json=payload, headers=headers) as resp:
                response_data = await resp.read()
                image = Image.open(BytesIO(response_data))
                
        logging.info(f"Generated image via CustomAPI model {self.model_version}")
        return ImageOutput(fmt="pil", ext="png", data=image)

```

### YAML Configuration

```yaml

# file: configs/custom_generation.yaml

image_generator:
  class_path: tools.image_generator_custom_api.ImageGeneratorCustomAPI
  init_args:
    api_key: ${CUSTOM_API_KEY_ENV_VAR}
    model_version: "v2-beta"
  max_requests_per_minute: 5
  max_requests_per_day: 500

video_generator:
  class_path: tools.video_generator_veo_google_api.VideoGeneratorVeo
  init_args:
    api_key: ${GOOGLE_API_KEY}

```

### Pipeline Usage

```python
from utils.provider_presets import load_yaml
from tools.render_backend import RenderBackend

# Load configuration

config = load_yaml("configs/custom_generation.yaml")

# Factory instantiates your custom class automatically

backend = RenderBackend.from_config(config)

# Use in pipeline

image_output = await backend.image_generator.generate_single_image(
    prompt="A futuristic cityscape with neon lights",
    aspect_ratio="16:9"
)

```

## Summary

- **Implement the contract**: Create a class in `tools/` with `__init__(..., rate_limiter=None)` and async `generate_single_image` or `generate_single_video` methods returning `ImageOutput` or `VideoOutput`.
- **Configure via YAML**: Set `class_path` to your module's dotted path and provide `init_args` for construction.
- **Enable rate limiting**: Add `max_requests_per_minute` or `max_requests_per_day` to the config section to automatically receive a `RateLimiter` instance.
- **Zero pipeline changes**: The `RenderBackend.from_config` factory in [`tools/render_backend.py`](https://github.com/HKUDS/ViMax/blob/main/tools/render_backend.py) handles instantiation via `importlib`, allowing hot-swapping of models without modifying core logic.

## Frequently Asked Questions

### What methods must my custom generator implement?

Your custom image generator must implement an async `generate_single_image` method accepting `prompt`, `reference_image_paths`, `aspect_ratio`, and returning `ImageOutput`. For video, implement `generate_single_video` returning `VideoOutput`. Both must accept `api_key` and an optional `rate_limiter` in their constructors.

### Where should I place my custom generator class files?

Place them in any module importable from the repository root, such as [`tools/my_generator.py`](https://github.com/HKUDS/ViMax/blob/main/tools/my_generator.py). Ensure the `class_path` in your YAML config matches the Python dotted path (e.g., `tools.my_generator.MyImageGen`).

### How does ViMax handle rate limiting for custom APIs?

When you specify `max_requests_per_minute` or `max_requests_per_day` in the YAML config, the `RenderBackend` constructs a `RateLimiter` from [`utils/rate_limiter.py`](https://github.com/HKUDS/ViMax/blob/main/utils/rate_limiter.py) and injects it into your generator's constructor. Your implementation should call `await self.rate_limiter.acquire()` before making API requests.

### Can I use the same custom model for both image and video generation?

No, ViMax treats these as separate concerns requiring distinct interfaces. If your underlying service supports both modalities, create separate classes implementing `generate_single_image` and `generate_single_video`, then reference both in the YAML config under `image_generator` and `video_generator` sections respectively.