74 lines
2.0 KiB
Python
74 lines
2.0 KiB
Python
import torch
|
|
|
|
from comfy.model_management import get_torch_device
|
|
from comfy_extras.nodes_model_advanced import ModelSamplingDiscreteDistilled, LCM
|
|
|
|
|
|
def is_lcm(coreml_model):
|
|
return "timestep_cond" in coreml_model.expected_inputs
|
|
|
|
|
|
def get_w_embedding(w, embedding_dim=512, dtype=torch.float32):
|
|
assert len(w.shape) == 1
|
|
w = w * 1000.0
|
|
|
|
half_dim = embedding_dim // 2
|
|
emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
|
|
emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
|
|
emb = w.to(dtype)[:, None] * emb[None, :]
|
|
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
|
|
if embedding_dim % 2 == 1: # zero pad
|
|
emb = torch.nn.functional.pad(emb, (0, 1))
|
|
assert emb.shape == (w.shape[0], embedding_dim)
|
|
return emb
|
|
|
|
|
|
def model_function_wrapper(w_embedding):
|
|
def wrapper(model_function, params):
|
|
x = params["input"]
|
|
t = params["timestep"]
|
|
c = params["c"]
|
|
|
|
context = c.get("c_crossattn")
|
|
|
|
if context is None:
|
|
return torch.zeros_like(x)
|
|
|
|
return model_function(x, t, **c, timestep_cond=w_embedding)
|
|
|
|
return wrapper
|
|
|
|
|
|
def lcm_patch(model):
|
|
m = model.clone()
|
|
sampling_type = LCM
|
|
sampling_base = ModelSamplingDiscreteDistilled
|
|
|
|
class ModelSamplingAdvanced(sampling_base, sampling_type):
|
|
pass
|
|
|
|
model_sampling = ModelSamplingAdvanced()
|
|
m.add_object_patch("model_sampling", model_sampling)
|
|
|
|
return m
|
|
|
|
|
|
def add_lcm_model_options(model_patcher, cfg, latent_image):
|
|
mp = model_patcher.clone()
|
|
|
|
latent = latent_image["samples"].to(get_torch_device())
|
|
batch_size = latent.shape[0]
|
|
dtype = latent.dtype
|
|
device = get_torch_device()
|
|
|
|
w = torch.tensor(cfg).repeat(batch_size)
|
|
w_embedding = get_w_embedding(w, embedding_dim=256).to(device=device, dtype=dtype)
|
|
|
|
model_options = {
|
|
"model_function_wrapper": model_function_wrapper(w_embedding),
|
|
"sampler_cfg_function": lambda x: x["cond"].to(device),
|
|
}
|
|
mp.model_options |= model_options
|
|
|
|
return mp
|