Update README.md

README.md

@@ -33,54 +33,32 @@ __Goals of elastic models:__
 
 ## Inference
 
-To infer our models, you just need to replace `transformers` import with `elastic_models.transformers`:
+To infer our models, you just need to replace `transformers` import with `elastic_models.diffusers`:
 
 ```python
 import torch
-from transformers import AutoTokenizer
-from elastic_models.transformers import AutoModelForCausalLM
+from elastic_models.diffusers import FluxPipeline
 
 # Currently we require to have your HF token
 # as we use original weights for part of layers and
 # model confugaration as well
-model_name = "mistralai/Mistral-7B-Instruct-v0.3"
+mode_name = 'black-forest-labs/FLUX.1-dev'
 hf_token = ''
-hf_cache_dir = ''
 device = torch.device("cuda")
 
-# Create mode
-tokenizer = AutoTokenizer.from_pretrained(
-    model_name, token=hf_token
-)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name, 
-    token=hf_token,
-    cache_dir=hf_cache_dir,
+pipeline = FluxPipeline.from_pretrained(
+    mode_name,
     torch_dtype=torch.bfloat16,
-    attn_implementation="sdpa",
-    mode='s'
-).to(device)
-model.generation_config.pad_token_id = tokenizer.eos_token_id
-
-# Inference simple as transformers library
-prompt = "Describe basics of DNNs quantization."
-inputs = tokenizer(prompt, return_tensors="pt")
-inputs.to(device)
-
-with torch.inference_mode:
-    generate_ids = model.generate(**inputs, max_length=500)
-
-input_len = inputs['input_ids'].shape[1]
-generate_ids = generate_ids[:, input_len:]
-output = tokenizer.batch_decode(
-    generate_ids,
-    skip_special_tokens=True, 
-    clean_up_tokenization_spaces=False
-)[0]
-
-# Validate answer
-print(f"# Q:\n{prompt}\n")
-print(f"# A:\n{output}\n")
+    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