diff --git a/CHANGELOG.md b/CHANGELOG.md index 417165276..9affd53db 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Change Log for SD.Next -## Update for 2024-11-08 +## Update for 2024-11-10 Smaller release just few days after the last one, but with some important fixes and improvements. This release can be considered an LTS release before we kick off the next round of major updates. diff --git a/modules/errors.py b/modules/errors.py index c4d66c351..527884cf1 100644 --- a/modules/errors.py +++ b/modules/errors.py @@ -59,7 +59,7 @@ def exception(suppress=[]): console.print_exception(show_locals=False, max_frames=16, extra_lines=2, suppress=suppress, theme="ansi_dark", word_wrap=False, width=min([console.width, 200])) -def profile(profiler, msg: str): +def profile(profiler, msg: str, n: int = 5): profiler.disable() import io import pstats @@ -83,7 +83,7 @@ def profile(profiler, msg: str): and 'rich' not in x and x.strip() != '' ] - txt = '\n'.join(lines[:min(5, len(lines))]) + txt = '\n'.join(lines[:min(n, len(lines))]) log.debug(f'Profile {msg}: {txt}') diff --git a/modules/pulid/encoders_transformer.py b/modules/pulid/encoders_transformer.py index d1ecef2c6..ae245044b 100644 --- a/modules/pulid/encoders_transformer.py +++ b/modules/pulid/encoders_transformer.py @@ -186,23 +186,79 @@ class IDFormer(nn.Module): ) def forward(self, x, y): - latents = self.latents.repeat(x.size(0), 1, 1) - num_duotu = x.shape[1] if x.ndim == 3 else 1 - x = self.id_embedding_mapping(x) x = x.reshape(-1, self.num_id_token * num_duotu, self.dim) - latents = torch.cat((latents, x), dim=1) - for i in range(5): vit_feature = getattr(self, f'mapping_{i}')(y[i]) ctx_feature = torch.cat((x, vit_feature), dim=1) for attn, ff in self.layers[i * self.depth: (i + 1) * self.depth]: latents = attn(ctx_feature, latents) + latents latents = ff(latents) + latents - latents = latents[:, :self.num_queries] latents = latents @ self.proj_out return latents + + +class IDEncoder(nn.Module): + def __init__(self, width=1280, context_dim=2048, num_token=5): + super().__init__() + self.num_token = num_token + self.context_dim = context_dim + h1 = min((context_dim * num_token) // 4, 1024) + h2 = min((context_dim * num_token) // 2, 1024) + self.body = nn.Sequential( + nn.Linear(width, h1), + nn.LayerNorm(h1), + nn.LeakyReLU(), + nn.Linear(h1, h2), + nn.LayerNorm(h2), + nn.LeakyReLU(), + nn.Linear(h2, context_dim * num_token), + ) + + for i in range(5): + setattr( + self, + f'mapping_{i}', + nn.Sequential( + nn.Linear(1024, 1024), + nn.LayerNorm(1024), + nn.LeakyReLU(), + nn.Linear(1024, 1024), + nn.LayerNorm(1024), + nn.LeakyReLU(), + nn.Linear(1024, context_dim), + ), + ) + + setattr( + self, + f'mapping_patch_{i}', + nn.Sequential( + nn.Linear(1024, 1024), + nn.LayerNorm(1024), + nn.LeakyReLU(), + nn.Linear(1024, 1024), + nn.LayerNorm(1024), + nn.LeakyReLU(), + nn.Linear(1024, context_dim), + ), + ) + + def forward(self, x, y): + # x shape [N, C] + x = self.body(x) + x = x.reshape(-1, self.num_token, self.context_dim) + + hidden_states = () + for i, emb in enumerate(y): + hidden_state = getattr(self, f'mapping_{i}')(emb[:, :1]) + getattr(self, f'mapping_patch_{i}')( + emb[:, 1:] + ).mean(dim=1, keepdim=True) + hidden_states += (hidden_state,) + hidden_states = torch.cat(hidden_states, dim=1) + + return torch.cat([x, hidden_states], dim=1) diff --git a/modules/pulid/pulid_sdxl.py b/modules/pulid/pulid_sdxl.py index 8bd28dbe2..d2a761654 100644 --- a/modules/pulid/pulid_sdxl.py +++ b/modules/pulid/pulid_sdxl.py @@ -19,23 +19,28 @@ from insightface.app import FaceAnalysis from eva_clip import create_model_and_transforms from eva_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD -from encoders_transformer import IDFormer -from attention_processor import AttnProcessor2_0 as AttnProcessor -from attention_processor import IDAttnProcessor2_0 as IDAttnProcessor +from encoders_transformer import IDFormer, IDEncoder class StableDiffusionXLPuLIDPipeline: - def __init__(self, pipe: StableDiffusionXLPipeline, device: torch.device, dtype: torch.dtype=None, providers: list=None, offload: bool=True, 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, sdp: bool=True, version: str='v1.1'): 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.sdp = sdp + self.version = version + self.folder = 'models--ToTheBeginning--PuLID' + self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config) - self.id_adapter = IDFormer().to(self.device, self.dtype) + if self.version == 'v1.1': + self.id_adapter = IDFormer().to(self.device, self.dtype) + else: + self.id_adapter = IDEncoder().to(self.device, self.dtype) self.providers = providers or ['CUDAExecutionProvider', 'CPUExecutionProvider'] + self.hack_unet_attn_layers(self.pipe.unet) # preprocessors # face align and parsing @@ -63,11 +68,11 @@ class StableDiffusionXLPuLIDPipeline: self.eva_transform_std = eva_transform_std # antelopev2 - local_dir = os.path.join(self.cache_dir, 'pulid', 'models', 'antelopev2') + local_dir = os.path.join(self.cache_dir, self.folder, 'models', 'antelopev2') _loc = snapshot_download('DIAMONIK7777/antelopev2', local_dir=local_dir) self.app = FaceAnalysis( name='antelopev2', - root=os.path.join(self.cache_dir, 'pulid'), + root=os.path.join(self.cache_dir, self.folder), providers=self.providers, ) self.app.prepare(ctx_id=0, det_size=(640, 640)) @@ -119,6 +124,12 @@ class StableDiffusionXLPuLIDPipeline: return torch.cat([sigmas, sigmas.new_zeros([1])]) def hack_unet_attn_layers(self, unet): + if self.sdp: + from attention_processor import AttnProcessor2_0 as AttnProcessor + from attention_processor import IDAttnProcessor2_0 as IDAttnProcessor + else: + from attention_processor import AttnProcessor + from attention_processor import IDAttnProcessor id_adapter_attn_procs = {} for name, _ in unet.attn_processors.items(): cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim @@ -143,8 +154,12 @@ class StableDiffusionXLPuLIDPipeline: self.id_adapter_attn_layers = nn.ModuleList(unet.attn_processors.values()) def load_pretrain(self): - ckpt_path = hf_hub_download('guozinan/PuLID', 'pulid_v1.1.safetensors', local_dir=os.path.join(self.cache_dir, 'pulid')) - state_dict = load_file(ckpt_path) + if self.version == 'v1.1': + ckpt_path = hf_hub_download('guozinan/PuLID', 'pulid_v1.1.safetensors', local_dir=os.path.join(self.cache_dir, self.folder)) + state_dict = load_file(ckpt_path) + else: + ckpt_path = hf_hub_download('guozinan/PuLID', 'pulid_v1.bin', local_dir=os.path.join(self.cache_dir, self.folder)) + state_dict = torch.load(ckpt_path, map_location="cpu") state_dict_dict = {} for k, v in state_dict.items(): module = k.split('.')[0] @@ -371,7 +386,10 @@ class StableDiffusionXLPuLIDPipeline: else: mask_args = None + # actual sampling loop latents = self.sampler(self.sample, noisy_latent, sigmas, extra_args=sampler_kwargs, disable=False, mask_args=mask_args) + + # process output if output_type == 'latent': images = self.pipe.image_processor.postprocess(latents, output_type='latent') elif output_type == 'np': diff --git a/scripts/pulid_ext.py b/scripts/pulid_ext.py index 6599fa2e8..181e954db 100644 --- a/scripts/pulid_ext.py +++ b/scripts/pulid_ext.py @@ -89,6 +89,8 @@ class Script(scripts.Script): with gr.Row(): 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(): + version = gr.Dropdown(label="Version", value='v1.1', choices=['v1.0', 'v1.1']) with gr.Row(): restore = gr.Checkbox(label='Restore pipe on end', value=False) offload = gr.Checkbox(label='Offload face module', value=True) @@ -97,9 +99,20 @@ class Script(scripts.Script): 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, restore, offload] + return [strength, zero, sampler, ortho, gallery, restore, offload, version] - 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 + 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, + version: str = 'v1.1' + ): # pylint: disable=arguments-differ, unused-argument images = [] try: if len(gallery) == 0: @@ -154,11 +167,13 @@ class Script(scripts.Script): ctx = contextlib.nullcontext() if debug else contextlib.redirect_stdout(stdout) with ctx: shared.sd_model = self.pulid.StableDiffusionXLPuLIDPipeline( - pipe =shared.sd_model, + pipe=shared.sd_model, device=devices.device, dtype=devices.dtype, providers=devices.onnx, offload=offload, + version=version, + sdp=shared.opts.cross_attention_optimization == "Scaled-Dot-Product", cache_dir=shared.opts.hfcache_dir, ) shared.sd_model.no_recurse = True @@ -172,7 +187,7 @@ 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]} offload={offload}') + shared.log.info(f'PuLID: class={shared.sd_model.__class__.__name__} version="{version}" 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' @@ -224,7 +239,7 @@ 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, restore, _offload = args + _strength, _zero, _sampler, _ortho, _gallery, restore, _offload, _version = args if hasattr(shared.sd_model, 'pipe') and shared.sd_model_type == "sdxl": shared.opts.data['mask_apply_overlay'] = self.mask_apply_overlay restore = getattr(p, 'pulid_restore', restore)