Update README.md

README.md

@@ -11,6 +11,8 @@ pipeline_tag: text-to-image
 
 ## Overview
 
+---
+
 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.
@@ -20,12 +22,15 @@ ElasticModels are the models produced by TheStage AI ANNA: Automated Neural Netw
 
 Models can be accessed via TheStage AI Python SDK: ElasticModels, or deployed as Docker containers with REST API endpoints (see Deploy section).
 
+## Installation
+
 ---
 
-## Installation
 
 ### System Requirements
 
+---
+
 | **Property**| **Value** |
  | ---  | ---  |
 | **GPU** | L40s, RTX 5090, H100, B200 |
@@ -36,6 +41,8 @@ Models can be accessed via TheStage AI Python SDK: ElasticModels, or deployed as
 
 ### TheStage AI Access token setup
 
+---
+
 Install TheStage AI CLI and setup API token:
 
 ```bash
@@ -45,6 +52,8 @@ thestage config set --access-token <YOUR_ACCESS_TOKEN>
 
 ### ElasticModels installation
 
+---
+
 Install TheStage Elastic Models package:
 
 ```bash
@@ -52,9 +61,10 @@ pip install 'thestage-elastic-models[nvidia]' \
     --extra-index-url https://thestage.jfrog.io/artifactory/api/pypi/pypi-thestage-ai-production/simple
 ```
 
+## Usage example
+
 ---
 
-## Usage example
 
 Elastic Models provides the same interface as HuggingFace Diffusers. Here is an example of how to use the FLUX.1-dev model:
 
@@ -84,9 +94,10 @@ for prompt, output_image in zip(prompts, output.images):
 ```
 
 
+## Quality Benchmarks
+
 ---
 
-## 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.
 
@@ -94,6 +105,8 @@ We have used PartiPrompts and DrawBench datasets to evaluate the quality of imag
 
 ### Quality Benchmark Results
 
+---
+
 | **Metric/Model Size**| **S**| **M**| **L**| **XL**| **Original** |
  | ---  | ---  | ---  | ---  | ---  | ---  |
 | **ARNIQA (PartiPrompts)** | 64.1 | 63.2 | 61.9 | 66.8 | 66.9 |
@@ -106,17 +119,19 @@ We have used PartiPrompts and DrawBench datasets to evaluate the quality of imag
 | **SSIM (PartiPrompts)** | 0.72 | 0.72 | 0.76 | 1.0 | 1.0 |
 
 
+## Datasets
+
 ---
 
-## 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
+
 ---
 
-## 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.
@@ -125,9 +140,10 @@ We have used PartiPrompts and DrawBench datasets to evaluate the quality of imag
 - **SSIM**: Structural Similarity Index measuring perceptual similarity between generated by accelerated model and original model images.
 
 
+## Latency Benchmarks
+
 ---
 
-## 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.
 
@@ -135,6 +151,8 @@ We have measured the latency of different sizes of FLUX.1-dev model (S, M, L, XL
 
 ### 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** |
@@ -145,9 +163,10 @@ Latency (in seconds) for generating a 1024x1024 image using different model size
 | **GeForce RTX 5090** | 5.79 | N/A | N/A | N/A | N/A |
 
 
+## Benchmarking Methodology
+
 ---
 
-## 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.
 
@@ -163,9 +182,10 @@ The benchmarking was performed on a single GPU with a batch size of 1. Each mode
 >    - Record the end time and calculate the latency for that iteration.
 > 5. Calculate the average latency over all iterations.
 
+## Reproduce benchmarking
+
 ---
 
-## Reproduce benchmarking
 
 ```python
 import torch
@@ -224,9 +244,10 @@ print(f"Average Latency over {num_runs} runs: {average_latency} seconds")
 ```
 
 
+## Serving with Docker Image
+
 ---
 
-## 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.
@@ -234,15 +255,12 @@ 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:
+Pull docker image and start inference container:
 
 ```bash
-docker pull <IMAGE_NAME>
+docker pull public.ecr.aws/i3f7g5s7/thestage/elastic-models:0.2.0-diffusers-24.09c
 ```
 ```bash
 docker run --rm -ti \
@@ -255,7 +273,7 @@ docker run --rm -ti \
   -e HUGGINGFACE_ACCESS_TOKEN=<HUGGINGFACE_ACCESS_TOKEN> \
   -e THESTAGE_AUTH_TOKEN=<THESTAGE_ACCESS_TOKEN> \
   -v /mnt/hf_cache:/root/.cache/huggingface \
-  <IMAGE_NAME_DEPENDING_ON_YOUR_GPU>
+  public.ecr.aws/i3f7g5s7/thestage/elastic-models:0.2.0-diffusers-24.09c
 ```
 
 | **Parameter**              | **Description**                                                                                      |
@@ -265,11 +283,11 @@ docker run --rm -ti \
 | `<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
 
 ---
 
-## Invocation
 
 You can invoke the endpoint using CURL as follows:
 
@@ -343,16 +361,21 @@ with open("thestage_image.webp", "wb") as f:
     f.write(response.content)
 ```
 
+## Endpoint Parameters
+
 ---
 
-## Endpoint Parameters
 
 ### Method
 
+---
+
 > **POST** `/v1/images/generations`
 
 ### Header Parameters
 
+---
+
 > `Authorization`: `string`
 >
 > Bearer token for authentication. Should match the `AUTH_TOKEN` set during container startup.
@@ -367,6 +390,8 @@ with open("thestage_image.webp", "wb") as f:
 
 ### Input Body
 
+---
+
 > `prompt` : `string`
 >
 > The text prompt to generate an image for.
@@ -400,14 +425,17 @@ with open("thestage_image.webp", "wb") as f:
 >
 > Guidance scale for classifier-free guidance. Higher values increase adherence to the prompt.
 
+## Deploy on Modal
+
 ---
 
-## 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
 
+---
+
 ```shell
 git clone https://github.com/TheStageAI/ElasticModels.git
 cd ElasticModels/examples/modal
@@ -415,6 +443,8 @@ cd ElasticModels/examples/modal
 
 ### Configuration of environment variables
 
+---
+
 Set your environment variables in `modal_serving.py`:
 
 ```python
@@ -433,6 +463,8 @@ ENVS = {
 
 ### Configuration of GPUs
 
+---
+
 Set your desired GPU type and autoscaling variables in `modal_serving.py`:
 
 ```python
@@ -459,6 +491,8 @@ def serve():
 
 ### Run serving
 
+---
+
 ```shell
 modal serve modal_serving.py
 ```
@@ -466,6 +500,8 @@ modal serve modal_serving.py
 
 ## Links
 
+---
+
 * __Platform__: [app.thestage.ai](https://app.thestage.ai)
 * __Subscribe for updates__: [TheStageAI X](https://x.com/TheStageAI)
 * __Contact email__: contact@thestage.ai