Revert README.md to pre-March-3 version (undo broken template changes)

README.md

@@ -1,74 +1,45 @@
 ---
+license: apache-2.0
 base_model:
 - black-forest-labs/FLUX.1-dev
+base_model_relation: quantized
 pipeline_tag: text-to-image
-language:
-- en
 ---
 
-# Elastic model: FLUX.1-dev
 
-## Overview
+# Elastic model: Fastest self-serving models. FLUX.1-dev.
 
-----
-
-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:
-
-- **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%.
-
-Models can be accessed via TheStage AI Python SDK: ElasticModels, or deployed as Docker containers with REST API endpoints (see Deploy section).
-
-## Installation
-
----
+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:
 
-### System Requirements
+* __XL__: Mathematically equivalent neural network, optimized with our DNN compiler. 
 
-| **Property**| **Value** |
- | ---  | ---  |
-| **GPU** | L40s, RTX 5090, H100, B200 |
-| **Python Version** | 3.10-3.12 |
-| **CPU** | Intel/AMD x86_64 |
-| **CUDA Version** | 12.8+ |
+* __L__: Near lossless model, with less than 1% degradation obtained on corresponding benchmarks.
 
+* __M__: Faster model, with accuracy degradation less than 1.5%.
 
-### TheStage AI Access token setup
+* __S__: The fastest model, with accuracy degradation less than 2%.
 
-Install TheStage AI CLI and setup API token:
 
-```bash
-pip install thestage
-thestage config set --api-token <YOUR_ACCESS_TOKEN>
-```
+__Goals of Elastic Models:__
 
-### ElasticModels installation
+* 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.
 
-Install TheStage Elastic Models package:
-
-```bash
-pip install 'thestage-elastic-models[nvidia]' \
-    --extra-index-url https://thestage.jfrog.io/artifactory/api/pypi/pypi-thestage-ai-production/simple
-```
+> 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.
 
-If you want to run on Nvidia Blackwell architecture, you need to install package as follows:
+-----
 
-```bash
-pip install 'thestage-elastic-models[blackwell]' \
-    --extra-index-url https://thestage.jfrog.io/artifactory/api/pypi/pypi-thestage-ai-production/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
-```
 
-## Usage example
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/67991798ae62bd1f17cc22ed/2FXY0tqSGqZq76j5Tz4Vi.png)
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/6799fc8e150f5a4014b030ca/CuuzzA_csoRnzbaZq1U1x.png)
 
-----
+## Inference
 
-Elastic Models provides the same interface as HuggingFace Diffusers. Here is an example of how to use the FLUX.1-dev model:
+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
@@ -82,8 +53,6 @@ pipeline = FluxPipeline.from_pretrained(
     mode_name,
     torch_dtype=torch.bfloat16,
     token=hf_token,
-    # 'original' for original model
-    # 'S', 'M', 'L', 'XL' for accelerated models
     mode='S'
 )
 pipeline.to(device)
