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

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, 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 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
RateLimiter	Async throttling helper for API quotas	utils/rate_limiter.py
YAML Config	Declares class_path and init_args for each generator	configs/script2video.yaml
Output Interfaces	Data containers (ImageOutput, VideoOutput) returned by generators	interfaces/image_output.py, 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:

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

And provide this async method:

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:

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

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:

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), add a class_path pointing to your module and class, plus init_args for instantiation:

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

# 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

# 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

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

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}