# How to Load and Use Pretrained RWKV-CLIP Weights for Production Inference

> Learn to load and use pretrained RWKV-CLIP weights for production inference. Set env vars, create model skeleton, and run efficient vision text embeddings.

- Repository: [DeepGlint/rwkv-clip](https://github.com/deepglint/rwkv-clip)
- Tags: how-to-guide
- Published: 2026-02-28

---

**Load pretrained RWKV-CLIP weights by setting required environment variables, initializing the model skeleton with `create_RWKV_Model`, and running inference through wrapper functions that handle vision and text embeddings with proper pooling.**

The deepglint/rwkv-clip repository implements a multimodal architecture that combines a **Vision-RWKV** backbone with a **Text-RWKV** encoder to mirror the CLIP framework. Unlike monolithic training artifacts, this implementation separates model definitions from learned parameters, requiring specific initialization steps to properly load and use pretrained RWKV-CLIP weights for inference in production environments.

## Architecture Overview

RWKV-CLIP replaces the traditional Transformer blocks found in standard CLIP with RWKV (Receptance Weighted Key Value) layers for both vision and language processing. The vision encoder processes image patches through spatial mixing operations, while the text encoder handles token sequences with adaptive pooling.

Two critical wrapper functions bridge the gap between raw RWKV outputs and CLIP-compatible embeddings:

- **`WarperCLIP_V_T_RWKV_method`** – Handles the vision pathway, managing conditional CLS token logic and applying **GlobalAveragePooling** to produce a single image embedding.
- **`WarperCLIP_V_T_RWKV_text_change_head`** – Processes text hidden states through normalization and `nn.AdaptiveAvgPool1d(1)` to match the original CLIP text encoder behavior.

## Three-Step Weight Loading Process

Loading pretrained weights requires precise orchestration of environment configuration, model instantiation, and state dict cleaning.

### 1. Configure Environment Variables

The vision RWKV blocks require specific CUDA kernel parameters set before model import. Define these in your shell or Python environment:

```bash
export Image_T_max=256
export Image_HEAD_SIE=64

```

These values must match the checkpoint's training configuration. `Image_T_max` controls the maximum token length for image patches, while `Image_HEAD_SIE` defines the head size dimension.

### 2. Initialize the Model Skeleton

Instantiate the model architecture using `create_RWKV_Model` from [`model/utils.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/utils.py) (line 14). This function constructs both encoders based on hyperparameters passed via an args object:

```python
from model import create_RWKV_Model

model = create_RWKV_Model(
    args, 
    model_weight_path="/path/to/rwkv_clip.pth"
)

```

The function automatically calls `load_model_weight` (defined at line 95 in [`model/utils.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/utils.py)), which strips `module.` and `_orig_mod.` prefixes introduced by `torch.nn.DataParallel` or `DistributedDataParallel` during distributed training.

### 3. Prepare for Production Inference

Move the model to GPU and enable evaluation mode to disable dropout and activate cuDNN optimizations:

```python
model.cuda()
model.eval()

```