@@ -95,387 +64,69 @@ for prompt, output_image in zip(prompts, output.images):
     output_image.save((prompt.replace(' ', '_') + '.png'))
 ```
 
+### Installation
 
-## Quality Benchmarks
-
-------------
-
-We have used PartiPrompts and DrawBench datasets to evaluate the quality of images generated by different sizes of FLUX.1-dev models (S, M, L, XL) compared to the original model. The evaluation metrics include ARNIQA, CLIP IQA, PSNR, SSIM, and VQA Faithfulness.
-
-![Quality Benchmarking](https://cdn.thestage.ai/production/cms_file_upload/1767818654-9b8b5892-54bb-4987-bcf9-db2b72af2bec/flux-dev-quality.png)
-
-### Quality Benchmark Results
-
-| **Metric/Model Size**| **S**| **M**| **L**| **XL**| **Original** |
- | ---  | ---  | ---  | ---  | ---  | ---  |
-| **ARNIQA (PartiPrompts)** | 64.1 | 63.2 | 61.9 | 66.8 | 66.9 |
-| **ARNIQA (DrawBench)** | 64.3 | 63.5 | 63.6 | 68.2 | 68.5 |
-| **CLIP IQA (PartiPrompts)** | 85.5 | 86.4 | 83.8 | 88.3 | 87.9 |
-| **CLIP IQA (DrawBench)** | 86.4 | 86.5 | 84.5 | 89.5 | 90.0 |
-| **VQA Faithfulness (PartiPrompts)** | 87.5 | 85.5 | 85.5 | 85.5 | 88.6 |
-| **VQA Faithfulness (DrawBench)** | 69.3 | 64.7 | 64.8 | 67.8 | 65.2 |
-| **PSNR (PartiPrompts)** | 30.22 | 30.24 | 30.38 | N/A | N/A |
-| **SSIM (PartiPrompts)** | 0.72 | 0.72 | 0.76 | 1.0 | 1.0 |
-
-
-## Datasets
-
--------
-
-- **PartiPrompts**: A benchmark dataset created by Google Research, containing 1,632 diverse and challenging prompts that test various aspects of text-to-image generation models. It includes categories such as abstract concepts, complex compositions, properties and attributes, counting and numbers, text rendering, artistic styles, and fine-grained details.
-
-- **DrawBench**: A comprehensive benchmark dataset developed by Google Research, containing 200 carefully curated prompts designed to test specific capabilities and challenge areas of diffusion models. It includes categories such as colors, counting, conflicting requirements, DALL-E inspired prompts, detailed descriptions, misspellings, positional relationships, rare words, Reddit user prompts, and text generation.
-
-## Metrics
-
-----------
-
-- **ARNIQA**: No-reference image quality assessment metric that predicts perceptual quality without reference images.
-- **CLIP_IQA**: No-reference image quality metric using contrastive learning to assess image quality without references.
-- **VQA Faithfulness**: Metric measuring how accurately generated images represent the text prompts.
-- **PSNR**: Peak Signal-to-Noise Ratio measuring similarity between generated by accelerated model and original model images.
-- **SSIM**: Structural Similarity Index measuring perceptual similarity between generated by accelerated model and original model images.
-
-
-## Latency Benchmarks
-
------
-
-We have measured the latency of different sizes of FLUX.1-dev model (S, M, L, XL, original) on various GPUs. The measurements were taken for generating images of size 1024x1024 pixels.
-
-![Latency Benchmarking](https://cdn.thestage.ai/production/cms_file_upload/1767818701-7eef4234-acf0-4f2f-8577-ac6c2527cbaa/flux-dev-performance.png)
-
-### Latency Benchmark Results
-
-Latency (in seconds) for generating a 1024x1024 image using different model sizes on various hardware setups.
-
-| **GPU/Model Size**| **S**| **M**| **L**| **XL**| **Original** |
- | ---  | ---  | ---  | ---  | ---  | ---  |
-| **H100** | 2.88 | 3.06 | 3.25 | 4.18 | 6.46 |
-| **L40s** | 9.22 | 10.07 | 10.67 | 14.39 | 16 |
-| **B200** | 1.89 | 2.04 | 2.12 | 2.23 | 4.4 |
-| **GeForce RTX 5090** | 5.53 | N/A | N/A | N/A | N/A |
-
-
-## Benchmarking Methodology
-
-----
-
-The benchmarking was performed on a single GPU with a batch size of 1. Each model was run for 10 iterations, and the average latency was calculated.
-
-> **Algorithm summary:**
-> 1. Load the FLUX.1-dev model with the specified size (S, M, L, XL, original).
-> 2. Move the model to the GPU.
-> 3. Prepare a sample prompt for image generation.
-> 4. Run the model for a number of iterations (e.g., 10) 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 image using the model.
->    - Synchronize the GPU again.
->    - Record the end time and calculate the latency for that iteration.
-> 5. Calculate the average latency over all iterations.
-
-## Reproduce benchmarking
-
-----
-
-```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,
-    # 'original' for original model
-    # 'S', 'M', 'L', 'XL' for accelerated models
-    mode='S'
-)
-pipeline.to(device)
-
-prompt = ["Kitten eating a banana"]
-generate_kwargs={
-    "height": 1024,
-    "width": 1024,
-    "num_inference_steps": 4,
-    "cfg_scale": 0.0
-}
-
-def evaluate_pipline():
-    torch.cuda.synchronize()
-    start_time = time.time()
-    output = pipeline(
-        prompt=prompt,
-        **generate_kwargs
-    )
-    torch.cuda.synchronize()
-    end_time = time.time()
-
-    return end_time - start_time
-
-# Warm-up
-for _ in range(5):
-    evaluate_pipline()
-
-# Benchmarking
-num_runs = 10
-total_time = 0.0
-
-for _ in range(num_runs):
-    latency = evaluate_pipline()
-    total_time += latency
-
-average_latency = total_time / num_runs
-print(f"Average Latency over {num_runs} runs: {average_latency} seconds")
-```
-
-
-## Serving with Docker Image
-
-------------
-
-For serving with Nvidia GPUs, we provide ready-to-go Docker containers with OpenAI-compatible API endpoints.
-Using our containers you can set up an inference endpoint on any desired cloud/serverless providers as well as on-premise servers.
-You can also use this container to run inference through TheStage AI platform.
-
-### Prebuilt image from ECR
-
-| **GPU** | **Docker image name** |
-| --- | --- |
-| H100, L40s | `public.ecr.aws/i3f7g5s7/thestage/elastic-models:0.1.2-diffusers-nvidia-24.09b` |
-| B200, RTX 5090 | `public.ecr.aws/i3f7g5s7/thestage/elastic-models:0.1.2-diffusers-blackwell-24.09b` |
-
-Pull docker image for your Nvidia GPU and start inference container:
-
-```bash
-docker pull <IMAGE_NAME>
-```
-```bash
-docker run --rm -ti \
-  --name serving_thestage_model \
-  -p 8000:80 \
-  -e AUTH_TOKEN=<AUTH_TOKEN> \
-  -e MODEL_REPO=black-forest-labs/FLUX.1-dev \
-  -e MODEL_SIZE=<MODEL_SIZE> \
-  -e MODEL_BATCH=<MAX_BATCH_SIZE> \
-  -e HUGGINGFACE_ACCESS_TOKEN=<HUGGINGFACE_ACCESS_TOKEN> \
-  -e THESTAGE_AUTH_TOKEN=<THESTAGE_ACCESS_TOKEN> \
-  -v /mnt/hf_cache:/root/.cache/huggingface \
-  <IMAGE_NAME_DEPNDING_ON_YOUR_GPU>
-```
-
-| **Parameter**              | **Description**                                                                                      |
-|----------------------------|------------------------------------------------------------------------------------------------------|
-| `<MODEL_SIZE>`             | Available: S, M, L, XL.                                                                              |
-| `<MAX_BATCH_SIZE>`         | Maximum batch size to process in parallel.                                                           |
-| `<HUGGINGFACE_ACCESS_TOKEN>` | Hugging Face access token.                                                                         |
-| `<THESTAGE_ACCESS_TOKEN>`  | TheStage token generated on the platform (Profile -> Access tokens).                                 |
-| `<AUTH_TOKEN>`             | Token for endpoint authentication. You can set it to any random string; it must match the value used by the client. |
-| `<IMAGE_NAME>`             | Image name which you have pulled.                                                                    |
-
-## Invocation
-
-------
-
-You can invoke the endpoint using CURL as follows:
-
-```bash
-curl -X POST <http://127.0.0.1:8000/v1/images/generations>  \
-    -H "Authorization: Bearer <AUTH_TOKEN>"  \
-    -H "Content-Type: application/json" \
-    -H "X-Model-Name: flux-1-dev-<MODEL_SIZE>-bs<MAX_BATCH_SIZE>" \
-    -d '{
-          "prompt": "Cat eating banana",
-          "seed": 12,
-          "aspect_ratio": "1:1",
-          "guidance_scale": 6.5,
-          "num_inference_steps": 4
-        }' \
-    --output sunset.webp -D -
-```
-
-Or using Python requests:
 
-```python
-import requests
-import json
-url = "http://127.0.0.1:8000/v1/images/generations"
-payload = json.dumps({
-  "prompt": "sunset",
-  "seed": 12,
-  "aspect_ratio": "1:1",
-  "guidance_scale": 6.5,
-  "num_inference_steps": 4
-})
-headers = {
-  'Authorization: ''Bearer <AUTH_TOKEN>'',
-  'Content-Type': 'application/json',
-  'X-Model-Name': 'flux-1-dev-<MODEL_SIZE>-bs<MAX_BATCH_SIZE>'
-}
-response = requests.request("POST", url, headers=headers, data=payload)
-with open("sunset.webp", "wb") as f:
-    f.write(response.content)
-```
+__System requirements:__
+* GPUs: H100, L40s, B200
+* CPU: AMD, Intel
+* Python: 3.10-3.12
 
-Or using OpenAI python client:
 
-```python
-import os, base64, pathlib, json
-from openai import OpenAI
+To work with our models just run these lines in your terminal:
 
