File size: 4,715 Bytes
7766b8e ba2a039 7766b8e 7d464d8 1011ea0 7766b8e da420ac 93900c5 7766b8e f7baf62 7766b8e f5db075 7766b8e ade2e09 7766b8e 21a480e 7766b8e 4143be0 b6454aa 4143be0 6aebbc5 7766b8e 7d93b81 7766b8e c1e6f35 4df7a38 765e122 50206fc 7766b8e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
---
license: apache-2.0
base_model:
- black-forest-labs/FLUX.1-dev
base_model_relation: quantized
pipeline_tag: text-to-image
---
# Elastic model: Fastest self-serving models. FLUX.1-dev.
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%.
__Goals of Elastic Models:__
* Provide the fastest models and service for self-hosting.
* Provide flexibility in cost vs quality selection for inference.
* Provide clear quality and latency benchmarks.
* Provide interface of HF libraries: transformers and diffusers with a single line of code.
* Provide models supported on a wide range of hardware, which are pre-compiled and require no JIT.
> It's important to note that specific quality degradation can vary from model to model. For instance, with an S model, you can have 0.5% degradation as well.
-----
## Inference
Currently, our demo model only supports 1024x1024, 768x768 and 512x512 outputs without batching (for B200 - only 1024x1024). This will be updated in the near future.
To infer our models, you just need to replace `diffusers` import with `elastic_models.diffusers`:
```python
import torch
from elastic_models.diffusers import FluxPipeline
mode_name = 'black-forest-labs/FLUX.1-dev'
hf_token = ''
device = torch.device("cuda")
pipeline = FluxPipeline.from_pretrained(
mode_name,
torch_dtype=torch.bfloat16,
token=hf_token,
mode='S'
)
pipeline.to(device)
prompts = ["Kitten eating a banana"]
output = pipeline(prompt=prompts)
for prompt, output_image in zip(prompts, output.images):
output_image.save((prompt.replace(' ', '_') + '.png'))
```
### Installation
__System requirements:__
* GPUs: H100, L40s, B200
* CPU: AMD, Intel
* Python: 3.10-3.12
To work with our models just run these lines in your terminal:
```shell
pip install thestage
pip install elastic_models[nvidia]\
--index-url https://thestage.jfrog.io/artifactory/api/pypi/pypi-thestage-ai-production/simple\
--extra-index-url https://pypi.nvidia.com\
--extra-index-url https://pypi.org/simple
# or for blackwell support
pip install elastic_models[blackwell]\
--index-url https://thestage.jfrog.io/artifactory/api/pypi/pypi-thestage-ai-production/simple\
--extra-index-url https://pypi.nvidia.com\
--extra-index-url https://pypi.org/simple
pip install -U --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128
pip install -U --pre torchvision --index-url https://download.pytorch.org/whl/nightly/cu128
pip install flash_attn==2.7.3 --no-build-isolation
pip uninstall apex
```
Then go to [app.thestage.ai](https://app.thestage.ai), login and generate API token from your profile page. Set up API token as follows:
```shell
thestage config set --api-token <YOUR_API_TOKEN>
```
Congrats, now you can use accelerated models!
----
## Benchmarks
Benchmarking is one of the most important procedures during model acceleration. We aim to provide clear performance metrics for models using our algorithms.
### Quality benchmarks
For quality evaluation we have used: PSNR, SSIM and CLIP score. PSNR and SSIM were computed using outputs of original model.
| Metric/Model | S | M | L | XL | Original |
|---------------|---|---|---|----|----------|
| PSNR | 30.22 | 30.24 | 30.38 | inf | inf |
| SSIM | 0.72 | 0.72 | 0.76 | 1.0 | 1.0 |
| CLIP | 12.49 | 12.51 | 12.69 | 12.41 | 12.41|
### Latency benchmarks
Time in seconds to generate one image 1024x1024
| GPU/Model | S | M | L | XL | Original |
|-----------|-----|---|---|----|----------|
| H100 | 2.71 | 3.0 | 3.18 | 4.17 | 6.46 |
| L40s | 8.5 | 9.29 | 9.29 | 13.2 | 16|
| B200 | 1.89 | 2.04 | 2.12 | 2.23 | 4.4|
| GeForce RTX 5090 | 5.53 | - | - | - | -|
## Links
* __Platform__: [app.thestage.ai](https://app.thestage.ai)
<!-- * __Elastic models Github__: [app.thestage.ai](app.thestage.ai) -->
* __Subscribe for updates__: [TheStageAI X](https://x.com/TheStageAI)
* __Contact email__: contact@thestage.ai |