optimize sd/paint_by_example modle VRAM usage
This commit is contained in:
@@ -22,16 +22,30 @@ class PaintByExample(InpaintModel):
|
||||
use_gpu = device == torch.device('cuda') and torch.cuda.is_available()
|
||||
torch_dtype = torch.float16 if use_gpu and fp16 else torch.float32
|
||||
model_kwargs = {"local_files_only": kwargs.get('local_files_only', False)}
|
||||
|
||||
if kwargs['disable_nsfw'] or kwargs.get('cpu_offload', False):
|
||||
logger.info("Disable Paint By Example Model NSFW checker")
|
||||
model_kwargs.update(dict(
|
||||
safety_checker=None,
|
||||
requires_safety_checker=False
|
||||
))
|
||||
|
||||
self.model = DiffusionPipeline.from_pretrained(
|
||||
"Fantasy-Studio/Paint-by-Example",
|
||||
torch_dtype=torch_dtype,
|
||||
**model_kwargs
|
||||
)
|
||||
self.model = self.model.to(device)
|
||||
|
||||
self.model.enable_attention_slicing()
|
||||
if kwargs.get('enable_xformers', False):
|
||||
self.model.enable_xformers_memory_efficient_attention()
|
||||
|
||||
# TODO: gpu_id
|
||||
if kwargs.get('cpu_offload', False) and torch.cuda.is_available():
|
||||
if kwargs.get('cpu_offload', False) and use_gpu:
|
||||
self.model.image_encoder = self.model.image_encoder.to(device)
|
||||
self.model.enable_sequential_cpu_offload(gpu_id=0)
|
||||
else:
|
||||
self.model = self.model.to(device)
|
||||
|
||||
def forward(self, image, mask, config: Config):
|
||||
"""Input image and output image have same size
|
||||
|
||||
Reference in New Issue
Block a user