ElasticModels are the models produced by TheStage AI ANNA: Automated Neural Networks Accelerator. ANNA allows you to control model size, latency and quality with a simple slider movement, routing different compression algorithms to different layers. For each model, we have produced a series of optimized models:
S: The fastest model, with accuracy degradation less than 2%.
Models can be accessed via TheStage AI Python SDK: ElasticModels.
Installation
System Requirements
Property
Value
GPU
H100, B200
Python Version
3.10-3.12
CPU
Intel/AMD x86_64
CUDA Version
12.8+
TheStage AI Access token setup
Install TheStage AI CLI and setup API token:
bash
1pip install thestage
2thestage config set --access-token <YOUR_ACCESS_TOKEN>
Elastic Models provides the same interface as HuggingFace Diffusers. To infer our models, just replace diffusers import with elastic_models.diffusers. Here is an example of how to use the Wan2.2-T2V-A14B model:
python
1import torch
2from elastic_models.diffusers import WanPipeline
3from diffusers.utils import export_to_video
45model_name ='Wan-AI/Wan2.2-T2V-A14B-Diffusers'6device = torch.device("cuda")7dtype = torch.bfloat16
89pipe = WanPipeline.from_pretrained(10 model_name,11 torch_dtype=dtype,12# 'original' for original model13# 'S' for accelerated model14 mode='S'15)16pipe.vae.enable_tiling()17pipe.vae.enable_slicing()18pipe.to(device)1920prompt ="A beautiful woman in a red dress dancing"2122with torch.no_grad():23 output = pipe(24 prompt=prompt,25 negative_prompt="",26 height=480,27 width=480,28 num_frames=81,29 num_inference_steps=40,30 guidance_scale=3.0,31 guidance_scale_2=2.0,32 generator=torch.Generator("cuda").manual_seed(42),33)3435 video = output.frames[0]36 export_to_video(video,"output.mp4", fps=16)
Compiled versions are currently available only for 81-frame generations at 480x480 resolution. Other versions are not yet accessible. Stay tuned for updates!
Quality Benchmarks
We have used VBench to evaluate the quality of videos generated by different sizes of Wan2.2-T2V-A14B models compared to the original model. The evaluation metrics cover subject consistency, background consistency, motion smoothness, dynamic degree, aesthetic quality, and imaging quality.
Quality Benchmarking
Quality Benchmark Results
Metric/Model Size
S
M
L
XL
Original
Subject Consistency
0.96
N/A
N/A
N/A
0.96
Background Consistency
0.96
N/A
N/A
N/A
0.96
Motion Smoothness
0.98
N/A
N/A
N/A
0.98
Dynamic Degree
0.29
N/A
N/A
N/A
0.29
Aesthetic Quality
0.62
N/A
N/A
N/A
0.62
Imaging Quality
0.68
N/A
N/A
N/A
0.68
Dataset
VBench: A comprehensive benchmark suite for video generative models. It evaluates videos across multiple dimensions of quality including temporal consistency, motion naturalness, aesthetic appeal, and prompt faithfulness. See VBench GitHub for details.
Metrics
Subject Consistency: Measures how consistently the main subject appears across frames.
Background Consistency: Measures temporal consistency of the background across frames.
Motion Smoothness: Evaluates the naturalness and smoothness of motion between frames.
Dynamic Degree: Measures the amount of motion/dynamics present in the generated video.
Aesthetic Quality: Assesses the overall visual appeal and aesthetic of generated frames.
Imaging Quality: Evaluates the per-frame visual quality and sharpness of generated content.
Latency Benchmarks
We have measured the latency of different sizes of Wan2.2-T2V-A14B model (S, original) on various GPUs. The measurements were taken for generating videos at 480x480 resolution with 81 frames.
Latency Benchmarking
Latency Benchmark Results
Latency (in seconds) for generating a 480x480 video with 81 frames on various hardware setups.
GPU/Model Size
S
M
L
XL
Original
H100
90
N/A
N/A
N/A
180
B200
99
N/A
N/A
N/A
117
Benchmarking Methodology
The benchmarking was performed on a single GPU with a batch size of 1. Each model was run for 5 iterations, and the average latency was calculated.
Algorithm summary:
Load the Wan2.2-T2V-A14B model with the specified size (S, original).
Move the model to the GPU.
Prepare a sample prompt for video generation.
Run the model for a number of iterations (e.g., 5) and measure the time taken for each iteration. On each iteration:
Synchronize the GPU to flush any previous operations.
Record the start time.
Generate the video using the model.
Synchronize the GPU again.
Record the end time and calculate the latency for that iteration.
Calculate the average latency over all iterations.