diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 7ec0dd08a..7537d0209 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -71,7 +71,7 @@ def process_base(p: processing.StableDiffusionProcessing): guidance_rescale=p.diffusers_guidance_rescale, denoising_start=0 if use_refiner_start else p.refiner_start if use_denoise_start else None, denoising_end=p.refiner_start if use_refiner_start else 1 if use_denoise_start else None, - output_type='latent' if hasattr(shared.sd_model, 'vae') else 'np', + output_type='latent', clip_skip=p.clip_skip, desc='Base', ) @@ -217,7 +217,7 @@ def process_hires(p: processing.StableDiffusionProcessing, output): eta=shared.opts.scheduler_eta, guidance_scale=p.image_cfg_scale if p.image_cfg_scale is not None else p.cfg_scale, guidance_rescale=p.diffusers_guidance_rescale, - output_type='latent' if hasattr(shared.sd_model, 'vae') else 'np', + output_type='latent', clip_skip=p.clip_skip, image=output.images, strength=p.denoising_strength, @@ -278,7 +278,7 @@ def process_refine(p: processing.StableDiffusionProcessing, output): for i in range(len(output.images)): image = output.images[i] noise_level = round(350 * p.denoising_strength) - output_type='latent' if hasattr(shared.sd_refiner, 'vae') else 'np' + output_type='latent' if 'Upscale' in shared.sd_refiner.__class__.__name__ or 'Flux' in shared.sd_refiner.__class__.__name__: image = processing_vae.vae_decode(latents=image, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.width, height=p.height) p.extra_generation_params['Noise level'] = noise_level @@ -346,7 +346,11 @@ def process_decode(p: processing.StableDiffusionProcessing, output): if not hasattr(output, 'images') and hasattr(output, 'frames'): shared.log.debug(f'Generated: frames={len(output.frames[0])}') output.images = output.frames[0] - if hasattr(shared.sd_model, "vae") and output.images is not None and len(output.images) > 0: + model = shared.sd_model if not is_refiner_enabled(p) else shared.sd_refiner + if not hasattr(model, 'vae'): + if hasattr(model, 'pipe') and hasattr(model.pipe, 'vae'): + model = model.pipe + if hasattr(model, "vae") and output.images is not None and len(output.images) > 0: if p.hr_resize_mode > 0 and (p.hr_upscaler != 'None' or p.hr_resize_mode == 5): width = max(getattr(p, 'width', 0), getattr(p, 'hr_upscale_to_x', 0)) height = max(getattr(p, 'height', 0), getattr(p, 'hr_upscale_to_y', 0)) @@ -355,7 +359,7 @@ def process_decode(p: processing.StableDiffusionProcessing, output): height = getattr(p, 'height', 0) results = processing_vae.vae_decode( latents = output.images, - model = shared.sd_model if not is_refiner_enabled(p) else shared.sd_refiner, + model = model, full_quality = p.full_quality, width = width, height = height, diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 75347f416..5e6fa68f4 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -35,7 +35,7 @@ def create_latents(image, p, dtype=None, device=None): def full_vae_decode(latents, model): t0 = time.time() - if not hasattr(model, 'vae'): + if model is None or not hasattr(model, 'vae'): shared.log.error('VAE not found in model') return [] if debug: @@ -170,7 +170,14 @@ def vae_decode(latents, model, output_type='np', full_quality=True, width=None, if latents.shape[-1] <= 4: # not a latent, likely an image decoded = latents.float().cpu().numpy() elif full_quality and hasattr(shared.sd_model, "vae"): - decoded = full_vae_decode(latents=latents, model=shared.sd_model) + parent = shared.sd_model if hasattr(shared.sd_model, 'vae') else None + if hasattr(shared.sd_model, 'vae'): + parent = shared.sd_model + elif hasattr(shared.sd_model, 'pipe') and hasattr(shared.sd_model.pipe, 'vae'): + parent = shared.sd_model.pipe + else: + parent = None + decoded = full_vae_decode(latents=latents, model=parent) else: decoded = taesd_vae_decode(latents=latents) diff --git a/modules/pulid/attention_processor.py b/modules/pulid/attention_processor.py index 9756decc1..fa9e4ff82 100644 --- a/modules/pulid/attention_processor.py +++ b/modules/pulid/attention_processor.py @@ -345,10 +345,7 @@ class IDAttnProcessor2_0(torch.nn.Module): value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # the output of sdp = (batch, num_heads, seq_len, head_dim) - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - + hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False) hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) hidden_states = hidden_states.to(query.dtype) @@ -363,17 +360,15 @@ class IDAttnProcessor2_0(torch.nn.Module): dtype=id_embedding.dtype, device=id_embedding.device, ) - id_key = self.id_to_k(torch.cat((id_embedding, zero_tensor), dim=1)).to(query.dtype) - id_value = self.id_to_v(torch.cat((id_embedding, zero_tensor), dim=1)).to(query.dtype) + id_cat = torch.cat((id_embedding, zero_tensor), dim=1) + id_key = self.id_to_k(id_cat).to(query.dtype) + id_value = self.id_to_v(id_cat).to(query.dtype) id_key = id_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) id_value = id_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # the output of sdp = (batch, num_heads, seq_len, head_dim) - id_hidden_states = F.scaled_dot_product_attention( - query, id_key, id_value, attn_mask=None, dropout_p=0.0, is_causal=False - ) - + id_hidden_states = F.scaled_dot_product_attention(query, id_key, id_value, attn_mask=None, dropout_p=0.0, is_causal=False) id_hidden_states = id_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) id_hidden_states = id_hidden_states.to(query.dtype) diff --git a/modules/pulid/eva_clip/pretrained.py b/modules/pulid/eva_clip/pretrained.py index a1e55dcf3..bb87c540c 100644 --- a/modules/pulid/eva_clip/pretrained.py +++ b/modules/pulid/eva_clip/pretrained.py @@ -2,7 +2,6 @@ import hashlib import os import urllib import warnings -from functools import partial from typing import Dict, Union from tqdm import tqdm @@ -277,7 +276,7 @@ def download_pretrained_from_url( loop.update(len(buffer)) if expected_sha256 and not hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256): - raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match") + raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match") return download_target diff --git a/modules/pulid/pulid_sdxl.py b/modules/pulid/pulid_sdxl.py index fade7509d..8bd28dbe2 100644 --- a/modules/pulid/pulid_sdxl.py +++ b/modules/pulid/pulid_sdxl.py @@ -5,6 +5,7 @@ import numpy as np import torch import torch.nn as nn from diffusers import DPMSolverMultistepScheduler, StableDiffusionXLPipeline +from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput from huggingface_hub import hf_hub_download, snapshot_download from safetensors.torch import load_file @@ -24,14 +25,17 @@ from attention_processor import IDAttnProcessor2_0 as IDAttnProcessor class StableDiffusionXLPuLIDPipeline: - def __init__(self, pipe: StableDiffusionXLPipeline, device: torch.device, sampler=None, cache_dir=None): + def __init__(self, pipe: StableDiffusionXLPipeline, device: torch.device, dtype: torch.dtype=None, providers: list=None, offload: bool=True, sampler=None, cache_dir=None): super().__init__() self.device = device + self.dtype = dtype or torch.float16 self.pipe = pipe self.cache_dir = cache_dir + self.offload = offload self.hack_unet_attn_layers(self.pipe.unet) self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config) - self.id_adapter = IDFormer().to(self.device) + self.id_adapter = IDFormer().to(self.device, self.dtype) + self.providers = providers or ['CUDAExecutionProvider', 'CPUExecutionProvider'] # preprocessors # face align and parsing @@ -43,13 +47,12 @@ class StableDiffusionXLPuLIDPipeline: save_ext='png', device=self.device, ) - self.face_helper.face_parse = None self.face_helper.face_parse = init_parsing_model(model_name='bisenet', device=self.device) # clip-vit backbone - model, _, _ = create_model_and_transforms('EVA02-CLIP-L-14-336', 'eva_clip', force_custom_clip=True) - model = model.visual - self.clip_vision_model = model.to(self.device) + eva_precision = 'fp16' if self.dtype == torch.float16 or self.dtype == torch.bfloat16 else 'fp32' + eva_model, _, _ = create_model_and_transforms('EVA02-CLIP-L-14-336', 'eva_clip', force_custom_clip=True, precision=eva_precision, device=self.device) + self.clip_vision_model = eva_model.visual.to(dtype=self.dtype) eva_transform_mean = getattr(self.clip_vision_model, 'image_mean', OPENAI_DATASET_MEAN) eva_transform_std = getattr(self.clip_vision_model, 'image_std', OPENAI_DATASET_STD) if not isinstance(eva_transform_mean, (list, tuple)): @@ -60,13 +63,12 @@ class StableDiffusionXLPuLIDPipeline: self.eva_transform_std = eva_transform_std # antelopev2 - # snapshot_download('DIAMONIK7777/antelopev2', local_dir='models/antelopev2') local_dir = os.path.join(self.cache_dir, 'pulid', 'models', 'antelopev2') _loc = snapshot_download('DIAMONIK7777/antelopev2', local_dir=local_dir) self.app = FaceAnalysis( name='antelopev2', root=os.path.join(self.cache_dir, 'pulid'), - providers=['CUDAExecutionProvider', 'CPUExecutionProvider'], + providers=self.providers, ) self.app.prepare(ctx_id=0, det_size=(640, 640)) self.handler_ante = insightface.model_zoo.get_model(os.path.join(local_dir, 'glintr100.onnx')) @@ -89,8 +91,12 @@ class StableDiffusionXLPuLIDPipeline: self.log_sigmas = self.sigmas.log() self.sigma_data = 1.0 + # default scheduler if sampler is not None: self.sampler = sampler + else: + from modules.pulid import sampling + self.sampler = sampling.sample_dpmpp_sde @property def sigma_min(self): @@ -130,7 +136,7 @@ class StableDiffusionXLPuLIDPipeline: id_adapter_attn_procs[name] = IDAttnProcessor( hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, - ).to(unet.device) + ).to(unet.device, unet.dtype) else: id_adapter_attn_procs[name] = AttnProcessor() unet.set_attn_processor(id_adapter_attn_procs) @@ -144,7 +150,7 @@ class StableDiffusionXLPuLIDPipeline: module = k.split('.')[0] state_dict_dict.setdefault(module, {}) new_k = k[len(module) + 1 :] - state_dict_dict[module][new_k] = v + state_dict_dict[module][new_k] = v.to(self.dtype) for module in state_dict_dict: getattr(self, module).load_state_dict(state_dict_dict[module], strict=True) @@ -161,24 +167,17 @@ class StableDiffusionXLPuLIDPipeline: """ id_cond_list = [] id_vit_hidden_list = [] + self.face_helper.face_det.to(self.device) + self.clip_vision_model.to(self.device) for _ii, image in enumerate(image_list): self.face_helper.clean_all() image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) # get antelopev2 embedding face_info = self.app.get(image_bgr) if len(face_info) > 0: - face_info = sorted( - face_info, key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]) - )[ - -1 - ] # only use the maximum face + face_info = sorted(face_info, key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[-1] # only use the maximum face id_ante_embedding = face_info['embedding'] - self.debug_img_list.append( - image[ - int(face_info['bbox'][1]) : int(face_info['bbox'][3]), - int(face_info['bbox'][0]) : int(face_info['bbox'][2]), - ] - ) + self.debug_img_list.append(image[int(face_info['bbox'][1]) : int(face_info['bbox'][3]), int(face_info['bbox'][0]) : int(face_info['bbox'][2])]) else: id_ante_embedding = None @@ -210,13 +209,9 @@ class StableDiffusionXLPuLIDPipeline: self.debug_img_list.append(tensor2img(face_features_image, rgb2bgr=False)) # transform img before sending to eva-clip-vit - face_features_image = resize( - face_features_image, self.clip_vision_model.image_size, InterpolationMode.BICUBIC - ) - face_features_image = normalize(face_features_image, self.eva_transform_mean, self.eva_transform_std) - id_cond_vit, id_vit_hidden = self.clip_vision_model( - face_features_image, return_all_features=False, return_hidden=True, shuffle=False - ) + face_features_image = resize(face_features_image, self.clip_vision_model.image_size, InterpolationMode.BICUBIC) + face_features_image = normalize(face_features_image, self.eva_transform_mean, self.eva_transform_std).to(self.dtype) + id_cond_vit, id_vit_hidden = self.clip_vision_model(face_features_image, return_all_features=False, return_hidden=True, shuffle=False) id_cond_vit_norm = torch.norm(id_cond_vit, 2, 1, True) id_cond_vit = torch.div(id_cond_vit, id_cond_vit_norm) @@ -225,19 +220,25 @@ class StableDiffusionXLPuLIDPipeline: id_cond_list.append(id_cond) id_vit_hidden_list.append(id_vit_hidden) - id_uncond = torch.zeros_like(id_cond_list[0]) + self.id_adapter.to(self.device) + id_uncond = torch.zeros_like(id_cond_list[0]).to(self.dtype) id_vit_hidden_uncond = [] for layer_idx in range(0, len(id_vit_hidden_list[0])): - id_vit_hidden_uncond.append(torch.zeros_like(id_vit_hidden_list[0][layer_idx])) + id_vit_hidden_uncond.append(torch.zeros_like(id_vit_hidden_list[0][layer_idx]).to(self.dtype)) - id_cond = torch.stack(id_cond_list, dim=1) + id_cond = torch.stack(id_cond_list, dim=1).to(self.dtype) id_vit_hidden = id_vit_hidden_list[0] for i in range(1, len(image_list)): for j, x in enumerate(id_vit_hidden_list[i]): - id_vit_hidden[j] = torch.cat([id_vit_hidden[j], x], dim=1) + id_vit_hidden[j] = torch.cat([id_vit_hidden[j], x], dim=1).to(self.dtype) id_embedding = self.id_adapter(id_cond, id_vit_hidden) uncond_id_embedding = self.id_adapter(id_uncond, id_vit_hidden_uncond) + if self.offload: + self.face_helper.face_det.to('cpu') + self.id_adapter.to('cpu') + self.clip_vision_model.to('cpu') + # return id_embedding return uncond_id_embedding, id_embedding @@ -314,6 +315,7 @@ class StableDiffusionXLPuLIDPipeline: id_embedding=None, uncond_id_embedding=None, id_scale: float=1.0, + output_type: str='pil', callback_on_step_end=None, ): self.step = 0 # pylint: disable=attribute-defined-outside-init @@ -370,24 +372,15 @@ class StableDiffusionXLPuLIDPipeline: mask_args = None latents = self.sampler(self.sample, noisy_latent, sigmas, extra_args=sampler_kwargs, disable=False, mask_args=mask_args) - latents = latents.to(dtype=self.pipe.vae.dtype, device=self.device) / self.pipe.vae.config.scaling_factor - images = self.pipe.vae.decode(latents).sample - images = self.pipe.image_processor.postprocess(images, output_type='pil') - - # Pixel space final mask - # if mask_image is not None: - # # TODO: Fix XYZ - # from PIL import Image - # mask_image = np.asarray(mask_image.convert("L")) - # mask_image = mask_image / mask_image.max() - # mask_image = mask_image.reshape(1,mask_image.shape[0],mask_image.shape[1],1) - # image = np.asarray(image).astype(mask_image.dtype) - # images = np.asarray(images).astype(mask_image.dtype) - # images = ((1 - mask_image) * image) + (mask_image * images) - # images = images[0].round().astype(np.uint8) - # images = [Image.fromarray(images)] - - return images + if output_type == 'latent': + images = self.pipe.image_processor.postprocess(latents, output_type='latent') + elif output_type == 'np': + images = self.pipe.image_processor.postprocess(latents, output_type='np') + else: + latents = latents.to(dtype=self.pipe.vae.dtype, device=self.device) / self.pipe.vae.config.scaling_factor + images = self.pipe.vae.decode(latents).sample + images = self.pipe.image_processor.postprocess(images, output_type='pil') + return StableDiffusionXLPipelineOutput(images) class StableDiffusionXLPuLIDPipelineImage(StableDiffusionXLPuLIDPipeline): diff --git a/modules/sd_models.py b/modules/sd_models.py index 4d3b2ad60..801c3705d 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -402,9 +402,11 @@ def apply_balanced_offload(sd_model): def apply_balanced_offload_to_module(pipe): if hasattr(pipe, "pipe"): apply_balanced_offload_to_module(pipe.pipe) - if not hasattr(pipe, "_internal_dict"): - return - for module_name in pipe._internal_dict.keys(): # pylint: disable=protected-access + if hasattr(pipe, "_internal_dict"): + keys = pipe._internal_dict.keys() # pylint: disable=protected-access + else: + keys = get_signature(shared.sd_model).keys() + for module_name in keys: # pylint: disable=protected-access module = getattr(pipe, module_name, None) if isinstance(module, torch.nn.Module): checkpoint_name = pipe.sd_checkpoint_info.name if getattr(pipe, "sd_checkpoint_info", None) is not None else None diff --git a/scripts/pulid_ext.py b/scripts/pulid_ext.py index d0d738b6a..6599fa2e8 100644 --- a/scripts/pulid_ext.py +++ b/scripts/pulid_ext.py @@ -1,5 +1,6 @@ import io import os +import time import contextlib import gradio as gr import numpy as np @@ -18,6 +19,7 @@ class Script(scripts.Script): self.pulid = None self.cache = None self.mask_apply_overlay = shared.opts.mask_apply_overlay + self.preprocess = 0 super().__init__() self.register() # pulid is script with processing override so xyz doesnt execute @@ -88,15 +90,16 @@ class Script(scripts.Script): sampler = gr.Dropdown(label="Sampler", value='dpmpp_sde', choices=['dpmpp_2m', 'dpmpp_2m_sde', 'dpmpp_2s_ancestral', 'dpmpp_3m_sde', 'dpmpp_sde', 'euler', 'euler_ancestral']) ortho = gr.Dropdown(label="Ortho", choices=['off', 'v1', 'v2'], value='v2') with gr.Row(): - cache = gr.Checkbox(label='Keep model', value=False) + restore = gr.Checkbox(label='Restore pipe on end', value=False) + offload = gr.Checkbox(label='Offload face module', value=True) with gr.Row(): files = gr.File(label='Input images', file_count='multiple', file_types=['image'], type='file', interactive=True, height=100) with gr.Row(): gallery = gr.Gallery(show_label=False, value=[], visible=False, container=False, rows=1) files.change(fn=self.load_images, inputs=[files], outputs=[gallery]) - return [strength, zero, sampler, ortho, gallery, cache] + return [strength, zero, sampler, ortho, gallery, restore, offload] - def run(self, p: processing.StableDiffusionProcessing, strength: float = 0.8, zero: int = 20, sampler: str = 'dpmpp_sde', ortho: str = 'v2', gallery: list = [], cache: bool = False): # pylint: disable=arguments-differ, unused-argument + def run(self, p: processing.StableDiffusionProcessing, strength: float = 0.8, zero: int = 20, sampler: str = 'dpmpp_sde', ortho: str = 'v2', gallery: list = [], restore: bool = False, offload: bool = True): # pylint: disable=arguments-differ, unused-argument images = [] try: if len(gallery) == 0: @@ -135,13 +138,13 @@ class Script(scripts.Script): shared.log.warning('PuLID: batch size not supported') p.batch_size = 1 + self.mask_apply_overlay = shared.opts.mask_apply_overlay + shared.opts.data['mask_apply_overlay'] = False strength = getattr(p, 'pulid_strength', strength) zero = getattr(p, 'pulid_zero', zero) ortho = getattr(p, 'pulid_ortho', ortho) sampler = getattr(p, 'pulid_sampler', sampler) sampler_fn = getattr(self.pulid.sampling, f'sample_{sampler}', None) - self.mask_apply_overlay = shared.opts.mask_apply_overlay - shared.opts.data['mask_apply_overlay'] = False if sampler_fn is None: sampler_fn = self.pulid.sampling.sample_dpmpp_2m_sde @@ -153,6 +156,9 @@ class Script(scripts.Script): shared.sd_model = self.pulid.StableDiffusionXLPuLIDPipeline( pipe =shared.sd_model, device=devices.device, + dtype=devices.dtype, + providers=devices.onnx, + offload=offload, cache_dir=shared.opts.hfcache_dir, ) shared.sd_model.no_recurse = True @@ -166,13 +172,20 @@ class Script(scripts.Script): return None shared.sd_model.sampler = sampler_fn - shared.log.info(f'PuLID: class={shared.sd_model.__class__.__name__} strength={strength} zero={zero} ortho={ortho} sampler={sampler_fn} images={[i.shape for i in images]}') + shared.log.info(f'PuLID: class={shared.sd_model.__class__.__name__} strength={strength} zero={zero} ortho={ortho} sampler={sampler_fn} images={[i.shape for i in images]} offload={offload}') self.pulid.attention.NUM_ZERO = zero self.pulid.attention.ORTHO = ortho == 'v1' self.pulid.attention.ORTHO_v2 = ortho == 'v2' images = [self.pulid.resize(image, 1024) for image in images] shared.sd_model.debug_img_list = [] + + # get id embedding used for attention + t0 = time.time() uncond_id_embedding, id_embedding = shared.sd_model.get_id_embedding(images) + if offload: + devices.torch_gc() + t1 = time.time() + self.preprocess = t1-t0 p.seed = processing_helpers.get_fixed_seed(p.seed) if direct: # run pipeline directly @@ -211,20 +224,18 @@ class Script(scripts.Script): return processed def after(self, p: processing.StableDiffusionProcessing, processed: processing.Processed, *args): # pylint: disable=unused-argument - _strength, _zero, _sampler, _ortho, _gallery, cache = args + _strength, _zero, _sampler, _ortho, _gallery, restore, _offload = args if hasattr(shared.sd_model, 'pipe') and shared.sd_model_type == "sdxl": shared.opts.data['mask_apply_overlay'] = self.mask_apply_overlay - cache = getattr(p, 'pulid_cache', cache) - if cache: - shared.log.debug(f'PuLID cache: class={shared.sd_model.__class__.__name__}') - return processed - if hasattr(shared.sd_model, 'app'): - shared.sd_model.app = None - shared.sd_model.ip_adapter = None - shared.sd_model.face_helper = None - shared.sd_model.clip_vision_model = None - shared.sd_model.handler_ante = None - shared.sd_model = shared.sd_model.pipe - devices.torch_gc(force=True) - shared.log.debug(f'PuLID restore: class={shared.sd_model.__class__.__name__}') + restore = getattr(p, 'pulid_restore', restore) + if restore: + if hasattr(shared.sd_model, 'app'): + shared.sd_model.app = None + shared.sd_model.ip_adapter = None + shared.sd_model.face_helper = None + shared.sd_model.clip_vision_model = None + shared.sd_model.handler_ante = None + shared.sd_model = shared.sd_model.pipe + devices.torch_gc(force=True) + shared.log.debug(f'PuLID complete: class={shared.sd_model.__class__.__name__} preprocess={self.preprocess:.2f} pipe={"restore" if restore else "cache"}') return processed