The model now accepts batched inputs through its forward method, which internally routes images through `WarperCLIP_V_T_RWKV_method` and text through `WarperCLIP_V_T_RWKV_text_change_head` (lines 46 and 66 in [`model/utils.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/utils.py)).

## Production Inference Implementation

### Minimal PyTorch Script

This self-contained snippet demonstrates the complete inference pipeline:

```python
import os
import torch
from model import create_RWKV_Model, WarperCLIP_V_T_RWKV_method, WarperCLIP_V_T_RWKV_text_change_head
from model.open_clip.tokenizer import tokenize
from model.open_clip.transform import image_transform
from PIL import Image

# Required environment setup

os.environ["Image_T_max"] = "256"
os.environ["Image_HEAD_SIE"] = "64"

# Define configuration stub (match your checkpoint's training args)

class Args:
    input_size = 224
    image_patch_size = 16
    image_embed_dims = 384
    image_hidden_rate = 4
    image_depth = 12
    image_num_heads = 6
    image_output_cls_token = False
    image_with_cls_token = False
    data_type = "utf-8"
    ctx_len = 77
    vocab_size = 49408
    text_initialization = True
    head_size = 64
    text_num_head = 0
    head_size_divisor = 8
    n_layer = 12
    n_embd = 384
    dim_att = 0
    dim_ffn = 0
    pre_ffn = 0
    pos_emb = 0
    head_qk = 0
    tiny_att_dim = 0
    tiny_att_layer = -999

args = Args()

# Build and load model

model = create_RWKV_Model(args, model_weight_path="/path/to/rwkv_clip.pth")
model.eval().cuda()

# Preprocess inputs

transform = image_transform(args.input_size, False)
raw_image = Image.open("example.jpg").convert("RGB")
image_tensor = transform(raw_image).unsqueeze(0).cuda()

texts = ["a photo of a cat."]
token_ids = tokenize(texts).cuda()

# Forward pass returns (image_emb, text_emb, logit_scale)

image_emb, text_emb, logit_scale = model(image_tensor, token_ids)

# Compute cosine similarity

image_norm = torch.nn.functional.normalize(image_emb, dim=-1)
text_norm = torch.nn.functional.normalize(text_emb, dim=-1)
similarity = logit_scale * image_norm @ text_norm.t()

```

### Zero-Shot Classification Pipeline

For batch classification across multiple categories, use the `zero_shot_classifier` utility from [`zero_shot.py`](https://github.com/deepglint/rwkv-clip/blob/main/zero_shot.py) (line 59):

```python
import json
from zero_shot import zero_shot_classifier

# Load class definitions and templates

with open('utils/label.json') as f:
    class_labels = json.load(f)['imagenet']
with open('utils/template.json') as f:
    templates = json.load(f)['imagenet']

# Pre-compute classifier weights (run once)

classifier_weights = zero_shot_classifier(model, class_labels, templates, args)
classifier_weights = classifier_weights.cuda()

# Inference per image

image_emb, _, _ = model(image_tensor, token_ids)  # token_ids can be dummy here

image_emb = torch.nn.functional.normalize(image_emb, dim=-1)
logits = image_emb @ classifier_weights
predicted_class = class_labels[logits.argmax().item()]

```

### FastAPI Deployment Example

Deploy as a scalable REST endpoint:

```python
from fastapi import FastAPI, File, UploadFile
from PIL import Image
import io, torch, os

app = FastAPI()

# Initialize once at startup

os.environ["Image_T_max"] = "256"
os.environ["Image_HEAD_SIE"] = "64"

model = create_RWKV_Model(args, model_weight_path="/path/to/rwkv_clip.pth")
model.eval().cuda()
transform = image_transform(args.input_size, False)

@app.post("/similarity")
async def similarity(image: UploadFile = File(...), text: str = "a photo of a cat."):
    # Image preprocessing

    contents = await image.read()
    pil_img = Image.open(io.BytesIO(contents)).convert("RGB")
    img_tensor = transform(pil_img).unsqueeze(0).cuda()
    
    # Text tokenization

    tokens = tokenize([text]).cuda()
    
    # Forward and normalize

    img_emb, txt_emb, scale = model(img_tensor, tokens)
    img_emb = torch.nn.functional.normalize(img_emb, dim=-1)
    txt_emb = torch.nn.functional.normalize(txt_emb, dim=-1)
    
    score = (scale * img_emb @ txt_emb.t()).item()
    return {"similarity_score": score}

```

## Production Optimization Guidelines

Optimize your deployment with these specific configurations:

- **GPU Placement** – Call `model.cuda()` once after loading weights. Use `model.eval()` to ensure deterministic batch normalization and disable training-specific layers.
- **Precision Modes** – The repository supports `fp16`, `bf16`, and `fp32`. Cast tensors using `.half()` or `.bfloat16()` for reduced memory bandwidth, or specify `--precision` flags if wrapping the training scripts.
- **Batch Processing** – The vision encoder is fully convolutional, allowing variable batch sizes up to GPU memory limits. Process multiple images per call to maximize throughput.
- **Determinism** – Reuse the `setup_seed` helper from the training utilities if reproducible outputs are required, though this is optional for standard inference.
- **Checkpoint Compatibility** – The `load_model_weight` function automatically handles checkpoints saved under `DataParallel` wrappers by removing prefix strings, ensuring compatibility regardless of training distribution strategy.

## Summary

- **Set environment variables** `Image_T_max` and `Image_HEAD_SIE` before importing the model to configure vision block parameters.
- **Initialize with `create_RWKV_Model`** from [`model/utils.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/utils.py), passing your configuration args and checkpoint path to handle both skeleton creation and weight loading.
- **Enable production mode** by calling `model.cuda()` and `model.eval()` immediately after initialization.
- **Process inputs through the model's forward method**, which automatically routes vision data through `WarperCLIP_V_T_RWKV_method` and text through `WarperCLIP_V_T_RWKV_text_change_head`.
- **Normalize embeddings** and apply the returned `logit_scale` to compute CLIP-style similarity scores, or use `zero_shot_classifier` for efficient batch classification.

## Frequently Asked Questions

### Where are the pretrained RWKV-CLIP weights stored?

The deepglint/rwkv-clip repository contains only model architecture definitions. Pretrained learned weights are distributed as separate checkpoint files (e.g., `rwkv_clip.pth`) that you download independently and load via the `model_weight_path` argument in `create_RWKV_Model`.

### Why must I set `Image_T_max` and `Image_HEAD_SIE` environment variables?

These variables configure the CUDA kernels and head size dimensions for the Vision-RWKV blocks at import time. The vision encoder in [`model/Image_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/Image_rwkv.py) reads these values during module initialization to allocate appropriate tensor shapes for spatial mixing operations.

### How does RWKV-CLIP handle checkpoints saved with DataParallel?

The `load_model_weight` function in [`model/utils.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/utils.py) (line 95) automatically cleans state dictionaries by stripping `module.` and `_orig_mod.` prefixes. This ensures checkpoints saved during distributed training with `torch.nn.DataParallel` or `DistributedDataParallel` load seamlessly into single-GPU or non-distributed multi-GPU inference environments.

### Can I run RWKV-CLIP inference in mixed precision?

Yes. The model architecture supports floating-point formats including `fp16`, `bf16`, and `fp32`. You can manually cast the model and inputs using `.half()` or `.bfloat16()`, or specify precision modes through the `--precision` command-line flag when adapting the provided training scripts for inference services.