-BASE_URL = "http://<your_ip>/v1"
-API_KEY  = ""
-MODEL    = "flux-1-dev-<MODEL_SIZE>-bs<MAX_BATCH_SIZE>"
+```shell
+pip install thestage
+pip install 'thestage-elastic-models[nvidia]' --extra-index-url https://thestage.jfrog.io/artifactory/api/pypi/pypi-thestage-ai-production/simple
 
-client = OpenAI(
-    api_key=API_KEY,
-    base_url=BASE_URL,
-    default_headers={"X-Model-Name": MODEL}
-)
+# or for blackwell support
+pip install 'thestage-elastic-models[blackwell]' --extra-index-url https://thestage.jfrog.io/artifactory/api/pypi/pypi-thestage-ai-production/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
 
-response = client.with_raw_response.images.generate(
-    model=MODEL,
-    prompt="Cat eating banana",
-    n=1,
-    extra_body={
-        "seed": 111,
-        "aspect_ratio": "1:1",
-        "guidance_scale": 3.5,
-        "num_inference_steps": 4
-    },
-)
 
-with open("thestage_image.webp", "wb") as f:
-    f.write(response.content)
+pip install flash_attn==2.7.3 --no-build-isolation
+pip uninstall apex
 ```
 
-## Endpoint Parameters
-
--------------
-
-### Method
-
-> **POST** `/v1/images/generations`
-
-### Header Parameters
-
-> `Authorization`: `string`
->
-> Bearer token for authentication. Should match the `AUTH_TOKEN` set during container startup.
-
-> `Content-Type`: `string`
->
-> Must be set to `application/json`.
-
-> `X-Model-Name`: `string`
->
-> Specifies the model to use for generation. Format: `flux-1-dev-<size>-bs<batch_size>`, where `<size>` is one of `S`, `M`, `L`, `XL`, `original` and `<batch_size>` is the maximum batch size configured during container startup.
-
-### Input Body
-
-> `prompt` : `string`
->
-> The text prompt to generate an image for.
-
-> `seed`: `int32`
->
->  Random seed for generation.
-
-> `num_inference_steps`: `int32`
->
-> Number of diffusion steps to use for generation. Higher values yield better quality but take longer. Default is 28
-
-> `aspect_ratio`: `string`
->
->  Aspect ratio of the generated image. Supported values:
->  ```
->  "1:1": (1024, 1024),
->  "16:9": (1280, 736),
->  "21:9": (1280, 544),
->  "3:2": (1248, 832),
->  "2:3": (832, 1248),
->  "4:3": (1184, 896),
->  "3:4": (896, 1184),
->  "5:4": (1152, 928),
->  "4:5": (928, 1152),
->  "9:16": (736, 1280),
->  "9:21": (544, 1280)
->  ```
-
-> `guidance_scale`: float32
->
-> Guidance scale for classifier-free guidance. Higher values increase adherence to the prompt.
-
-## Deploy on Modal
-
------------------------
-
-For more details please use the tutorial [Modal deployment](https://docs.thestage.ai/tutorials/source/modal_thestage.html)
-
-### Clone modal serving code
+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
-git clone https://github.com/TheStageAI/ElasticModels.git
-cd ElasticModels/examples/modal
+thestage config set --api-token <YOUR_API_TOKEN>
 ```
 
-### Configuration of environment variables
+Congrats, now you can use accelerated models!
 
-Set your environment variables in `modal_serving.py`:
+----
 
-```python
-# modal_serving.py
-
-ENVS = {
-    "MODEL_REPO": "black-forest-labs/FLUX.1-dev",
-    "MODEL_BATCH": "4",
-    "THESTAGE_AUTH_TOKEN": "",
-    "HUGGINGFACE_ACCESS_TOKEN": "",
-    "PORT": "80",
-    "PORT_HEALTH": "80",
-    "HF_HOME": "/cache/huggingface",
-}
-```
+## Benchmarks
 
-### Configuration of GPUs
+Benchmarking is one of the most important procedures during model acceleration. We aim to provide clear performance metrics for models using our algorithms.
 
-Set your desired GPU type and autoscaling setup. variables in `modal_serving.py`:
+### Quality benchmarks
 
-```python
-# modal_serving.py
-
-@app.function(
-    image=image,
-    gpu="B200",
-    min_containers=8,
-    max_containers=8,
-    timeout=10000,
-    ephemeral_disk=600 * 1024,
-    volumes={"/opt/project/.cache": HF_CACHE},
-    startup_timeout=60*20
-)
-@modal.web_server(
-    80,
-    label="black-forest-labs/FLUX.1-dev-test",
-    startup_timeout=60*20
-)
-def serve():
-    pass
-```
+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|
 
-### Run serving
 
-```shell
-modal serve modal_serving.py
-```
+### 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
+* __Contact email__: contact@thestage.ai