let5see/IOPaint
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add enable_low_mem

Browse Source
This commit is contained in:
Qing
2024-01-09 22:42:48 +08:00
parent a49c3f86d3
commit 8ed969eec1
7 changed files with 147 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
iopaint/model/sdxl.py
View File

@@ -9,7 +9,7 @@ from loguru import logger
from iopaint.schema import InpaintRequest, ModelType
from .base import DiffusionInpaintModel
from .utils import handle_from_pretrained_exceptions, get_torch_dtype
from .utils import handle_from_pretrained_exceptions, get_torch_dtype, enable_low_mem
class SDXL(DiffusionInpaintModel):
@@ -47,10 +47,7 @@ class SDXL(DiffusionInpaintModel):
variant="fp16",
)
if torch.backends.mps.is_available():
# MPS: Recommended RAM < 64 GB https://huggingface.co/docs/diffusers/optimization/mps
# CUDA: Don't enable attention slicing if you're already using `scaled_dot_product_attention` (SDPA) from PyTorch 2.0 or xFormers. https://huggingface.co/docs/diffusers/v0.25.0/en/api/pipelines/stable_diffusion/image_variation#diffusers.StableDiffusionImageVariationPipeline.enable_attention_slicing
self.model.enable_attention_slicing()
enable_low_mem(self.model, kwargs.get("low_mem", False))
if kwargs.get("cpu_offload", False) and use_gpu:
logger.info("Enable sequential cpu offload")