From 125319988984987801dc4b4ab1e5ed36e9b211c5 Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Fri, 10 Feb 2023 03:30:20 -0800 Subject: Working UniPC (for batch size 1) --- modules/models/diffusion/uni_pc/sampler.py | 85 ++++++++++++++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 modules/models/diffusion/uni_pc/sampler.py (limited to 'modules/models/diffusion/uni_pc/sampler.py') diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py new file mode 100644 index 00000000..7cccd8a2 --- /dev/null +++ b/modules/models/diffusion/uni_pc/sampler.py @@ -0,0 +1,85 @@ +"""SAMPLING ONLY.""" + +import torch + +from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC + +class UniPCSampler(object): + def __init__(self, model, **kwargs): + super().__init__() + self.model = model + to_torch = lambda x: x.clone().detach().to(torch.float32).to(model.device) + self.before_sample = None + self.after_sample = None + self.register_buffer('alphas_cumprod', to_torch(model.alphas_cumprod)) + + def register_buffer(self, name, attr): + if type(attr) == torch.Tensor: + if attr.device != torch.device("cuda"): + attr = attr.to(torch.device("cuda")) + setattr(self, name, attr) + + def set_hooks(self, before, after): + self.before_sample = before + self.after_sample = after + + @torch.no_grad() + def sample(self, + S, + batch_size, + shape, + conditioning=None, + callback=None, + normals_sequence=None, + img_callback=None, + quantize_x0=False, + eta=0., + mask=None, + x0=None, + temperature=1., + noise_dropout=0., + score_corrector=None, + corrector_kwargs=None, + verbose=True, + x_T=None, + log_every_t=100, + unconditional_guidance_scale=1., + unconditional_conditioning=None, + # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ... + **kwargs + ): + if conditioning is not None: + if isinstance(conditioning, dict): + cbs = conditioning[list(conditioning.keys())[0]].shape[0] + if cbs != batch_size: + print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}") + else: + if conditioning.shape[0] != batch_size: + print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}") + + # sampling + C, H, W = shape + size = (batch_size, C, H, W) + + device = self.model.betas.device + if x_T is None: + img = torch.randn(size, device=device) + else: + img = x_T + + ns = NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod) + + model_fn = model_wrapper( + lambda x, t, c: self.model.apply_model(x, t, c), + ns, + model_type="noise", + guidance_type="classifier-free", + #condition=conditioning, + #unconditional_condition=unconditional_conditioning, + guidance_scale=unconditional_guidance_scale, + ) + + uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, condition=conditioning, unconditional_condition=unconditional_conditioning, before_sample=self.before_sample, after_sample=self.after_sample) + x = uni_pc.sample(img, steps=S, skip_type="time_uniform", method="multistep", order=3, lower_order_final=True) + + return x.to(device), None -- cgit v1.2.3 From 21880eb9e57b884635a07d2360831b4186afddf4 Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Fri, 10 Feb 2023 04:47:08 -0800 Subject: Fix logspam and live previews --- modules/models/diffusion/uni_pc/sampler.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) (limited to 'modules/models/diffusion/uni_pc/sampler.py') diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py index 7cccd8a2..219e9862 100644 --- a/modules/models/diffusion/uni_pc/sampler.py +++ b/modules/models/diffusion/uni_pc/sampler.py @@ -19,9 +19,10 @@ class UniPCSampler(object): attr = attr.to(torch.device("cuda")) setattr(self, name, attr) - def set_hooks(self, before, after): - self.before_sample = before - self.after_sample = after + def set_hooks(self, before_sample, after_sample, after_update): + self.before_sample = before_sample + self.after_sample = after_sample + self.after_update = after_update @torch.no_grad() def sample(self, @@ -50,9 +51,17 @@ class UniPCSampler(object): ): if conditioning is not None: if isinstance(conditioning, dict): - cbs = conditioning[list(conditioning.keys())[0]].shape[0] + ctmp = conditioning[list(conditioning.keys())[0]] + while isinstance(ctmp, list): ctmp = ctmp[0] + cbs = ctmp.shape[0] if cbs != batch_size: print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}") + + elif isinstance(conditioning, list): + for ctmp in conditioning: + if ctmp.shape[0] != batch_size: + print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}") + else: if conditioning.shape[0] != batch_size: print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}") @@ -60,6 +69,7 @@ class UniPCSampler(object): # sampling C, H, W = shape size = (batch_size, C, H, W) + print(f'Data shape for UniPC sampling is {size}, eta {eta}') device = self.model.betas.device if x_T is None: @@ -79,7 +89,7 @@ class UniPCSampler(object): guidance_scale=unconditional_guidance_scale, ) - uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, condition=conditioning, unconditional_condition=unconditional_conditioning, before_sample=self.before_sample, after_sample=self.after_sample) + uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, condition=conditioning, unconditional_condition=unconditional_conditioning, before_sample=self.before_sample, after_sample=self.after_sample, after_update=self.after_update) x = uni_pc.sample(img, steps=S, skip_type="time_uniform", method="multistep", order=3, lower_order_final=True) return x.to(device), None -- cgit v1.2.3 From 79ffb9453f8eddbdd4e316b9d9c75812b0eea4e1 Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Fri, 10 Feb 2023 05:27:05 -0800 Subject: Add UniPC sampler settings --- modules/models/diffusion/uni_pc/sampler.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) (limited to 'modules/models/diffusion/uni_pc/sampler.py') diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py index 219e9862..e66a21e3 100644 --- a/modules/models/diffusion/uni_pc/sampler.py +++ b/modules/models/diffusion/uni_pc/sampler.py @@ -3,6 +3,7 @@ import torch from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC +from modules import shared class UniPCSampler(object): def __init__(self, model, **kwargs): @@ -89,7 +90,7 @@ class UniPCSampler(object): guidance_scale=unconditional_guidance_scale, ) - uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, condition=conditioning, unconditional_condition=unconditional_conditioning, before_sample=self.before_sample, after_sample=self.after_sample, after_update=self.after_update) - x = uni_pc.sample(img, steps=S, skip_type="time_uniform", method="multistep", order=3, lower_order_final=True) + uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=shared.opts.uni_pc_thresholding, variant=shared.opts.uni_pc_variant, condition=conditioning, unconditional_condition=unconditional_conditioning, before_sample=self.before_sample, after_sample=self.after_sample, after_update=self.after_update) + x = uni_pc.sample(img, steps=S, skip_type=shared.opts.uni_pc_skip_type, method="multistep", order=shared.opts.uni_pc_order, lower_order_final=shared.opts.uni_pc_lower_order_final) return x.to(device), None -- cgit v1.2.3 From fb274229b2c5c1a89dac0b3da28c08c92d71fd95 Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Fri, 10 Feb 2023 14:30:35 -0800 Subject: bug fix --- modules/models/diffusion/uni_pc/sampler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'modules/models/diffusion/uni_pc/sampler.py') diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py index e66a21e3..0bef6eed 100644 --- a/modules/models/diffusion/uni_pc/sampler.py +++ b/modules/models/diffusion/uni_pc/sampler.py @@ -70,7 +70,7 @@ class UniPCSampler(object): # sampling C, H, W = shape size = (batch_size, C, H, W) - print(f'Data shape for UniPC sampling is {size}, eta {eta}') + print(f'Data shape for UniPC sampling is {size}') device = self.model.betas.device if x_T is None: -- cgit v1.2.3 From 716a69237cefb385f71105dbbf50e92d664e0f42 Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Sat, 11 Feb 2023 06:18:34 -0800 Subject: support SD2.X models --- modules/models/diffusion/uni_pc/sampler.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) (limited to 'modules/models/diffusion/uni_pc/sampler.py') diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py index 0bef6eed..708a9b2b 100644 --- a/modules/models/diffusion/uni_pc/sampler.py +++ b/modules/models/diffusion/uni_pc/sampler.py @@ -80,10 +80,13 @@ class UniPCSampler(object): ns = NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod) + # SD 1.X is "noise", SD 2.X is "v" + model_type = "v" if self.model.parameterization == "v" else "noise" + model_fn = model_wrapper( lambda x, t, c: self.model.apply_model(x, t, c), ns, - model_type="noise", + model_type=model_type, guidance_type="classifier-free", #condition=conditioning, #unconditional_condition=unconditional_conditioning, -- cgit v1.2.3 From 5fef67f6ee949a61826a3a043ea8610fd89fc371 Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Fri, 10 Mar 2023 19:56:14 -0500 Subject: Requested changes --- modules/models/diffusion/uni_pc/sampler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'modules/models/diffusion/uni_pc/sampler.py') diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py index 708a9b2b..6bb3bb21 100644 --- a/modules/models/diffusion/uni_pc/sampler.py +++ b/modules/models/diffusion/uni_pc/sampler.py @@ -93,7 +93,7 @@ class UniPCSampler(object): guidance_scale=unconditional_guidance_scale, ) - uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=shared.opts.uni_pc_thresholding, variant=shared.opts.uni_pc_variant, condition=conditioning, unconditional_condition=unconditional_conditioning, before_sample=self.before_sample, after_sample=self.after_sample, after_update=self.after_update) + uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, variant=shared.opts.uni_pc_variant, condition=conditioning, unconditional_condition=unconditional_conditioning, before_sample=self.before_sample, after_sample=self.after_sample, after_update=self.after_update) x = uni_pc.sample(img, steps=S, skip_type=shared.opts.uni_pc_skip_type, method="multistep", order=shared.opts.uni_pc_order, lower_order_final=shared.opts.uni_pc_lower_order_final) return x.to(device), None -- cgit v1.2.3