104 lines
2.5 KiB
Python
104 lines
2.5 KiB
Python
from functools import wraps
|
|
from math import pi, sin
|
|
from typing import Callable
|
|
|
|
import torch
|
|
|
|
import latent_preview
|
|
|
|
RECENTER: float = 0.0
|
|
NORMALIZATION: float = 0.0
|
|
|
|
|
|
def disable_all():
|
|
global RECENTER
|
|
RECENTER = 0.0
|
|
global NORMALIZATION
|
|
NORMALIZATION = 0.0
|
|
|
|
|
|
_prepare_callback: Callable = latent_preview.prepare_callback
|
|
|
|
|
|
@wraps(_prepare_callback)
|
|
def hijack_prepare(*args, **kwargs):
|
|
original_callback: Callable = _prepare_callback(*args, **kwargs)
|
|
|
|
@torch.inference_mode()
|
|
@wraps(original_callback)
|
|
def hijack_callback(step, x0, x, total_steps):
|
|
original_callback(step, x0, x, total_steps)
|
|
|
|
ratio: float = step / total_steps
|
|
strength: float = 1.0 - sin(ratio * pi / 2.0)
|
|
latent: torch.Tensor = x0.detach().clone()
|
|
|
|
batches: int = latent.shape[0]
|
|
channels: int = latent.shape[1]
|
|
|
|
for b in range(batches):
|
|
_std: float = latent[b].std()
|
|
|
|
for c in range(channels):
|
|
bias: float = latent[b][c].std() / _std * strength
|
|
x0[b][c] += (0.0 - latent[b][c].mean()) * bias * RECENTER
|
|
|
|
magnitude = float(latent[b][c].max() - latent[b][c].min())
|
|
factor = 1.0 / 0.13025
|
|
scale = max(factor / magnitude - 1.0, 0.0)
|
|
x0[b][c] *= scale * NORMALIZATION * strength + 1.0
|
|
|
|
return hijack_callback
|
|
|
|
|
|
latent_preview.prepare_callback = hijack_prepare
|
|
|
|
|
|
class DiffusionCG:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"latent": ("LATENT",),
|
|
"recenter": (
|
|
"FLOAT",
|
|
{"default": 0.00, "min": 0.00, "max": 1.00, "step": 0.05},
|
|
),
|
|
"normalization": (
|
|
"FLOAT",
|
|
{"default": 0.00, "min": 0.00, "max": 1.00, "step": 0.05},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "hook"
|
|
CATEGORY = "sampling"
|
|
|
|
def hook(self, latent, recenter: float, normalization: float):
|
|
global RECENTER
|
|
RECENTER = recenter
|
|
global NORMALIZATION
|
|
NORMALIZATION = normalization
|
|
return (latent,)
|
|
|
|
|
|
class DisableCG:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"latent": ("LATENT",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "hook"
|
|
CATEGORY = "sampling"
|
|
|
|
def hook(self, latent):
|
|
disable_all()
|
|
return (latent,)
|