| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import os |
| 4 | |
| 5 | import pybase64 as base64 |
| 6 | import torch |
| 7 | |
| 8 | from vllm import LLM |
| 9 | |
| 10 | # This example shows how to perform an offline inference that generates |
| 11 | # multimodal data. In this specific case this example will take a geotiff |
| 12 | # image as input, process it using the multimodal data processor, and |
| 13 | # perform inference. |
| 14 | # Requirements: |
| 15 | # - install TerraTorch v1.1 (or later): |
| 16 | # pip install terratorch>=v1.1 |
| 17 | |
| 18 | |
| 19 | def main(): |
| 20 | torch.set_default_dtype(torch.float16) |
| 21 | image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff" # noqa: E501 |
| 22 | |
| 23 | img_data = dict( |
| 24 | data=image_url, |
| 25 | data_format="url", |
| 26 | image_format="tiff", |
| 27 | out_data_format="b64_json", |
| 28 | ) |
| 29 | |
| 30 | prompt = dict(data=img_data) |
| 31 | |
| 32 | llm = LLM( |
| 33 | model="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11", |
| 34 | skip_tokenizer_init=True, |
| 35 | trust_remote_code=True, |
| 36 | enforce_eager=True, |
| 37 | # Limit the maximum number of parallel requests |
| 38 | # to avoid the model going OOM. |
| 39 | # The maximum number depends on the available GPU memory |
| 40 | max_num_seqs=32, |
| 41 | io_processor_plugin="terratorch_segmentation", |
| 42 | model_impl="terratorch", |
| 43 | enable_mm_embeds=True, |
| 44 | ) |
| 45 | |
| 46 | pooler_output = llm.encode(prompt, pooling_task="plugin") |
| 47 | output = pooler_output[0].outputs |
| 48 | |
| 49 | print(output) |
| 50 | decoded_data = base64.b64decode(output.data) |
| 51 | |
| 52 | file_path = os.path.join(os.getcwd(), "offline_prediction.tiff") |
| 53 | with open(file_path, "wb") as f: |
| 54 | f.write(decoded_data) |
| 55 | |
| 56 | print(f"Output file path: {file_path}") |
| 57 | |
| 58 | |
| 59 | if __name__ == "__main__": |
| 60 | main() |