Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

docs: added multi model load and store example #4575

Merged
merged 14 commits into from
Apr 8, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 82 additions & 0 deletions examples/model_load_store_lora/README.md
Original file line number Diff line number Diff line change
@@ -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

lycheel1 marked this conversation as resolved.
Show resolved Hide resolved
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.

lycheel1 marked this conversation as resolved.
Show resolved Hide resolved
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)
lycheel1 marked this conversation as resolved.
Show resolved Hide resolved
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')
```
lycheel1 marked this conversation as resolved.
Show resolved Hide resolved
lycheel1 marked this conversation as resolved.
Show resolved Hide resolved

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).
13 changes: 13 additions & 0 deletions examples/model_load_store_lora/bentofile.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
service: "service.py:StableDiffusion"
lycheel1 marked this conversation as resolved.
Show resolved Hide resolved
model:
tag: sd
labels:
owner: bentoml-team
project: gallery
include:
- "service.py"
python:
requirements_txt: "./requirements.txt"
lock_packages: false
docker:
distro: debian
18 changes: 18 additions & 0 deletions examples/model_load_store_lora/import_model.py
Original file line number Diff line number Diff line change
@@ -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}")
5 changes: 5 additions & 0 deletions examples/model_load_store_lora/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
diffusers
bentoml
transformers
torch
accelerate
51 changes: 51 additions & 0 deletions examples/model_load_store_lora/service.py
Original file line number Diff line number Diff line change
@@ -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},
larme marked this conversation as resolved.
Show resolved Hide resolved
)
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
Loading