diff --git a/examples/model_load_store_lora/README.md b/examples/model_load_store_lora/README.md new file mode 100644 index 00000000000..10c5fe743fb --- /dev/null +++ b/examples/model_load_store_lora/README.md @@ -0,0 +1,82 @@ +# Model Load and Store with Multiple Models + +In this example, we will use LoRA and StableDiffusion to demonstrate how to load and store multiple models both for local models and remote models. + +## Install Dependencies +```bash +python -m venv model-load-store-lora +source model-load-store-lora/bin/activate +pip install -r requirement.txt +``` + +## Import Models + +Run the import_model.py script in this example to save models to the BentoML Model Store. +```python +import bentoml +import shutil +import os +from diffusers import StableDiffusionPipeline + +model_id = 'CompVis/stable-diffusion-v1-4' +pipeline = StableDiffusionPipeline.from_pretrained(model_id, use_safetensors=True) + +local_lora = "./lora" +with bentoml.models.create( + name='sd', # All models are created under ~/bentoml/models/{name = sd} +) as model_ref: + local_diffusion_path = model_ref.path_of("diffusion") + local_lora_path = model_ref.path_of("lora") + pipeline.save_pretrained(local_diffusion_path) + shutil.copytree(local_lora, local_lora_path, dirs_exist_ok=True) + print(f"Model saved: {model_ref}") +``` + +When storing multiple models, you can simply create new directories corresponding to these models inside the `bentoml.models.create()` context manager. For local models, use `shutil.copytree()` to copy the local models and files. For remote models, use the respective library tools to save the models to the directory you just created. + +To view all the models in the Model Store, run bentoml models list: +```bash +$ bentoml models list + Tag Module Size Creation Time + sd:l42k7zxcmksxhyln 5.11 GiB 2024-03-15 00:24:22 +``` + +## Load Models +Now that the models are saved, we can get the directory of the saved models by using `model_ref = bentoml.models.get("sd:latest")`. You can get the base path by calling `model_ref.path`. This will return the sd directory you have just created. To +get each individual model subpath, use `model_ref.path_of("{subpath}")`. + +```python +class StableDiffusion: + model_ref = bentoml.models.get("sd:latest") + + def __init__(self) -> None: + # Load model into pipeline + self.diffusion_ref = model_ref.path_of("diffusion") + self.lora_ref = model_ref.path_of("lora") + self.stable_diffusion_txt2img = StableDiffusionPipeline.from_pretrained(self.diffusion_ref, use_safetensors=True) + self.stable_diffusion_txt2img.unet.load_attn_procs(self.lora_ref) + self.stable_diffusion_img2img = StableDiffusionImg2ImgPipeline( + vae=self.stable_diffusion_txt2img.vae, + text_encoder=self.stable_diffusion_txt2img.text_encoder, + tokenizer=self.stable_diffusion_txt2img.tokenizer, + unet=self.stable_diffusion_txt2img.unet, + scheduler=self.stable_diffusion_txt2img.scheduler, + safety_checker=None, + feature_extractor=None, + requires_safety_checker=False, + ) + self.stable_diffusion_txt2img.to('cuda') + self.stable_diffusion_img2img.to('cuda') +``` + +To serve this Service locally, run: +```bash +bentoml serve . +``` + +## Deploy to BentoCloud +Run the following command to deploy this example to BentoCloud for better management and scalability. [Sign up](https://www.bentoml.com/) if you haven't got a BentoCloud account. +```bash +bentoml deploy . +``` +For more information, see [Create Deployments](https://docs.bentoml.com/en/latest/bentocloud/how-tos/create-deployments.html). diff --git a/examples/model_load_store_lora/bentofile.yaml b/examples/model_load_store_lora/bentofile.yaml new file mode 100644 index 00000000000..d4041fa3d69 --- /dev/null +++ b/examples/model_load_store_lora/bentofile.yaml @@ -0,0 +1,13 @@ +service: "service.py:StableDiffusion" +model: + tag: sd +labels: + owner: bentoml-team + project: gallery +include: + - "service.py" +python: + requirements_txt: "./requirements.txt" + lock_packages: false +docker: + distro: debian diff --git a/examples/model_load_store_lora/import_model.py b/examples/model_load_store_lora/import_model.py new file mode 100644 index 00000000000..36124642f8f --- /dev/null +++ b/examples/model_load_store_lora/import_model.py @@ -0,0 +1,18 @@ +import shutil + +from diffusers import StableDiffusionPipeline + +import bentoml + +model_id = "CompVis/stable-diffusion-v1-4" +pipeline = StableDiffusionPipeline.from_pretrained(model_id, use_safetensors=True) + +local_lora = "./lora" +with bentoml.models.create( + name="sd", +) as model_ref: + local_diffusion_path = model_ref.path_of("diffusion") + local_lora_path = model_ref.path_of("lora") + pipeline.save_pretrained(local_diffusion_path) + shutil.copytree(local_lora, local_lora_path, dirs_exist_ok=True) + print(f"Model saved: {model_ref}") diff --git a/examples/model_load_store_lora/requirements.txt b/examples/model_load_store_lora/requirements.txt new file mode 100644 index 00000000000..a5c279a38eb --- /dev/null +++ b/examples/model_load_store_lora/requirements.txt @@ -0,0 +1,5 @@ +diffusers +bentoml +transformers +torch +accelerate diff --git a/examples/model_load_store_lora/service.py b/examples/model_load_store_lora/service.py new file mode 100644 index 00000000000..0f221b3cff9 --- /dev/null +++ b/examples/model_load_store_lora/service.py @@ -0,0 +1,51 @@ +import torch +from PIL.Image import Image + +import bentoml + + +@bentoml.service( + resources={ + "memory": "500MiB", + "gpu": 1, + "gpu_type": "nvidia-tesla-t4", + }, + traffic={"timeout": 600}, +) +class StableDiffusion: + model_ref = bentoml.models.get("sd:latest") + + def __init__(self) -> None: + from diffusers import StableDiffusionPipeline + + # Load model into pipeline + self.diffusion_ref = self.model_ref.path_of("diffusion") + self.lora_ref = self.model_ref.path_of("lora") + self.stable_diffusion_txt2img = StableDiffusionPipeline.from_pretrained( + self.diffusion_ref, use_safetensors=True + ) + self.stable_diffusion_txt2img.unet.load_attn_procs(self.lora_ref) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.stable_diffusion_txt2img.to(device) + + @bentoml.api + def txt2img( + self, + prompt: str = "A Pokemon with blue eyes.", + height: int = 768, + width: int = 768, + num_inference_steps: int = 30, + guidance_scale: float = 7.5, + eta: int = 0, + ) -> Image: + input_data = dict( + prompt=prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + eta=eta, + ) + res = self.stable_diffusion_txt2img(**input_data) + image = res[0][0] + return image