pullid support sdpa add both v1.0 and v1.1 models

Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2024-11-10 09:42:05 -05:00
parent 6510a13904
commit b75567eefe
5 changed files with 113 additions and 24 deletions
+1 -1
View File
@@ -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.
+2 -2
View File
@@ -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}')
+62 -6
View File
@@ -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)
+28 -10
View File
@@ -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':
+20 -5
View File
@@ -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)