File size: 4,185 Bytes
7766b8e ba2a039 7766b8e 7d464d8 1011ea0 7766b8e da420ac c1e6f35 7766b8e ade2e09 7766b8e 21a480e 7766b8e 7d93b81 7766b8e c1e6f35 4df7a38 765e122 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 |
---
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 outputs without batching. 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
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|
## 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__: [email protected] |