File size: 3,076 Bytes
d382778 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """SAMPLING ONLY."""
import torch
from .dpm_solver_v3 import NoiseScheduleVP, model_wrapper, DPM_Solver_v3
class DPMSolverv3Sampler:
def __init__(self, ckp_path, stats_dir, model, steps, guidance_scale, **kwargs):
super().__init__()
self.model = model
to_torch = lambda x: x.clone().detach().to(torch.float32).to(model.device)
self.alphas_cumprod = to_torch(model.alphas_cumprod)
self.device = self.model.betas.device
self.guidance_scale = guidance_scale
self.ns = NoiseScheduleVP("discrete", alphas_cumprod=self.alphas_cumprod)
assert stats_dir is not None, f"No statistics file found in {stats_base}."
print("Use statistics", stats_dir)
self.dpm_solver_v3 = DPM_Solver_v3(
statistics_dir=stats_dir,
noise_schedule=self.ns,
steps=steps,
t_start=None,
t_end=None,
skip_type="time_uniform",
degenerated=False,
device=self.device,
)
self.steps = steps
@torch.no_grad()
def sample(
self,
batch_size,
shape,
conditioning=None,
x_T=None,
unconditional_conditioning=None,
use_corrector=False,
half=False,
# this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
**kwargs,
):
if conditioning is not None:
if isinstance(conditioning, dict):
cbs = conditioning[list(conditioning.keys())[0]].shape[0]
if cbs != batch_size:
print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
else:
if conditioning.shape[0] != batch_size:
print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
# sampling
C, H, W = shape
size = (batch_size, C, H, W)
if x_T is None:
img = torch.randn(size, device=self.device)
else:
img = x_T
if conditioning is None:
model_fn = model_wrapper(
lambda x, t, c: self.model.apply_model(x, t, c),
self.ns,
model_type="noise",
guidance_type="uncond",
)
ORDER = 3
else:
model_fn = model_wrapper(
lambda x, t, c: self.model.apply_model(x, t, c),
self.ns,
model_type="noise",
guidance_type="classifier-free",
condition=conditioning,
unconditional_condition=unconditional_conditioning,
guidance_scale=self.guidance_scale,
)
ORDER = 2
x = self.dpm_solver_v3.sample(
img,
model_fn,
order=ORDER,
p_pseudo=False,
c_pseudo=True,
lower_order_final=True,
use_corrector=use_corrector,
half=half,
)
return x.to(self.device), None
|