345 lines
14 KiB
Python
345 lines
14 KiB
Python
import copy
|
|
import numpy as np
|
|
import torch
|
|
import comfy
|
|
import comfy.sampler_helpers
|
|
import comfy.patcher_extension
|
|
import comfy.model_patcher
|
|
from comfy.samplers import Sampler, CFGGuider, finalize_default_conds, get_area_and_mult, cond_cat
|
|
from tqdm.auto import trange
|
|
from .model_base import BasePiFlow
|
|
from .piflow_policies.base import BasePolicy
|
|
|
|
|
|
def deepcopy_no_tensors(x):
|
|
memo = {}
|
|
|
|
# Pre-register all tensors we can reach so deepcopy will reuse them.
|
|
def register_tensors(obj):
|
|
if torch.is_tensor(obj):
|
|
memo[id(obj)] = obj
|
|
return
|
|
if isinstance(obj, dict):
|
|
for k, v in obj.items():
|
|
register_tensors(k)
|
|
register_tensors(v)
|
|
elif isinstance(obj, (list, tuple, set, frozenset)):
|
|
for v in obj:
|
|
register_tensors(v)
|
|
|
|
register_tensors(x)
|
|
return copy.deepcopy(x, memo)
|
|
|
|
|
|
def calc_cond_batch(model: BasePiFlow, conds: list[list[dict]], x_in: torch.Tensor, timestep, model_options: dict[str]):
|
|
handler: comfy.context_windows.ContextHandlerABC = model_options.get("context_handler", None)
|
|
if handler is None or not handler.should_use_context(model, conds, x_in, timestep, model_options):
|
|
return _calc_cond_batch_outer(model, conds, x_in, timestep, model_options)
|
|
return handler.execute(_calc_cond_batch_outer, model, conds, x_in, timestep, model_options)
|
|
|
|
|
|
def _calc_cond_batch_outer(model: BasePiFlow, conds: list[list[dict]], x_in: torch.Tensor, timestep, model_options):
|
|
executor = comfy.patcher_extension.WrapperExecutor.new_executor(
|
|
_calc_cond_batch,
|
|
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.CALC_COND_BATCH, model_options, is_model_options=True)
|
|
)
|
|
return executor.execute(model, conds, x_in, timestep, model_options)
|
|
|
|
|
|
class CompositePolicy(BasePolicy):
|
|
def __init__(self, entries: list[tuple[BasePolicy, list[int] | None, torch.Tensor]]):
|
|
self.entries = entries
|
|
|
|
@staticmethod
|
|
def _narrow_area(tensor: torch.Tensor, area: list[int] | None):
|
|
if area is None:
|
|
return tensor
|
|
|
|
dims = len(area) // 2
|
|
for i in range(dims):
|
|
tensor_dim = i + 2
|
|
if tensor.ndim > tensor_dim and tensor.shape[tensor_dim] > 1:
|
|
tensor = tensor.narrow(tensor_dim, area[i + dims], area[i])
|
|
return tensor
|
|
|
|
def pi(self, x_t, sigma_t):
|
|
out = torch.zeros_like(x_t)
|
|
counts = torch.ones_like(x_t) * 1e-37
|
|
|
|
for policy, area, mult in self.entries:
|
|
if area is None:
|
|
out += policy.pi(x_t, sigma_t) * mult
|
|
counts += mult
|
|
continue
|
|
|
|
out_view = out
|
|
counts_view = counts
|
|
dims = len(area) // 2
|
|
for i in range(dims):
|
|
out_view = out_view.narrow(i + 2, area[i + dims], area[i])
|
|
counts_view = counts_view.narrow(i + 2, area[i + dims], area[i])
|
|
|
|
x_view = self._narrow_area(x_t, area)
|
|
sigma_view = self._narrow_area(sigma_t, area)
|
|
out_view += policy.pi(x_view, sigma_view) * mult
|
|
counts_view += mult
|
|
|
|
return out / counts
|
|
|
|
def detach(self):
|
|
return CompositePolicy([(policy.detach(), area, mult) for policy, area, mult in self.entries])
|
|
|
|
def temperature_(self, temperature):
|
|
for policy, _, _ in self.entries:
|
|
if hasattr(policy, "temperature_"):
|
|
policy.temperature_(temperature)
|
|
return self
|
|
|
|
|
|
def _calc_cond_batch(model: BasePiFlow, conds: list[list[dict]], x_in: torch.Tensor, timestep, model_options):
|
|
out_conds = [[] for _ in conds]
|
|
# separate conds by matching hooks
|
|
hooked_to_run: dict[comfy.hooks.HookGroup, list[tuple[tuple, int]]] = {}
|
|
default_conds = []
|
|
has_default_conds = False
|
|
|
|
for i in range(len(conds)):
|
|
cond = conds[i]
|
|
default_c = []
|
|
if cond is not None:
|
|
for x in cond:
|
|
if 'default' in x:
|
|
default_c.append(x)
|
|
has_default_conds = True
|
|
continue
|
|
p = get_area_and_mult(x, x_in, timestep)
|
|
if p is None:
|
|
continue
|
|
if p.hooks is not None:
|
|
model.current_patcher.prepare_hook_patches_current_keyframe(timestep, p.hooks, model_options)
|
|
hooked_to_run.setdefault(p.hooks, list())
|
|
hooked_to_run[p.hooks] += [(p, i)]
|
|
default_conds.append(default_c)
|
|
|
|
if has_default_conds:
|
|
finalize_default_conds(model, hooked_to_run, default_conds, x_in, timestep, model_options)
|
|
|
|
model.current_patcher.prepare_state(timestep, model_options)
|
|
|
|
# run every hooked_to_run separately
|
|
for hooks, to_run in hooked_to_run.items():
|
|
while len(to_run) > 0:
|
|
p, cond_index = to_run.pop()
|
|
input_x = p.input_x
|
|
c = cond_cat([p.conditioning])
|
|
timestep_ = timestep
|
|
|
|
transformer_options = model.current_patcher.apply_hooks(hooks=hooks)
|
|
if 'transformer_options' in model_options:
|
|
transformer_options = comfy.patcher_extension.merge_nested_dicts(
|
|
transformer_options, model_options['transformer_options'], copy_dict1=False)
|
|
|
|
if p.patches is not None:
|
|
transformer_options["patches"] = comfy.patcher_extension.merge_nested_dicts(
|
|
transformer_options.get("patches", {}),
|
|
p.patches
|
|
)
|
|
|
|
transformer_options["cond_or_uncond"] = [cond_index]
|
|
transformer_options["uuids"] = [p.uuid]
|
|
transformer_options["sigmas"] = timestep
|
|
|
|
c['transformer_options'] = transformer_options
|
|
|
|
if p.control is not None:
|
|
c['control'] = p.control.get_control(input_x, timestep_, c, 1, transformer_options)
|
|
|
|
if 'model_function_wrapper' in model_options:
|
|
output = model_options['model_function_wrapper'](model.apply_model, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": [cond_index]})
|
|
else:
|
|
output = model.apply_model(input_x, timestep_, **c)
|
|
|
|
out_conds[cond_index].append((output, p.area, p.mult))
|
|
|
|
policies = []
|
|
for entries in out_conds:
|
|
if len(entries) == 0:
|
|
raise RuntimeError("No active pi-Flow conditioning was available for this sampling step.")
|
|
if len(entries) == 1 and entries[0][1] is None and not torch.any(entries[0][2] == 0).item():
|
|
policies.append(entries[0][0])
|
|
else:
|
|
policies.append(CompositePolicy(entries))
|
|
|
|
if len(policies) == 1:
|
|
return policies[0]
|
|
return policies
|
|
|
|
|
|
def sampling_function(model, x, timestep, cond, model_options={}, seed=None):
|
|
return calc_cond_batch(model, [cond], x, timestep, model_options)
|
|
|
|
|
|
class _PolicySampler(CFGGuider):
|
|
|
|
def predict_noise(self, x, timestep, model_options={}, seed=None):
|
|
return sampling_function(
|
|
self.inner_model, x, timestep, self.conds.get("positive", None), model_options=model_options, seed=seed)
|
|
|
|
|
|
class _PiFlowSampler(Sampler):
|
|
|
|
def __init__(
|
|
self, model_sampling, h=0.0, substeps=128, gm_temperature=1.0):
|
|
self.model_sampling = model_sampling
|
|
self.h = h
|
|
self.substeps = substeps
|
|
self.gm_temperature = gm_temperature
|
|
|
|
def calculate_sigmas_dst(self, sigmas, eps=1e-6):
|
|
sigmas_src = sigmas[:-1]
|
|
sigmas_to = sigmas[1:]
|
|
alphas_src = 1 - sigmas_src
|
|
alphas_to = 1 - sigmas_to
|
|
if self.h <= 0.0:
|
|
m = torch.ones_like(sigmas_src)
|
|
else:
|
|
h2 = self.h * self.h
|
|
m = (sigmas_to * alphas_src / (sigmas_src * alphas_to).clamp(min=eps)) ** h2
|
|
|
|
sigmas_to_mul_m = sigmas_to * m
|
|
sigmas_dst = sigmas_to_mul_m / (alphas_to + sigmas_to_mul_m).clamp(min=eps)
|
|
|
|
return sigmas_dst, m
|
|
|
|
def sample(self, model_wrap, sigmas, extra_args, callback, noise,
|
|
latent_image=None, denoise_mask=None, disable_pbar=False):
|
|
noise = self.model_sampling.unpatchify(noise)
|
|
if latent_image is not None:
|
|
latent_image = self.model_sampling.unpatchify(latent_image)
|
|
if denoise_mask is not None:
|
|
denoise_mask = self.model_sampling.unpatchify(denoise_mask)
|
|
|
|
x_t_src = self.model_sampling.noise_scaling(
|
|
sigmas[0], noise, latent_image)
|
|
|
|
sigmas_dst, m_vals = self.calculate_sigmas_dst(sigmas)
|
|
|
|
seed = extra_args.get("seed", None)
|
|
if seed is not None:
|
|
seed = (seed + 1) & 0xffffffffffffffff
|
|
generator = torch.Generator(device=x_t_src.device).manual_seed(seed)
|
|
else:
|
|
generator = None
|
|
|
|
num_batches = x_t_src.size(0)
|
|
ndim = x_t_src.dim()
|
|
|
|
# pi-Flow sampling loop
|
|
nfe = len(sigmas) - 1
|
|
|
|
for step_id in trange(nfe, disable=disable_pbar):
|
|
sigma_t_src = sigmas[step_id].expand(num_batches).reshape(num_batches, *((ndim - 1) * [1]))
|
|
sigma_t_dst = sigmas_dst[step_id].expand(num_batches).reshape(num_batches, *((ndim - 1) * [1]))
|
|
t_src = self.model_sampling.timestep(sigma_t_src.flatten())
|
|
raw_t_src = self.model_sampling.unwarp_t(sigma_t_src.flatten())
|
|
raw_t_dst = self.model_sampling.unwarp_t(sigma_t_dst.flatten())
|
|
sigma_t_to = sigmas[step_id + 1]
|
|
alpha_t_to = 1 - sigma_t_to
|
|
m = m_vals[step_id]
|
|
|
|
policy = model_wrap(x_t_src, t_src, **extra_args)
|
|
if hasattr(policy, "temperature_") and step_id != nfe - 1:
|
|
policy.temperature_(self.gm_temperature)
|
|
|
|
x_t_dst, _, _ = policy.integrate(
|
|
x_t_src, sigma_t_src, raw_t_src, raw_t_dst,
|
|
self.model_sampling, total_substeps=self.substeps,
|
|
denoise_mask=denoise_mask, latent_image=latent_image)
|
|
|
|
noise = torch.randn(
|
|
x_t_dst.size(), dtype=x_t_dst.dtype, device=x_t_dst.device, generator=generator)
|
|
|
|
x_t_to = (alpha_t_to + sigma_t_to * m) * x_t_dst + sigma_t_to * (1 - m.square()).clamp(
|
|
min=0).sqrt() * noise
|
|
|
|
if callback is not None:
|
|
u = (x_t_src - x_t_to) / (sigma_t_src - sigma_t_to).clamp(min=1e-6)
|
|
x_t_0 = x_t_to - u * sigma_t_to
|
|
callback( # for preview
|
|
step_id,
|
|
self.model_sampling.patchify(x_t_0),
|
|
self.model_sampling.patchify(x_t_src),
|
|
nfe)
|
|
|
|
x_t_src = x_t_to
|
|
|
|
samples = self.model_sampling.inverse_noise_scaling(sigmas[-1], x_t_src)
|
|
|
|
samples = self.model_sampling.patchify(samples)
|
|
return samples
|
|
|
|
|
|
class PiFlowSampler:
|
|
def __init__(
|
|
self, model, steps, device, h=0.0, substeps=128, final_step_size_scale=1.0, denoise=1.0, model_options={}):
|
|
self.model = model
|
|
self.device = device
|
|
self.h = h
|
|
self.substeps = substeps
|
|
self.set_steps(steps, final_step_size_scale, denoise)
|
|
self.denoise = denoise
|
|
self.model_options = model_options
|
|
|
|
def set_steps(self, steps, final_step_size_scale=1.0, denoise=1.0):
|
|
self.steps = steps
|
|
if denoise <= 0.0:
|
|
self.sigmas = torch.FloatTensor([])
|
|
else:
|
|
model_sampling = self.model.get_model_object("model_sampling")
|
|
raw_timesteps = torch.from_numpy(np.linspace(
|
|
1, (final_step_size_scale - 1) / (steps + final_step_size_scale - 1),
|
|
steps, dtype=np.float32, endpoint=False)).clamp(min=0) * min(denoise, 1.0)
|
|
self.sigmas = torch.cat(
|
|
[model_sampling.warp_t(raw_timesteps), torch.FloatTensor([0.0])], dim=0)
|
|
|
|
def sample(
|
|
self, noise, conditioning, temperature, latent_image=None,
|
|
denoise_mask=None, callback=None, disable_pbar=False, seed=None):
|
|
policy_sampler = _PolicySampler(self.model)
|
|
|
|
model_sampling = self.model.get_model_object("model_sampling")
|
|
conditioning = deepcopy_no_tensors(conditioning)
|
|
for conditioning_single in conditioning:
|
|
conditioning_dict = conditioning_single[1]
|
|
if 'reference_latents' in conditioning_dict:
|
|
conditioning_dict['reference_latents'] = [
|
|
model_sampling.unpatchify(x) for x in conditioning_dict['reference_latents']]
|
|
|
|
policy_sampler.set_conds(conditioning, [])
|
|
policy_sampler.set_cfg(1.0)
|
|
sampler = _PiFlowSampler(
|
|
self.model.get_model_object("model_sampling"),
|
|
h=self.h, substeps=self.substeps, gm_temperature=temperature)
|
|
return policy_sampler.sample(noise, latent_image, sampler, self.sigmas, denoise_mask, callback, disable_pbar, seed)
|
|
|
|
|
|
def sample(
|
|
model, noise, steps, substeps, final_step_size_scale, diffusion_coefficient,
|
|
gm_temperature, manual_gm_temperature,
|
|
conditioning, latent_image, denoise=1.0,
|
|
noise_mask=None, callback=None, disable_pbar=None, seed=None):
|
|
sampler = PiFlowSampler(
|
|
model, steps, device=model.load_device, h=diffusion_coefficient, substeps=substeps,
|
|
final_step_size_scale=final_step_size_scale, denoise=denoise)
|
|
if gm_temperature == 'auto':
|
|
temperature = min(max(0.1 * (steps - 1), 0), 1)
|
|
elif gm_temperature == 'manual':
|
|
temperature = manual_gm_temperature
|
|
else:
|
|
temperature = 1.0
|
|
samples = sampler.sample(
|
|
noise, conditioning, temperature=temperature, latent_image=latent_image,
|
|
denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
|
|
samples = samples.to(comfy.model_management.intermediate_device())
|
|
return samples
|