73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
import torch
|
|
|
|
import comfy.model_management as mm
|
|
|
|
|
|
def get_model_fn(model):
|
|
# sample, sample_null, cfg_scale
|
|
def model_fn(z, sigma, positive, negative, cfg):
|
|
model.dit.to(model.device)
|
|
if hasattr(model.dit, "cublas_half_matmul") and model.dit.cublas_half_matmul:
|
|
autocast_dtype = torch.float16
|
|
else:
|
|
autocast_dtype = torch.bfloat16
|
|
|
|
with torch.autocast(mm.get_autocast_device(model.device), dtype=autocast_dtype):
|
|
if cfg > 1.0:
|
|
out_cond = model.dit(z, sigma, **positive)
|
|
out_uncond = model.dit(z, sigma, **negative)
|
|
else:
|
|
out_cond = model.dit(z, sigma, **positive)
|
|
return out_cond
|
|
|
|
return out_uncond + cfg * (out_cond - out_uncond)
|
|
|
|
return model_fn
|
|
|
|
|
|
def get_sample_args(model, cond_embeds, uncond_embeds):
|
|
cond_args = {
|
|
"y_mask": [cond_embeds["attention_mask"].to(model.device)],
|
|
"y_feat": [cond_embeds["embeds"].to(model.device)]
|
|
}
|
|
|
|
uncond_args = {
|
|
"y_mask": [uncond_embeds["attention_mask"].to(model.device)],
|
|
"y_feat": [uncond_embeds["embeds"].to(model.device)]
|
|
}
|
|
return cond_args, uncond_args
|
|
|
|
|
|
def prepare_conds(positive, negative):
|
|
#For compatibility with Comfy CLIPTextEncode
|
|
if not isinstance(positive, dict):
|
|
positive = {
|
|
"embeds": positive[0][0],
|
|
"attention_mask": positive[0][1]["attention_mask"].bool(),
|
|
}
|
|
if not isinstance(negative, dict):
|
|
negative = {
|
|
"embeds": negative[0][0],
|
|
"attention_mask": negative[0][1]["attention_mask"].bool(),
|
|
}
|
|
return positive, negative
|
|
|
|
|
|
def generate_eta_values(steps, start_time, end_time, eta, eta_trend):
|
|
end_time = min(end_time, steps)
|
|
eta_values = [0] * steps
|
|
|
|
if eta_trend == 'constant':
|
|
for i in range(start_time, end_time):
|
|
eta_values[i] = eta
|
|
elif eta_trend == 'linear_increase':
|
|
for i in range(start_time, end_time):
|
|
progress = (i - start_time) / (end_time - start_time - 1)
|
|
eta_values[i] = eta * progress
|
|
elif eta_trend == 'linear_decrease':
|
|
for i in range(start_time, end_time):
|
|
progress = 1 - (i - start_time) / (end_time - start_time - 1)
|
|
eta_values[i] = eta * progress
|
|
|
|
return eta_values
|