Skip to content

Commit

Permalink
fix sd1.5 link, add http/2 test
Browse files Browse the repository at this point in the history
Signed-off-by: Keming <kemingy94@gmail.com>
  • Loading branch information
kemingy committed Sep 22, 2024
1 parent 8743df1 commit 36322de
Show file tree
Hide file tree
Showing 5 changed files with 96 additions and 75 deletions.
134 changes: 67 additions & 67 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

7 changes: 3 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ from mosec.mixin import MsgpackMixin
logger = get_logger()
```

Then, we **build an API** for clients to query a text prompt and obtain an image based on the [stable-diffusion-v1-5 model](https://huggingface.co/runwayml/stable-diffusion-v1-5) in just 3 steps.
Then, we **build an API** for clients to query a text prompt and obtain an image based on the [stable-diffusion-v1-5 model](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) in just 3 steps.

1) Define your service as a class which inherits `mosec.Worker`. Here we also inherit `MsgpackMixin` to employ the [msgpack](https://msgpack.org/index.html) serialization format<sup>(a)</sup></a>.

Expand All @@ -104,10 +104,9 @@ Then, we **build an API** for clients to query a text prompt and obtain an image
class StableDiffusion(MsgpackMixin, Worker):
def __init__(self):
self.pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
"sd-legacy/stable-diffusion-v1-5", torch_dtype=torch.float16
)
device = "cuda" if torch.cuda.is_available() else "cpu"
self.pipe = self.pipe.to(device)
self.pipe.enable_model_cpu_offload()
self.example = ["useless example prompt"] * 4 # warmup (batch_size=4)

def forward(self, data: List[str]) -> List[memoryview]:
Expand Down
6 changes: 3 additions & 3 deletions examples/stable_diffusion/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,10 +27,10 @@
class StableDiffusion(MsgpackMixin, Worker):
def __init__(self):
self.pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
"sd-legacy/stable-diffusion-v1-5",
torch_dtype=torch.float16,
)
device = "cuda" if torch.cuda.is_available() else "cpu"
self.pipe = self.pipe.to(device) # type: ignore
self.pipe.enable_model_cpu_offload()
self.example = ["useless example prompt"] * 4 # warmup (bs=4)

def forward(self, data: List[str]) -> List[memoryview]:
Expand Down
Loading

0 comments on commit 36322de

Please sign in to comment.