Elastic models 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. For each model, ANNA produces a series of optimized models:
-
XL: Mathematically equivalent neural network, optimized with our DNN compiler.
-
L: Near lossless model, with less than 1% degradation obtained on corresponding benchmarks.
-
M: Faster model, with accuracy degradation less than 1.5%.
-
S: The fastest model, with accuracy degradation less than 2%.
1import torch
2from elastic_models.diffusers import DiffusionPipeline
3from diffusers.video_processor import VideoProcessor
4from diffusers.utils import export_to_video
5
6mode_name = "genmo/mochi-1-preview"
7hf_token = ""
8device = torch.device("cuda")
9dtype = torch.bfloat16
10
11pipe = DiffusionPipeline.from_pretrained(
12 mode_name, torch_dtype=dtype, token=hf_token, mode="S"
13)
14pipe.enable_vae_tiling()
15pipe.to(device)
16
17prompt = "Kitten eating a banana"
18with torch.no_grad():
19 torch.cuda.synchronize()
20 (
21 prompt_embeds,
22 prompt_attention_mask,
23 negative_prompt_embeds,
24 negative_prompt_attention_mask,
25 ) = pipe.encode_prompt(prompt=prompt)
26 if prompt_attention_mask is not None and isinstance(
27 prompt_attention_mask, torch.Tensor
28 ):
29 prompt_attention_mask = prompt_attention_mask.to(dtype)
30
31 if negative_prompt_attention_mask is not None and isinstance(
32 negative_prompt_attention_mask, torch.Tensor
33 ):
34 negative_prompt_attention_mask = negative_prompt_attention_mask.to(dtype)
35
36 prompt_embeds = prompt_embeds.to(dtype)
37 negative_prompt_embeds = negative_prompt_embeds.to(dtype)
38
39 with torch.autocast("cuda", torch.bfloat16, enabled=True):
40 frames = pipe(
41 prompt_embeds=prompt_embeds,
42 prompt_attention_mask=prompt_attention_mask,
43 negative_prompt_embeds=negative_prompt_embeds,
44 negative_prompt_attention_mask=negative_prompt_attention_mask,
45 guidance_scale=4.5,
46 num_inference_steps=64,
47 height=480,
48 width=848,
49 num_frames=163,
50 generator=torch.Generator("cuda").manual_seed(0),
51 output_type="latent",
52 return_dict=False,
53 )[0]
54
55 video_processor = VideoProcessor(vae_scale_factor=8)
56 has_latents_mean = (
57 hasattr(pipe.vae.config, "latents_mean")
58 and pipe.vae.config.latents_mean is not None
59 )
60 has_latents_std = (
61 hasattr(pipe.vae.config, "latents_std")
62 and pipe.vae.config.latents_std is not None
63 )
64
65 if has_latents_mean and has_latents_std:
66 latents_mean = (
67 torch.tensor(pipe.vae.config.latents_mean)
68 .view(1, 12, 1, 1, 1)
69 .to(frames.device, frames.dtype)
70 )
71 latents_std = (
72 torch.tensor(pipe.vae.config.latents_std)
73 .view(1, 12, 1, 1, 1)
74 .to(frames.device, frames.dtype)
75 )
76 frames = frames * latents_std / pipe.vae.config.scaling_factor + latents_mean
77 else:
78 frames = frames / pipe.vae.config.scaling_factor
79
80 with torch.autocast("cuda", torch.bfloat16, enabled=False):
81 video = pipe.vae.decode(frames.to(pipe.vae.dtype), return_dict=False)[0]
82
83 video = video_processor.postprocess_video(video)[0]
84 torch.cuda.synchronize()
85 export_to_video(video, "mochi.mp4", fps=30)
1pip install thestage
2pip install 'thestage-elastic-models[nvidia]' --extra-index-url https://thestage.jfrog.io/artifactory/api/pypi/pypi-thestage-ai-production/simple
3
4# or for blackwell support
5pip install 'thestage-elastic-models[blackwell]' --extra-index-url https://thestage.jfrog.io/artifactory/api/pypi/pypi-thestage-ai-production/simple
6pip install -U --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128
7pip install -U --pre torchvision --index-url https://download.pytorch.org/whl/nightly/cu128
8
9pip install flash_attn==2.7.3 --no-build-isolation
10pip uninstall apex
11pip install tensorrt==10.11.0.33 opencv-python==4.11.0.86 imageio-ffmpeg==0.6.0
Then go to
app.thestage.ai, login and generate API token from your profile page. Set up API token as follows:
Benchmarking is one of the most important procedures during model acceleration. We aim to provide clear performance metrics for models using our algorithms.
Time in seconds of generation.