Problem
The README currently documents single-GPU inference with offload mode enabled, specifically saying that VAE and text_encoder are offloaded to CPU for the RoboTwin and i2av paths:
- RoboTwin: approximately 24GB VRAM with offload mode enabled, VAE and
text_encoder offloaded to CPU.
- i2av: approximately 18GB VRAM with offload mode enabled, VAE and
text_encoder offloaded to CPU.
In our tests, following this broad offload behavior makes inference latency unacceptable for deployment/control use. VAE CPU offload is the critical blocker: the model initializes, but a single inference chunk does not complete within a 240s timeout.
Environment
- GPU: RTX 4090 24GB
- PyTorch: 2.9.0+cu128
- CUDA runtime: 12.8
- Attention backend:
flashattn
- Input size: 256x256
frame_chunk_size: 4
num_inference_steps: 5
action_num_inference_steps: 10
Observed results
Case 1: VAE + text_encoder offloaded to CPU
Configuration:
enable_offload=true
enable_vae_offload=true
enable_text_encoder_offload=true
Timing:
model_init: 2.462s
reset: 86.625s
one_chunk_inference: started but did not finish within 240s
This is not usable for real-time or closed-loop robot control.
Case 2: keep VAE on GPU, offload only text_encoder
Configuration:
enable_offload=true
enable_vae_offload=false
enable_text_encoder_offload=true
Timing:
model_load: 2.421s
reset: 89.740s
one_chunk_inference: 1.623s
Memory peak during one chunk:
max_allocated: 20210.8 MiB
max_reserved: 20292.0 MiB
A 32-chunk wrapper run with VAE on GPU completed 32/32 requests successfully. After the first reset-heavy request, subsequent chunks were about 1.4-2.0s each.
Expected behavior / request
Please fix readme instruction so that For 24GB GPUs, keeping the VAE resident on GPU appears to be necessary for practical inference.
Related: #41 also points out that the original code offloads VAE to CPU.
Problem
The README currently documents single-GPU inference with offload mode enabled, specifically saying that VAE and
text_encoderare offloaded to CPU for the RoboTwin and i2av paths:text_encoderoffloaded to CPU.text_encoderoffloaded to CPU.In our tests, following this broad offload behavior makes inference latency unacceptable for deployment/control use. VAE CPU offload is the critical blocker: the model initializes, but a single inference chunk does not complete within a 240s timeout.
Environment
flashattnframe_chunk_size: 4num_inference_steps: 5action_num_inference_steps: 10Observed results
Case 1: VAE + text_encoder offloaded to CPU
Configuration:
Timing:
This is not usable for real-time or closed-loop robot control.
Case 2: keep VAE on GPU, offload only text_encoder
Configuration:
Timing:
Memory peak during one chunk:
A 32-chunk wrapper run with VAE on GPU completed 32/32 requests successfully. After the first reset-heavy request, subsequent chunks were about 1.4-2.0s each.
Expected behavior / request
Please fix readme instruction so that For 24GB GPUs, keeping the VAE resident on GPU appears to be necessary for practical inference.
Related: #41 also points out that the original code offloads VAE to CPU.