UDiffText / util.py
ZYMPKU's picture
cpu
660f2af
raw
history blame
3.69 kB
import torch
from omegaconf import OmegaConf
from sgm.util import instantiate_from_config
from sgm.modules.diffusionmodules.sampling import *
SD_XL_BASE_RATIOS = {
"0.5": (704, 1408),
"0.52": (704, 1344),
"0.57": (768, 1344),
"0.6": (768, 1280),
"0.68": (832, 1216),
"0.72": (832, 1152),
"0.78": (896, 1152),
"0.82": (896, 1088),
"0.88": (960, 1088),
"0.94": (960, 1024),
"1.0": (1024, 1024),
"1.07": (1024, 960),
"1.13": (1088, 960),
"1.21": (1088, 896),
"1.29": (1152, 896),
"1.38": (1152, 832),
"1.46": (1216, 832),
"1.67": (1280, 768),
"1.75": (1344, 768),
"1.91": (1344, 704),
"2.0": (1408, 704),
"2.09": (1472, 704),
"2.4": (1536, 640),
"2.5": (1600, 640),
"2.89": (1664, 576),
"3.0": (1728, 576),
}
def init_model(cfgs):
model_cfg = OmegaConf.load(cfgs.model_cfg_path)
ckpt = cfgs.load_ckpt_path
model = instantiate_from_config(model_cfg.model)
model.init_from_ckpt(ckpt)
if cfgs.type == "train":
model.train()
else:
if cfgs.use_gpu:
model.to(torch.device("cuda", index=cfgs.gpu))
model.eval()
model.freeze()
return model
def init_sampling(cfgs):
discretization_config = {
"target": "sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization",
}
if cfgs.dual_conditioner:
guider_config = {
"target": "sgm.modules.diffusionmodules.guiders.DualCFG",
"params": {"scale": cfgs.scale},
}
sampler = EulerEDMDualSampler(
num_steps=cfgs.steps,
discretization_config=discretization_config,
guider_config=guider_config,
s_churn=0.0,
s_tmin=0.0,
s_tmax=999.0,
s_noise=1.0,
verbose=True,
device=torch.device("cuda", index=cfgs.gpu)
)
else:
guider_config = {
"target": "sgm.modules.diffusionmodules.guiders.VanillaCFG",
"params": {"scale": cfgs.scale[0]},
}
sampler = EulerEDMSampler(
num_steps=cfgs.steps,
discretization_config=discretization_config,
guider_config=guider_config,
s_churn=0.0,
s_tmin=0.0,
s_tmax=999.0,
s_noise=1.0,
verbose=True,
device=torch.device("cuda", index=cfgs.gpu)
)
return sampler
def deep_copy(batch):
c_batch = {}
for key in batch:
if isinstance(batch[key], torch.Tensor):
c_batch[key] = torch.clone(batch[key])
elif isinstance(batch[key], (tuple, list)):
c_batch[key] = batch[key].copy()
else:
c_batch[key] = batch[key]
return c_batch
def prepare_batch(cfgs, batch):
for key in batch:
if isinstance(batch[key], torch.Tensor) and cfgs.use_gpu:
batch[key] = batch[key].to(torch.device("cuda", index=cfgs.gpu))
if not cfgs.dual_conditioner:
batch_uc = deep_copy(batch)
if "ntxt" in batch:
batch_uc["txt"] = batch["ntxt"]
else:
batch_uc["txt"] = ["" for _ in range(len(batch["txt"]))]
if "label" in batch:
batch_uc["label"] = ["" for _ in range(len(batch["label"]))]
return batch, batch_uc, None
else:
batch_uc_1 = deep_copy(batch)
batch_uc_2 = deep_copy(batch)
batch_uc_1["ref"] = torch.zeros_like(batch["ref"])
batch_uc_2["ref"] = torch.zeros_like(batch["ref"])
batch_uc_1["label"] = ["" for _ in range(len(batch["label"]))]
return batch, batch_uc_1, batch_uc_2