add tomesd optimization

chainyo · May 3, 2023 · f2c8439 · f2c8439
1 parent 356a9c8
commit f2c8439
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/docker/api/requirements.txt b/docker/api/requirements.txt
@@ -8,4 +8,5 @@ pydantic>=1.8
 python-dotenv>=0.21
 python-jose>=3.3
 python-multipart>=0.0.5
+tomesd>=0.1.2
 transformers>=4.25
diff --git a/picaisso/api/diffusion_model.py b/picaisso/api/diffusion_model.py
@@ -7,6 +7,8 @@
 
 from diffusers import StableDiffusionPipeline
 
+import tomesd
+
 import torch
 from torch import autocast
 # Torch optimizations for inference
@@ -32,6 +34,7 @@ def __init__(self, model_name: str, dtype: str, n_steps: int, max_batch_size: in
         assert torch.cuda.is_available(), "CUDA is not available"
         self.pipeline = StableDiffusionPipeline.from_pretrained(self.model, torch_dtype=self.dtype)
         self.pipeline.to("cuda:0")
+        tomesd.apply_patch(self.pipeline, ratio=0.5)
 
 
     def schedule_processing_if_needed(self):